import numpy as np
import scipy
from scipy.optimize import minimize
import random
import math
import cmath

# Renormalization function takes 4 tensors Ta,Tb,Tc,Td with bond dimensions D**2 and adds the square coming from S1a S2b S3c S4d.
#all the decompositions are found for systematic checks
# The function takes Dc as the singular cut value to be used for the approximation.
# Renormalization returns the tensor T' made from A,B,C,D  corresponding S's decompositions.
def renormalization(Ta,Tb,Tc,Td,D,Dc):
    Ma=Ta.reshape((D**4,D**4))
    Mb=Tb.reshape((D**4,D**4))
    Mc=Tc.reshape((D**4,D**4))
    Md=Td.reshape((D**4,D**4))
    Ua, sa, Va=np.linalg.svd(Ma)
    Ub, sb, Vb=np.linalg.svd(Mb)
    Uc, sc, Vc=np.linalg.svd(Mc)
    Ud, sd, Vd=np.linalg.svd(Md)
    Xa1=np.zeros((D**4,D**4),complex)
    Xa2=np.zeros((D**4,D**4),complex)
    Xb1=np.zeros((D**4,D**4),complex)
    Xb2=np.zeros((D**4,D**4),complex)
    Xc3=np.zeros((D**4,D**4),complex)
    Xc4=np.zeros((D**4,D**4),complex)
    Xd3=np.zeros((D**4,D**4),complex)
    Xd4=np.zeros((D**4,D**4),complex)
    for a in range(D**4):
        for b in range(D**4):
                Xa1[a][b]=Ua[a][b]*math.sqrt(sa[b])
                Xa2[a][b]=Va[a][b]*math.sqrt(sa[a])
                Xb1[a][b]=Ub[a][b]*math.sqrt(sb[b])
                Xb2[a][b]=Vb[a][b]*math.sqrt(sb[a])
                Xc3[a][b]=Uc[a][b]*math.sqrt(sc[b])
                Xc4[a][b]=Vc[a][b]*math.sqrt(sc[a])
                Xd3[a][b]=Ud[a][b]*math.sqrt(sd[b])
                Xd4[a][b]=Vd[a][b]*math.sqrt(sd[a])    
    Sa1=Xa1.reshape((D**2,D**2,D**4))
    Sa2=Xa2.reshape((D**4,D**2,D**2))
    Sb1=Xb1.reshape((D**2,D**2,D**4))
    Sb2=Xb2.reshape((D**4,D**2,D**2))
    Sc3=Xc3.reshape((D**2,D**2,D**4))
    Sc4=Xc4.reshape((D**4,D**2,D**2))
    Sd3=Xd3.reshape((D**2,D**2,D**4))
    Sd4=Xd4.reshape((D**4,D**2,D**2))

    renormalizedT=np.zeros((Dc,Dc,Dc,Dc),complex)
    for q in range(Dc):
        for r in range(Dc):
            for s in range(Dc):
                for t in range(Dc):
                    for i in range(D**2):
                        for j in range(D**2):
                            for k in range(D**2):
                                for l in range(D**2):
                                    renormalizedT[q][r][s][t]+= Sa1[j][k][q]*Sb2[r][l][i]*Sc3[i][j][s]*Sd4[t][k][l]
                                    # print renormalizedT[q][r][s][t]     
    return renormalizedT

#tTr takes a 2x2 lattice of tensors (TA,TB,TC,TD) and gives the tensor trace over all bonds with 
#periodic boundary conditions
def tTr(TA,TB,TC,TD,Dim):        #tensor trace for a 2x2 lattice, TAxTBxTCxTD total tensor to be calculated the trace over, Dim=dimensionality of indexes of T's.
    res = 0.0
    for a in range(Dim):
        for b in range(Dim):
            for c in range(Dim):
                for d in range(Dim):
                    for e in range(Dim):
                        for f in range(Dim):
                            for g in range(Dim):
                                for h in range(Dim):
                                        res += TA[a,b,h,e]*TD[c,e,f,b]*TC[f,g,c,d]*TB[h,d,a,g]
                                        
    return res  



