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

S = 10
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

m_ED = np.zeros((S,3))
m_DMRG = np.zeros((S,3))

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)
mags = np.zeros(S)
spins = ["x","y","z"]
for s in range(S):
    for d in range(3):
        m_ED[s,d] = np.real(np.conj(vecs[:,0]).T@spin(spins[d],s,S)/2@vecs[:,0]) #spin/2 since pauli matrices shouldn't have 1/2

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 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")
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 = XYZChain(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)
E0, psi = engine.run()
m_DMRG[:,0] = psi.expectation_value("Sx")
m_DMRG[:,1] = psi.expectation_value("Sy")
m_DMRG[:,2] = psi.expectation_value("Sz")

print(m_ED,"\n")
print(m_DMRG)

for d in range(3):
    plt.plot(np.linspace(1,S,S,endpoint=True),m_ED[:,d],label=r"ED, $S^{}$".format(spins[d]))
    plt.scatter(np.linspace(1,S,S,endpoint=True),m_DMRG[:,d],marker=".",label="DMRG")
plt.legend(loc=1)
plt.xlabel(r"site $i$")
plt.ylabel(r"$S_i$")

from mpl_toolkits.mplot3d import Axes3D
from matplotlib.patches import FancyArrowPatch
from mpl_toolkits.mplot3d import proj3d
import matplotlib.colors as clr

my_cmap = clr.LinearSegmentedColormap.from_list("",[(0,"#EEEEEE"),(0.5,"#8888EE"),(1,"#0000EE")])

class Arrow3D(FancyArrowPatch):
    def __init__(self, xs, ys, zs, *args, **kwargs):
        super().__init__((0,0), (0,0), *args, **kwargs)
        self._verts3d = xs, ys, zs

    def do_3d_projection(self, renderer=None):
        xs3d, ys3d, zs3d = self._verts3d
        xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, self.axes.M)
        self.set_positions((xs[0],ys[0]),(xs[1],ys[1]))
        return np.min(zs)

def axes():
    arrow_prop_dict = dict(mutation_scale=20, arrowstyle='-|>', color="k", shrinkA=0, shrinkB=0,linewidth=2)
    a = Arrow3D([0,S+1], [0, 0], [0, 0], **arrow_prop_dict)
    ax.add_artist(a)
    arrow_prop_dict = dict(mutation_scale=20, arrowstyle='-|>', color="k", shrinkA=0, shrinkB=0,linewidth=2)
    a = Arrow3D([0,0], [0, 1], [0, 0], **arrow_prop_dict)
    ax.add_artist(a)
    arrow_prop_dict = dict(mutation_scale=20, arrowstyle='-|>', color="k", shrinkA=0, shrinkB=0,linewidth=2)
    a = Arrow3D([0,0], [0, 0], [0, 1], **arrow_prop_dict)
    ax.add_artist(a)

def arrow(site,vec):
    len = np.sqrt(vec[0]**2+vec[1]**2+vec[2]**2)
    arrow_prop_dict = dict(mutation_scale=20, arrowstyle='-|>', color=my_cmap(len), shrinkA=0, shrinkB=0,linewidth=2)
    a = Arrow3D([site,site+vec[0]], [0, vec[1]], [0, vec[2]], **arrow_prop_dict)
    ax.add_artist(a)

fig = plt.figure(figsize=(15,5))
ax = fig.add_subplot(111, projection="3d")
ax.set_title("ED")
ax.set_axis_off()
axes()
ax.text(S+1,0,0,"x")
ax.text(0,1,0,"y")
ax.text(0,0,1,"z")
lenh = np.sqrt(hx**2+hy**2+hz**2)
ax.text(hx/lenh,hy/lenh,hz/lenh,"h")
arrow_prop_dict = dict(mutation_scale=20, arrowstyle="-|>", color="r", shrinkA=0, shrinkB=0,linewidth=2)
a = Arrow3D([0,hx/lenh], [0, hy/lenh], [0, hz/lenh], **arrow_prop_dict)
ax.add_artist(a)

#plot magnetizations
scale = 3 #to make arrows the right scale
for i in range(S):
    arrow(i+1,scale*m_ED[i])

ax.set_xlim(-1,S+1)
ax.set_ylim(-1,1)
ax.set_zlim(-1,1)

ax.view_init(30, -45)


ax = fig.add_subplot(121, projection="3d")
ax.set_title("DMRG")
ax.set_axis_off()
axes()
ax.text(S+1,0,0,"x")
ax.text(0,1,0,"y")
ax.text(0,0,1,"z")
lenh = np.sqrt(hx**2+hy**2+hz**2)
ax.text(hx/lenh,hy/lenh,hz/lenh,"h")
arrow_prop_dict = dict(mutation_scale=20, arrowstyle="-|>", color="r", shrinkA=0, shrinkB=0,linewidth=2)
a = Arrow3D([0,hx/lenh], [0, hy/lenh], [0, hz/lenh], **arrow_prop_dict)
ax.add_artist(a)

#plot magnetizations
scale = 3 #to make arrows the right scale
for i in range(S):
    arrow(i+1,scale*m_DMRG[i])

ax.set_xlim(-1,S+1)
ax.set_ylim(-1,1)
ax.set_zlim(-1,1)

ax.view_init(30, -45)