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

S = 10

nZ = 61
hzs = np.linspace(-3,3,nZ)
m_ED = np.zeros(nZ)
m_DMRG = np.zeros(nZ)

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 = H_s(S,"x")

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

hZ = h_s(S,"z")

def Ham(S,hz,Jx=1):
    #build the Hamiltonian
    H = Jx*HX - hz*hZ
    return H

%%time

for i in range(nZ):
    h = Ham(S,hzs[i])
    _, vecs = np.linalg.eigh(h)
    mags = np.zeros(S)
    for s in range(S):
        mags[s] = np.real(np.conj(vecs[:,0]).T@spin("z",s,S)/2@vecs[:,0]) #spin/2 since pauli matrices shouldn't have 1/2
    m_ED[i] = np.mean(mags)

from tenpy.networks.mps import MPS
from tenpy.algorithms import dmrg

from tenpy.models.model import CouplingMPOModel
from tenpy.models.model import NearestNeighborModel
from tenpy.networks.site import SpinHalfSite

class TFIMChain(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.)
        hz = model_params.get("hz", 0.)
        # add terms
        for i in range(len(self.lat.unit_cell)):
            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=1, hz=hzs[0], bc_MPS="finite")
dmrg_params = {
    "trunc_params": {
        "chi_max": 100, #bond dimension
        "svd_min": 1*10**-10
    },
    "max_E_err": 0.0001, #energy convergence step threshold
    "max_S_err": 0.0001, #entropy convergence step threshold
    "max_sweeps": 100  #may or may not be enough to converge
}

M = TFIMChain(model_params)
psi = MPS.from_product_state(M.lat.mps_sites(), (["up", "down"] * S)[:S], M.lat.bc_MPS)

engine = dmrg.TwoSiteDMRGEngine(psi, M, dmrg_params)

Sz = []
for h in hzs:
    model_params["hz"] = h/2 #Since Sz = Sigmaz/2 -- spin-1/2
    M = TFIMChain(model_params)
    engine.init_env(model=M)  # (re)initialize DMRG environment with new model
    # this uses the result from the previous DMRG as first initial guess
    E0, psi = engine.run()
    Sz.append(psi.expectation_value("Sz"))
m_DMRG = np.mean(Sz,axis=1)

# length and runtime
# 10->17
# 12->37
# 16->80
# 24->230

plt.plot(hzs,m_ED,label="ED")
plt.scatter(hzs,m_DMRG,marker=".",c="k",label="DMRG")
plt.legend()
plt.xlabel(r"$h^z$")
plt.ylabel(r"$\langle S^z\rangle$")