D=2             #Bond dimensionality
Dc=D**2         #Defines the cutoff Dc for the summation of singular values. (USe D**2 to iterate the renormalization step)
physdim=7            #Physical index dimension
I=1             #moment of inertia of quantum rotors
G=1.0            #1/g coupling
T=np.random.random((D,D,D,D,physdim))        #Randomize T (original tensor making the TPS) 
TT=np.zeros((D,D,D,D,D,D,D,D))                #Initialize \mathbb{T} in the paper
TT1=np.zeros((D,D,D,D,D,D,D,D),dtype=complex)
TT2=np.zeros((D,D,D,D,D,D,D,D),dtype=complex)
TT0= np.zeros((D,D,D,D,D,D,D,D), dtype=complex)

#This for loops assign the correponding values to have the 4 types of possible tensors per site discussed in the paper.
for a in range(D):
    for b in range(D):
        for c in range(D):
            for d in range(D):
                for e in range(D):
                    for f in range(D):
                        for g in range(D):
                            for h in range(D):
                                for m in range(physdim):
                                    TT[a][b][c][d][e][f][g][h] += T[a][b][c][d][m]*T[e][f][g][h][m]
                                    TT0[a][b][c][d][e][f][g][h] += -(((-(physdim-1)/2.+m)**2)/(2.*I))*T[a][b][c][d][m]*T[e][f][g][h][m]
                                    if (m+1)<physdim:
                                        TT1[a][b][c][d][e][f][g][h] += (1/2.)*cmath.sqrt(-G)*T[a][b][c][d][m]*T[e][f][g][h][m+1]
                                        TT2[a][b][c][d][e][f][g][h] += (1/2.)*cmath.sqrt(-G)*T[a][b][c][d][m]*T[e][f][g][h][m+1]
                                    if m-1>=0:
                                        TT1[a][b][c][d][e][f][g][h] += (1/2.)*cmath.sqrt(-G)*T[a][b][c][d][m]*T[e][f][g][h][m-1]
                                        TT2[a][b][c][d][e][f][g][h] += -(1/2.)*cmath.sqrt(-G)*T[a][b][c][d][m]*T[e][f][g][h][m-1]
#TTsaved just saves our original TT in order to perform more iterations with it later.                                    
TTsaved=TT   
TTA=TT1
TTB=TT1
TTC=TT
TTD=TT
#Put interaction on sites TTA,TTB
for i in range(4):
    TT=renormalization(TT,TT,TT,TT,D,D**2)
    TTA=renormalization(TTA,TT,TT,TTB,D,D**2)
    TTB=renormalization(TTC,TT,TTB,TT,D,D**2)
    TTC=renormalization(TT,TTC,TTD,TT,D,D**2)
    TTD=renormalization(TT,TTA,TT,TTD,D,D**2)

result1=tTr(TTA,TTB,TTC,TTD,D**2)
result2=tTr(TT,TT,TT,TT,D**2)

TT=TTsaved   
TTA=TT2
TTB=TT2
TTC=TT
TTD=TT
for i in range(4):
    TT=renormalization(TT,TT,TT,TT,D,D**2)
    TTA=renormalization(TTA,TT,TT,TTB,D,D**2)
    TTB=renormalization(TTC,TT,TTB,TT,D,D**2)
    TTC=renormalization(TT,TTC,TTD,TT,D,D**2)
    TTD=renormalization(TT,TTA,TT,TTD,D,D**2)

result3=tTr(TTA,TTB,TTC,TTD,D**2)
result4=tTr(TT,TT,TT,TT,D**2)

energy12 = result1/result2 + result3/result4        #bond energy (interagction from TTA and TTB)

TT=TTsaved
TTA=TT0
TTB=TT
TTC=TT
TTD=TT
for i in range(4):
    TTsaved=renormalization(TT,TT,TT,TT,D,D**2)
    TTA=renormalization(TTA,TT,TT,TTB,D,D**2)
    TTB=renormalization(TTC,TT,TTB,TT,D,D**2)
    TTC=renormalization(TT,TTC,TTD,TT,D,D**2)
    TTD=renormalization(TT,TTA,TT,TTD,D,D**2)

result1=tTr(TTA,TTB,TTC,TTD,D**2)
result2=tTr(TT,TT,TT,TT,D**2)

energy0=result1/result2          #onsite energy

#energy of one site and one bond term
totalenergy = energy0 + energy12