import numpy as np
import matplotlib.pyplot as plt
%config InlineBackend.figure_format = "svg"

S = 8
Jx,Jy,Jz=0.1,0.2,0.5 #random selection of couplings
hx,hy,hz=0.1,-0.2,0.5 #random selection of fields

nT = 501
Tmax = 25
Ts = np.linspace(0,Tmax,nT)
m_ED = np.zeros((nT,3))
m_TEBD = np.zeros((nT,3))

from scipy.linalg import expm

s0 = np.array([[1,0],[0,1]],dtype=complex)
sx = np.array([[0,1],[1,0]],dtype=complex)
sy = np.array([[0,-1j],[1j,0]],dtype=complex)
sz = np.array([[1,0],[0,-1]],dtype=complex)

pauli = {"0" : s0,
         "x" : sx,
         "y" : sy,
         "z" : sz}

def spin(polarization,s,S):
    if(s==0):
        temp = pauli[polarization]
    else:
        temp = s0
    for i in range(1,S,1):
        if(i==s):
            temp = np.kron(temp,pauli[polarization])
        else:
            temp = np.kron(temp,s0)
    return temp

def H_s(S,dir): #open boundaries
    H = np.zeros((2**S,2**S),dtype=complex)
    for i in range(0,S-1,1):
        H += spin(dir,i,S)@spin(dir,i+1,S)
    #H += spin(dir,S-1,S)@spin(dir,0,S) #periodic boundaries
    return H

HX,HY,HZ = H_s(S,"x"),H_s(S,"y"),H_s(S,"z")

def h_s(S,dir):
    H = np.zeros((2**S,2**S),dtype=complex)
    for i in range(0,S,1):
        H += spin(dir,i,S)
    return H

hX,hY,hZ = h_s(S,"x"),h_s(S,"y"),h_s(S,"z")

def Ham(S):
    #build the Hamiltonian
    H = Jx*HX + Jy*HY + Jz*HZ - (hx*hX + hy*hY + hz*hZ)
    return H

%%time
h = Ham(S)
_, vecs = np.linalg.eigh(h)

#initial condition
psi = np.zeros(2**S)
psi[0] = 1


#time-evolve
U_dt = expm(-1j*h*Tmax/(nT-1))
for t in range(nT):
    mags = np.zeros((S,3))
    for s in range(S):
        mags[s,0] = np.real(np.conj(psi).T@spin("x",s,S)/2@psi) #spin/2 since pauli matrices shouldn't have 1/2
        mags[s,1] = np.real(np.conj(psi).T@spin("y",s,S)/2@psi) #spin/2 since pauli matrices shouldn't have 1/2
        mags[s,2] = np.real(np.conj(psi).T@spin("z",s,S)/2@psi) #spin/2 since pauli matrices shouldn't have 1/2
    m_ED[t,0] = np.mean(mags[:,0])
    m_ED[t,1] = np.mean(mags[:,1])
    m_ED[t,2] = np.mean(mags[:,2])
    psi = U_dt@psi

from tenpy.algorithms import dmrg
from tenpy.algorithms import tebd

from tenpy.networks.mps import MPS
from tenpy.networks.site import SpinHalfSite

from tenpy.models.model import CouplingMPOModel
from tenpy.models.model import NearestNeighborModel

class XYZChain(CouplingMPOModel, NearestNeighborModel):
    default_lattice = "Chain"
    force_default_lattice = True

    def init_sites(self, model_params):
        return SpinHalfSite(conserve="None")

    def init_terms(self, model_params):
        # read out parameters
        Jx = model_params.get("Jx", 1.)
        Jy = model_params.get("Jy", 0.)
        Jz = model_params.get("Jz", 0.)
        hx = model_params.get("hx", 0.)
        hy = model_params.get("hy", 0.)
        hz = model_params.get("hz", 0.)
        # add terms
        for i in range(len(self.lat.unit_cell)):
            self.add_onsite(-hx, i, "Sx")
            self.add_onsite(-hy, i, "Sy")
            self.add_onsite(-hz, i, "Sz")
        for i1, i2, dx in self.lat.pairs["nearest_neighbors"]:
            self.add_coupling(Jx, i1, "Sx", i2, "Sx", dx)
            self.add_coupling(Jy, i1, "Sy", i2, "Sy", dx)
            self.add_coupling(Jz, i1, "Sz", i2, "Sz", dx)

%%time

model_params = dict(L=S, Jx=Jx,Jy=Jy,Jz=Jz, hx=hx/2,hy=hy/2,hz=hz/2, bc_MPS="finite") #1/2 * since difference between sigma_i and S_i
tebd_params = {
    "trunc_params": {
        "chi_max": 100, #bond dimension
        "svd_min": 1*10**-10
    },
    "N_steps": 1,
    "dt": 4*Tmax/(nT-1), #4* since difference between sigma_i and S_i the S_i*S_i terms are 1/4 as small so to evolve the same duration we need 4x the time step
    "order": 2,
}

M = XYZChain(model_params)
psi = MPS.from_lat_product_state(M.lat, [["up"]])
engine = tebd.TEBDEngine(psi, M, tebd_params)

for step in range(nT):
    m_TEBD[step,0] = np.mean(psi.expectation_value("Sx"))
    m_TEBD[step,1] = np.mean(psi.expectation_value("Sy"))
    m_TEBD[step,2] = np.mean(psi.expectation_value("Sz"))
    engine.run()

plt.plot(Ts,m _ED[:,0],label="ED x")
plt.scatter(Ts[::10],m_TEBD[::10,0],marker=".",label="TEBD x")
plt.plot(Ts,m_ED[:,1],label="ED y")
plt.scatter(Ts[::10],m_TEBD[::10,1],marker=".",label="TEBD y")
plt.plot(Ts,m_ED[:,2],label="ED z")
plt.scatter(Ts[::10],m_TEBD[::10,2],marker=".",label="TEBD z")
plt.xlabel(r"$t$")
plt.ylabel(r"$\langle S\rangle$")
plt.legend()