S = 10 #sites

import numpy as np
import matplotlib.pyplot as plt
from scipy.linalg import expm as matrixexp
from scipy.linalg import logm as matrixlog
%config InlineBackend.figure_format = "svg"

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,
         "+" : (sx+1j*sy)/2,
         "-" : (sx-1j*sy)/2}

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 vacuum(S):
    vac = np.zeros(2**S,dtype=complex)
    vac[-1] = 1
    return vac

def random_state(S): #random state--not in a particular number sector!
    state = vacuum(S)
    b = bin(2**S+np.random.randint(0,2**S))
    for i in range(1,S+1,1):
        if(int(b[-i])==1):
            state = spin("+",i-1,S)@state
    return state

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

HXX = H_XX(S)

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

HZ = H_Z(S)

def Z_s(S):
    #create a dictionary for the Sz operators
    numbers = {}
    for i in range(0,S,1):
        numbers[str(i)] = spin("z",i,S)
    return numbers

Szs = Z_s(S)

def H_h(S,W):
    #build the disorder Hamiltonian
    H = np.zeros((2**S,2**S),dtype=complex)
    for i in range(0,S,1):
        w = np.random.uniform(-W,W)
        H += w*Szs[str(i)]
    return H

def half_basis(S):
    basis_states = {}
    for i in range(0,2**(S//2),1):
        state = vacuum(S//2)
        b = bin(2**(S//2)+i)
        for j in range(1,S//2+1,1):
            if(int(b[-j])==1):
                state = spin("+",j-1,S//2)@state
        basis_states[str(i)] = state
    return basis_states

basis_states = half_basis(S)

def half_trace(rho,S): #trace rho over the first s sites of a length S chain
    rho_A = np.zeros((2**(S//2),2**(S//2)),dtype=complex) #not bothering to reduce dimensionality or take advantage of rho being real
    for i in range(0,2**(S//2),1):
        psi_j = np.kron(np.identity(2**(S//2),dtype=complex),basis_states[str(i)])
        rho_A += psi_j@rho@np.conjugate(psi_j.T)
    return rho_A

# %%time

J = 1
Jps = [0,0.01,0.1,0.2]
W = 5

timesteps = 31
num_states = 150
times = np.logspace(np.log10(5*10**-2),np.log10(2*10**3),timesteps)
S_EE = np.zeros((4,timesteps),dtype=complex)

# for j in range(0,4,1):
#     for s in range(0,num_states,1):
#         H = J*HXX + Jps[j]*HZ + H_h(S,W) #create a Hamiltonian with disorder
#         psi0 = random_state(S) #create a random state
#         for t in range(0,timesteps,1):
#             psi = matrixexp(-1j*times[t]*H)@psi0 #time evolve starting from a random product state
#             rho = np.outer(np.conjugate(psi),psi) #since this is a pure state
#             rho_A = half_trace(rho,S)
#             S_EE[j,t] += -np.trace(rho_A@matrixlog(rho_A))/num_states

#31,10 - 1 min with L=8
#31,40 - 1 hr with L=10

#np.save("SEE.npy",S_EE)

S_EE = np.load("SEE.npy")

plt.plot(times,np.log10(np.exp(1))*np.real(S_EE[0,:]),marker="o",c="#0800FF",label=r"$J'={}$".format(Jps[0]))
plt.plot(times,np.log10(np.exp(1))*np.real(S_EE[1,:]),marker="s",c="#008000",label=r"$J'={}$".format(Jps[1]))
plt.plot(times,np.log10(np.exp(1))*np.real(S_EE[2,:]),marker="d",c="#FE0000",label=r"$J'={}$".format(Jps[2]))
plt.plot(times,np.log10(np.exp(1))*np.real(S_EE[3,:]),marker="^",c="#01BFBF",label=r"$J'={}$".format(Jps[3]))
plt.semilogx()
plt.xlabel(r"Time ($t$)")
 
plt.legend()
plt.savefig("entropy-obc-10.pdf")
plt.show()