%matplotlib inline
import sys
sys.path.append('./src')
# Importing standard Qiskit libraries and configuring account
#from qiskit import QuantumCircuit, execute, Aer, IBMQ
#from qiskit.compiler import transpile, assemble
#from qiskit.tools.jupyter import *
#from qiskit.visualization import *
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import numpy as np
import ClassicalHubbardEvolutionChain as chc

import random as rand
import scipy.linalg as la

def get_bin(x, n=0):
    """
    Get the binary representation of x.
    Parameters: x (int), n (int, number of digits)"""
    binry = format(x, 'b').zfill(n)
    sup = list( reversed( binry[0:int(len(binry)/2)] ) )
    sdn = list( reversed( binry[int(len(binry)/2):len(binry)] ) )
    return format(x, 'b').zfill(n)

#============ Run Classical Evolution ==============

#Define our basis states
#States for 3 electrons with net spin up
'''
states = [ [[1,1,0],[1,0,0]], [[1,1,0],[0,1,0]], [[1,1,0], [0,0,1]],
           [[1,0,1],[1,0,0]], [[1,0,1],[0,1,0]], [[1,0,1], [0,0,1]],
           [[0,1,1],[1,0,0]], [[0,1,1],[0,1,0]], [[0,1,1], [0,0,1]] ]
'''
#States for 2 electrons in singlet state

states = [ [[1,0,0],[1,0,0]], [[1,0,0],[0,1,0]], [[1,0,0],[0,0,1]],
           [[0,1,0],[1,0,0]], [[0,1,0],[0,1,0]], [[0,1,0],[0,0,1]],
           [[0,0,1],[1,0,0]], [[0,0,1],[0,1,0]], [[0,0,1],[0,0,1]] ]
#'''
#States for a single electron
#states = [ [[1,0,0],[0,0,0]], [[0,1,0],[0,0,0]], [[0,0,1],[0,0,0]] ]


#Possible initial wavefunctions
#wfk = [0., 0., 0., 0., 1., 0., 0., 0., 0.] #Half-filling initial state
wfk = [0., 0., 0., 0., 1.0, 0., 0., 0., 0.] #2 electron initial state
#wfk = [0., 1., 0.] #1 electron initial state

#System parameters
t = 1.0
U = 2.
classical_time_step = 0.01
classical_total_time = 500*0.01
times = np.arange(0., classical_total_time, classical_time_step)

evolution, engs = chc.sys_evolve(states, wfk, t, U, classical_total_time, classical_time_step)
#print(evolution)
probs = [np.sum(x) for x in evolution]
#print(probs)
#print(np.sum(evolution[0]))
print(engs)



states_list = []
nsite = 3
for state in range(0, 2**(2*nsite)):
    state_bin = get_bin(state, 2*nsite)
    state_list = [[],[]]
    for mode in range(0,nsite):
        state_list[0].append(int(state_bin[mode]))
        state_list[1].append(int(state_bin[mode+nsite]))
    #print(state_list)
    states_list.append(state_list)
#print(states_list[18])

#evolution2, engs2 = chc.sys_evolve(states_list, wfk_full, t, U, classical_total_time, classical_time_step)

#System parameters
t = 1.0
U = 2.
classical_time_step = 0.01
classical_total_time = 500*0.01
times = np.arange(0., classical_total_time, classical_time_step)

#Create full Hamiltonian
wfk_full = np.zeros(len(states_list))
#wfk_full[18] = 1.    #010010
wfk_full[21] = 1.    #010101
#wfk_full[2] = 1.   #000010
#evolution2, engs2 = chc.sys_evolve(states_list, wfk_full, t, U, classical_total_time, classical_time_step)


def repel(l,state):
    if state[0][l]==1 and state[1][l]==1:
        return state
    else:
        return []
#Check if two states are different by a single hop
def hop(psii, psij, hopping):
    #Check spin down
    hopp = 0
    if psii[0]==psij[0]:
        #Create array of indices with nonzero values
        '''
        indi = np.nonzero(psii[1])[0]
        indj = np.nonzero(psij[1])[0]
        if len(indi) != len(indj):
            return hopp
        print('ind_i: ',indi,'    ind_j: ',indj)
        for i in range(len(indi)):
            if abs(indi[i]-indj[i])==1:
                hopp = -hopping
                print('Hopping Found: ',psii,'     with: ',psij)
                return hopp
        '''
        hops = []
        for site in range(len(psii[0])):
            if psii[1][site] != psij[1][site]:
                hops.append(site)
        if len(hops)==2 and np.sum(psii[1]) == np.sum(psij[1]):
            if hops[1]-hops[0]==1:
                hopp = -hopping
                return hopp
    #Check spin up
    if psii[1]==psij[1]:
        '''
        indi = np.nonzero(psii[0])[0]
        indj = np.nonzero(psij[0])[0]
        if len(indi) != len(indj):
            return hopp
        print('ind_i: ',indi,'    ind_j: ',indj)
        for i in range(len(indi)):
            if abs(indi[i]-indj[i])==1:
                hopp = -hopping
                print('Hopping Found: ',psii,'     with: ',psij)
                return hopp
        '''
        hops = []
        for site in range(len(psii[1])):
            if psii[0][site] != psij[0][site]:
                hops.append(site)
        if len(hops)==2 and np.sum(psii[0])==np.sum(psij[0]):
            if hops[1]-hops[0]==1:
                hopp = -hopping
                return hopp
    return hopp

def get_hamiltonian(states, t, U):
    H = np.zeros((len(states),len(states)) )
    #Construct Hamiltonian matrix
    for i in range(len(states)):
        psi_i = states[i]
        for j in range(i, len(states)):
            psi_j = states[j]
            if j==i:
                for l in range(0,len(states[0][0])):
                    if psi_i == repel(l,psi_j):
                        H[i,j] = U
                        break
            else:
                #print('psi_i: ',psi_i,'    psi_j: ',psi_j)
                H[i,j] = hop(psi_i, psi_j, t)
                H[j,i] = H[i,j]
    return H

hamil = get_hamiltonian(states_list, t, U)
print(hamil)
print(states_list)
print()
print("Target state: ", states_list[21])


mapping = get_mapping(states_list)
print('Second mapping set')
print(mapping[1])
print(wfk_full)
mapped_wfk = np.zeros(6)


for i in range(len(mapping)):
    if 21 in mapping[i]:
        print("True for mapping set: ",i)
    

def get_mapping(states):
    num_sites = len(states[0][0])
    mode_list = []
    for i in range(0,2*num_sites):
        index_list = []
        for state_index in range(0,len(states)):
            state = states[state_index]
        #Check spin-up modes
            if i < num_sites:
                if state[0][i]==1:
                    index_list.append(state_index)
        #Check spin-down modes
            else:
                if state[1][i-num_sites]==1:
                    index_list.append(state_index)
        if index_list:
            mode_list.append(index_list)
    return mode_list

def wfk_energy(wfk, hamil):
    eng = np.dot(np.conj(wfk), np.dot(hamil, wfk))
    return eng

def get_variance(wfk, h):
    h_squared = np.matmul(h, h)
    eng_squared = np.vdot(wfk, np.dot(h_squared, wfk))
    squared_eng = np.vdot(wfk, np.dot(h, wfk))
    var = np.sqrt(eng_squared - squared_eng)
    return var


def sys_evolve(states, init_wfk, hopping, repulsion, total_time, dt):
    hamiltonian = get_hamiltonian(states, hopping, repulsion)
    t_operator = la.expm(-1j*hamiltonian*dt)

    mapping = get_mapping(states)
    print(mapping)

    #Initalize system
    tsteps = int(total_time/dt)
    evolve = np.zeros([tsteps, len(init_wfk)])
    mode_evolve = np.zeros([tsteps, len(mapping)])
    wfk = init_wfk
    energies = np.zeros(tsteps)

    #Store first time step in mode_evolve
    evolve[0] = np.multiply(np.conj(wfk), wfk)
    for i in range(0, len(mapping)):
        wfk_sum = 0.
        norm = 0.
        print("Mapping: ", mapping[i])
        for j in mapping[i]:
            print(evolve[0][j])
            wfk_sum += evolve[0][j]
        
        mode_evolve[0][i] = wfk_sum
    energies[0] = wfk_energy(wfk, hamiltonian)
    norm = np.sum(mode_evolve[0])
    mode_evolve[0][:] = mode_evolve[0][:] / norm
        
        #print('wfk_sum: ',wfk_sum,'    norm: ',norm)
        
        #print('Variance: ',get_variance(wfk, hamiltonian) )
#Now do time evolution
    print(mode_evolve[0])
    times = np.arange(0., total_time, dt)
    for t in range(1, tsteps):
        wfk = np.dot(t_operator, wfk)
        evolve[t] = np.multiply(np.conj(wfk), wfk)
        #print(evolve[t])
        energies[t] = wfk_energy(wfk, hamiltonian)
        for i in range(0, len(mapping)):
            norm = 0.
            wfk_sum = 0.
            for j in mapping[i]:
                wfk_sum += evolve[t][j]
            mode_evolve[t][i] = wfk_sum
        norm = np.sum(mode_evolve[t])
        mode_evolve[t][:] = mode_evolve[t][:] / norm

    #Return time evolution
    return mode_evolve, energies
evolution2, engs2 = sys_evolve(states_list, wfk_full, t, U, classical_total_time, classical_time_step)

print(evolution2)
#print(engs2)
colors = list(mcolors.TABLEAU_COLORS.keys())
fig2, ax2 = plt.subplots(figsize=(20,10))
for i in range(3):
    #Create string label
    strup = "Site "+str(i+1)+r'$\uparrow$'
    strdwn = "Site "+str(i+1)+r'$\downarrow$'
    ax2.plot(times, evolution2[:,i], linestyle='-', color=str(colors[i]), linewidth=2, label=strup)
    ax2.plot(times, evolution2[:,i+3], linestyle='--', color=str(colors[i]), linewidth=2, label=strdwn)
    
plt.legend()

#print(evolution[10])
#H = np.array([[0., -t, 0.],[-t, 0., -t],[0., -t, 0.]])
#print(np.dot(np.conj(evolution[10]), np.dot(H, evolution[10])))
print(engs[0])
plt.plot(times, engs)
plt.xlabel('Time')
plt.ylabel('Energy')

#Save data
import json
fname = './data/classical_010101.json'
data = {'times': list(times)}
for i in range(3):
    key1 = 'site_'+str(i)+'_up'
    key2 = 'site_'+str(i)+'_down'
    data[key1] = list(evolution2[:,i])
    data[key2] = list(evolution2[:,i+3])

with open(fname, 'w') as fp:
    json.dump(data, fp, indent=4)

#Create plots of the processed data
fig2, ax2 = plt.subplots(figsize=(20,10))
colors = list(mcolors.TABLEAU_COLORS.keys())
cos = np.cos(np.sqrt(2)*times)**2

sit1 = "Site "+str(1)+r'$\uparrow$'
sit2 = "Site "+str(2)+r'$\uparrow$'
sit3 = "Site "+str(3)+r'$\uparrow$'
#ax2.plot(times, evolution[:,0], marker='.', color='k', linewidth=2, label=sit1)
#ax2.plot(times, evolution[:,1], marker='.', color=str(colors[0]), linewidth=2, label=sit2)
#ax2.plot(times, evolution[:,2], marker='.', color=str(colors[1]), linewidth=1.5, label=sit3)
#ax2.plot(times, cos, label='cosdat')
#ax2.plot(times, np.zeros(len(times)))


for i in range(3):
    #Create string label
    strup = "Site "+str(i+1)+r'$\uparrow$'
    strdwn = "Site "+str(i+1)+r'$\downarrow$'
    ax2.plot(times, evolution[:,i], linestyle='-', color=str(colors[i]), linewidth=2, label=strup)
    ax2.plot(times, evolution[:,i+3], linestyle='--', color=str(colors[i]), linewidth=2, label=strdwn)




#ax2.set_ylim(0, 1.)
ax2.set_xlim(0, classical_total_time)
#ax2.set_xlim(0, 1.)
ax2.set_xticks(np.arange(0,classical_total_time, 0.2))
#ax2.set_yticks(np.arange(0,1.1, 0.05))
ax2.tick_params(labelsize=16)
ax2.set_title('Time Evolution of 3 Site One Dimensional Chain', fontsize=22)
ax2.set_xlabel('Time', fontsize=24)
ax2.set_ylabel('Probability', fontsize=24)
ax2.legend(fontsize=20)



#Try by constructing the matrix and finding the eigenvalues
N = 3
Nup = 2
Ndwn = N - Nup

t = 1.
U = 2.

#Check if two states are different by a single hop
def hop(psii, psij):
    #Check spin down
    hopp = 0
    if psii[0]==psij[0]:
        #Create array of indices with nonzero values
        indi = np.nonzero(psii[1])[0]
        indj = np.nonzero(psij[1])[0]
        for i in range(len(indi)):
            if abs(indi[i]-indj[i])==1:
                hopp = -t
                return hopp
    #Check spin up
    if psii[1]==psij[1]:
        indi = np.nonzero(psii[0])[0]
        indj = np.nonzero(psij[0])[0]
        for i in range(len(indi)):
            if abs(indi[i]-indj[i])==1:
                hopp = -t
                return hopp
    return hopp
            

#On-site terms
def repel(l,state):
    if state[0][l]==1 and state[1][l]==1:
        return state
    else:
        return []
        
#States for 3 electrons with net spin up
states = [ [[1,1,0],[1,0,0]], [[1,1,0],[0,1,0]], [[1,1,0], [0,0,1]],
           [[1,0,1],[1,0,0]], [[1,0,1],[0,1,0]], [[1,0,1], [0,0,1]],
           [[0,1,1],[1,0,0]], [[0,1,1],[0,1,0]], [[0,1,1], [0,0,1]] ]

#States for 2 electrons in singlet state
#states = [ [[1,0,0],[1,0,0]], [[1,0,0],[0,1,0]], [[1,0,0],[0,0,1]],
#           [[0,1,0],[1,0,0]], [[0,1,0],[0,1,0]], [[0,1,0],[0,0,1]],
#           [[0,0,1],[1,0,0]], [[0,0,1],[0,1,0]], [[0,0,1],[0,0,1]] ]

#States for a single electron
states = [ [[1,0,0],[0,0,0]], [[0,1,0],[0,0,0]], [[0,0,1],[0,0,0]] ]

H = np.zeros((len(states),len(states)) )
#Construct Hamiltonian matrix
for i in range(len(states)):
    psi_i = states[i]
    for j in range(len(states)):
        psi_j = states[j]
        
        if j==i:
            for l in range(0,N):
                if psi_i == repel(l,psi_j):
                    H[i,j] = U
                    break
        else:
            H[i,j] = hop(psi_i, psi_j)

print(H)
results = la.eig(H)

print()
for i in range(len(results[0])):
    print('Eigenvalue: ',results[0][i])
    print('Eigenvector: \n',results[1][i])
    print('Norm: ', np.linalg.norm(results[1][i]))
    print('Density Matrix: ')
    print(np.outer(results[1][i],results[1][i]))
    print()


dens_ops = []
eigs = []
for vec in results[1]:
    dens_ops.append(np.outer(results[1][i],results[1][i]))
    eigs.append(results[0][i])
    print(dens_ops)


dt = 0.01
tsteps = 450
times = np.arange(0., tsteps*dt, dt)

t_op = la.expm(-1j*H*dt)
#print(np.subtract(np.identity(len(H)), dt*H*1j))
#print(t_op)

wfk = [0., 0., 0., 0., 1., 0., 0., 0., 0.]
wfk = [0., 0., 0., 0., 1.0, 0., 0., 0., 0.] #2 electron initial state



evolve = np.zeros([tsteps, len(wfk)])
mode_evolve = np.zeros([tsteps, 6])
evolve[0] = wfk

#Figure out how to generalize this later
#'''[[0, 1, 2], [3, 4, 5], [6, 7, 8], [0, 3, 6], [1, 4, 7], [2, 5, 8]]
mode_evolve[0][0] = (evolve[0][0]+evolve[0][1]+evolve[0][2]) /2.
mode_evolve[0][1] = (evolve[0][3]+evolve[0][4]+evolve[0][5]) /2.
mode_evolve[0][2] = (evolve[0][6]+evolve[0][7]+evolve[0][8]) /2.
mode_evolve[0][3] = (evolve[0][0]+evolve[0][3]+evolve[0][6]) /2.
mode_evolve[0][4] = (evolve[0][1]+evolve[0][4]+evolve[0][7]) /2.
mode_evolve[0][5] = (evolve[0][2]+evolve[0][5]+evolve[0][8])  /2.
'''
mode_evolve[0][0] = (evolve[0][0]+evolve[0][1]+evolve[0][2]+evolve[0][3]+evolve[0][4]+evolve[0][5]) /3.
mode_evolve[0][1] = (evolve[0][0]+evolve[0][1]+evolve[0][2]+evolve[0][6]+evolve[0][7]+evolve[0][8]) /3.
mode_evolve[0][2] = (evolve[0][3]+evolve[0][4]+evolve[0][5]+evolve[0][6]+evolve[0][7]+evolve[0][8]) /3.
mode_evolve[0][3] = (evolve[0][0]+evolve[0][3]+evolve[0][6]) /3.
mode_evolve[0][4] = (evolve[0][1]+evolve[0][4]+evolve[0][7]) /3.
mode_evolve[0][5] = (evolve[0][2]+evolve[0][5]+evolve[0][8]) /3.
'''
print(mode_evolve[0])
#Define density matrices


for t in range(1, tsteps):
    #t_op = la.expm(-1j*H*t)
    wfk = np.dot(t_op, wfk)
    evolve[t] = np.multiply(np.conj(wfk), wfk)
    norm = np.sum(evolve[t])
    #print(evolve[t])
    #Store data in modes rather than basis defined in 'states' variable
    #''' #Procedure for two electrons
    mode_evolve[t][0] = (evolve[t][0]+evolve[t][1]+evolve[t][2]) / (2)
    mode_evolve[t][1] = (evolve[t][3]+evolve[t][4]+evolve[t][5]) / (2)
    mode_evolve[t][2] = (evolve[t][6]+evolve[t][7]+evolve[t][8]) / (2)
    mode_evolve[t][3] = (evolve[t][0]+evolve[t][3]+evolve[t][6]) / (2)
    mode_evolve[t][4] = (evolve[t][1]+evolve[t][4]+evolve[t][7]) / (2)
    mode_evolve[t][5] = (evolve[t][2]+evolve[t][5]+evolve[t][8]) / (2)
    '''
    mode_evolve[t][0] = (evolve[t][0]+evolve[t][1]+evolve[t][2]+evolve[t][3]+evolve[t][4]+evolve[t][5]) /3.
    mode_evolve[t][1] = (evolve[t][0]+evolve[t][1]+evolve[t][2]+evolve[t][6]+evolve[t][7]+evolve[t][8]) /3.
    mode_evolve[t][2] = (evolve[t][3]+evolve[t][4]+evolve[t][5]+evolve[t][6]+evolve[t][7]+evolve[t][8]) /3.
    mode_evolve[t][3] = (evolve[t][0]+evolve[t][3]+evolve[t][6]) /3.
    mode_evolve[t][4] = (evolve[t][1]+evolve[t][4]+evolve[t][7]) /3.
    mode_evolve[t][5] = (evolve[t][2]+evolve[t][5]+evolve[t][8]) /3.
    print(mode_evolve[t])
    '''
    #print(np.linalg.norm(evolve[t]))
    

#print(len(evolve[:,0]) )
#print(len(times))
#print(evolve[:,0])
#print(min(evolve[:,0]))



#Create plots of the processed data
fig2, ax2 = plt.subplots(figsize=(20,10))
colors = list(mcolors.TABLEAU_COLORS.keys())
cos = np.cos(np.sqrt(2)*times)**2

sit1 = "Site "+str(1)+r'$\uparrow$'
sit2 = "Site "+str(2)+r'$\uparrow$'
sit3 = "Site "+str(3)+r'$\uparrow$'
#ax2.plot(times, evolve[:,0], marker='.', color='k', linewidth=2, label=sit1)
#ax2.plot(times, evolve[:,1], marker='.', color=str(colors[0]), linewidth=2, label=sit2)
#ax2.plot(times, evolve[:,2], marker='.', color=str(colors[1]), linewidth=1.5, label=sit3)
#ax2.plot(times, cos, label='cosdat')
#ax2.plot(times, np.zeros(len(times)))

for i in range(3):
    #Create string label
    strup = "Site "+str(i+1)+r'$\uparrow$'
    strdwn = "Site "+str(i+1)+r'$\downarrow$'
    ax2.plot(times, mode_evolve[:,i], linestyle='-', color=str(colors[i]), linewidth=2, label=strup)
    ax2.plot(times, mode_evolve[:,i+3], linestyle='--', color=str(colors[i]), linewidth=2, label=strdwn)




#ax2.set_ylim(0, 1.)
ax2.set_xlim(0, tsteps*dt+dt/2.)
#ax2.set_xlim(0, 1.)
ax2.set_xticks(np.arange(0,tsteps*dt+dt, 0.2))
#ax2.set_yticks(np.arange(0,1.1, 0.05))
ax2.tick_params(labelsize=16)
ax2.set_title('Time Evolution of 3 Site One Dimensional Chain', fontsize=22)
ax2.set_xlabel('Time', fontsize=24)
ax2.set_ylabel('Probability', fontsize=24)
ax2.legend(fontsize=20)

dt = 0.1
tsteps = 50
times = np.arange(0., tsteps*dt, dt)

t_op = la.expm(-1j*H*dt)
#print(np.subtract(np.identity(len(H)), dt*H*1j))
#print(t_op)
#wfk = [0., 1., 0., 0., .0, 0., 0., 0., 0.] #Half-filling initial state
wfk0 = [0., 0., 0., 0., 1.0, 0., 0., 0., 0.] #2 electron initial state
#wfk0 = [0., 1., 0.] #1 electron initial state


evolve = np.zeros([tsteps, len(wfk0)])
mode_evolve = np.zeros([tsteps, 6])
evolve[0] = wfk0

#Figure out how to generalize this later
#'''
mode_evolve[0][0] = (evolve[0][0]+evolve[0][1]+evolve[0][2]) /2.
mode_evolve[0][1] = (evolve[0][3]+evolve[0][4]+evolve[0][5]) /2.
mode_evolve[0][2] = (evolve[0][6]+evolve[0][7]+evolve[0][8]) /2.
mode_evolve[0][3] = (evolve[0][0]+evolve[0][3]+evolve[0][6]) /2.
mode_evolve[0][4] = (evolve[0][1]+evolve[0][4]+evolve[0][7]) /2.
mode_evolve[0][5] = (evolve[0][2]+evolve[0][5]+evolve[0][8])  /2.
'''
mode_evolve[0][0] = (evolve[0][0]+evolve[0][1]+evolve[0][2]+evolve[0][3]+evolve[0][4]+evolve[0][5]) /3.
mode_evolve[0][1] = (evolve[0][0]+evolve[0][1]+evolve[0][2]+evolve[0][6]+evolve[0][7]+evolve[0][8]) /3.
mode_evolve[0][2] = (evolve[0][3]+evolve[0][4]+evolve[0][5]+evolve[0][6]+evolve[0][7]+evolve[0][8]) /3.
mode_evolve[0][3] = (evolve[0][0]+evolve[0][3]+evolve[0][6]) /3.
mode_evolve[0][4] = (evolve[0][1]+evolve[0][4]+evolve[0][7]) /3.
mode_evolve[0][5] = (evolve[0][2]+evolve[0][5]+evolve[0][8]) /3.
#'''

#Define density matrices


for t in range(1, tsteps):
    t_op = la.expm(-1j*H*t*dt)
    wfk = np.dot(t_op, wfk0)
    evolve[t] = np.multiply(np.conj(wfk), wfk)
    norm = np.sum(evolve[t])
    print(evolve[t])
    #Store data in modes rather than basis defined in 'states' variable
    #'''
    #Procedure for two electrons
    mode_evolve[t][0] = (evolve[t][0]+evolve[t][1]+evolve[t][2]) / (2)
    mode_evolve[t][1] = (evolve[t][3]+evolve[t][4]+evolve[t][5]) / (2)
    mode_evolve[t][2] = (evolve[t][6]+evolve[t][7]+evolve[t][8]) / (2)
    mode_evolve[t][3] = (evolve[t][0]+evolve[t][3]+evolve[t][6]) / (2)
    mode_evolve[t][4] = (evolve[t][1]+evolve[t][4]+evolve[t][7]) / (2)
    mode_evolve[t][5] = (evolve[t][2]+evolve[t][5]+evolve[t][8]) / (2)
    #Procedure for half-filling
    '''
    mode_evolve[t][0] = (evolve[t][0]+evolve[t][1]+evolve[t][2]+evolve[t][3]+evolve[t][4]+evolve[t][5]) /3.
    mode_evolve[t][1] = (evolve[t][0]+evolve[t][1]+evolve[t][2]+evolve[t][6]+evolve[t][7]+evolve[t][8]) /3.
    mode_evolve[t][2] = (evolve[t][3]+evolve[t][4]+evolve[t][5]+evolve[t][6]+evolve[t][7]+evolve[t][8]) /3.
    mode_evolve[t][3] = (evolve[t][0]+evolve[t][3]+evolve[t][6]) /3.
    mode_evolve[t][4] = (evolve[t][1]+evolve[t][4]+evolve[t][7]) /3.
    mode_evolve[t][5] = (evolve[t][2]+evolve[t][5]+evolve[t][8]) /3.
    #'''
    #print(mode_evolve[t])
    #print(np.linalg.norm(evolve[t]))
    

#print(len(evolve[:,0]) )
#print(len(times))
#print(evolve[:,0])
#print(min(evolve[:,0]))


#Create plots of the processed data
fig2, ax2 = plt.subplots(figsize=(20,10))
colors = list(mcolors.TABLEAU_COLORS.keys())
cos = np.cos(np.sqrt(2)*times)**2

sit1 = "Site "+str(1)+r'$\uparrow$'
sit2 = "Site "+str(2)+r'$\uparrow$'
sit3 = "Site "+str(3)+r'$\uparrow$'
#ax2.plot(times, evolve[:,0], marker='.', color='k', linewidth=2, label=sit1)
#ax2.plot(times, evolve[:,1], marker='.', color=str(colors[0]), linewidth=2, label=sit2)
#ax2.plot(times, evolve[:,2], marker='.', color=str(colors[1]), linewidth=1.5, label=sit3)
#ax2.plot(times, cos, label='cosdat')
#ax2.plot(times, np.zeros(len(times)))

for i in range(3):
    #Create string label
    strup = "Site "+str(i+1)+r'$\uparrow$'
    strdwn = "Site "+str(i+1)+r'$\downarrow$'
    ax2.plot(times, mode_evolve[:,i], linestyle='-', color=str(colors[i]), linewidth=2, label=strup)
    ax2.plot(times, mode_evolve[:,i+3], linestyle='--', color=str(colors[i]), linewidth=2, label=strdwn)


#ax2.set_ylim(0, 1.)
ax2.set_xlim(0, tsteps*dt+dt/2.)
#ax2.set_xlim(0, 1.)
ax2.set_xticks(np.arange(0,tsteps*dt+dt, 0.2))
#ax2.set_yticks(np.arange(0,1.1, 0.05))
ax2.tick_params(labelsize=16)
ax2.set_title('Time Evolution of 3 Site One Dimensional Chain', fontsize=22)
ax2.set_xlabel('Time', fontsize=24)
ax2.set_ylabel('Probability', fontsize=24)
ax2.legend(fontsize=20)

#Calculate total energy for 1D Hubbard model of N sites

#Number of sites in chain
N = 10
Ne = 10

#Filling factor
nu = Ne/N

#Hamiltonian parameters
t = 2.0 #Hopping
U = 4.0 #On-site repulsion

state = np.zeros(2*N)

#Add in space to populate array randomly, but do it by hand for now
#Populate Psi
n = 0
while n < Ne:
    site = rand.randint(0,N-1)
    spin = rand.randint(0,1)
    if state[site+spin*N]==0:
        state[site+spin*N] = 1
        n+=1
print(state)
#Loop over state and gather energy
E = 0

#############  ADD UP ENERGY FROM STATE  #############
#Add hoppings at edges
if state[0]==1:
    if state[1]==0:
        E+=t/2.
    if state[N-1]==0: #Periodic boundary
        E+=t/2.
if state[N]==1:
    if state[N+1]==0:
        E+=t/2.
    if state[-1]==0:
        E+=t/2.
print('Ends of energy are: ',E)
for i in range(1,N-1):
    print('i is: ',i)
    #Check spin up sites if site is occupied and if electron can hop
    if state[i]==1:
        if state[i+1]==0:
            E+=t/2.
        if state[i-1]==0:
            E+=t/2.
    #Check spin down sites if site is occupied and if electron can hop
    j = i+N
    if state[j]==1:
        if state[j+1]==0:
            E+=t/2.
        if state[j-1]==0:
            E+=t/2.
#Check Hubbard repulsion terms
for i in range(0,N):
    if state[i]==1 and state[i+N]==1:
        E+=U/4.
        
        
        

print('Energy is: ',E)
print('State is: ',state)

#Try by constructing the matrix and finding the eigenvalues
N = 3
Nup = 2
Ndwn = N - Nup

t = 1.0
U = 2.5

#To generalize, try using permutations function but for now hard code this
#Store a list of 2d list of all the states where the 2d list stores the spin-up occupation
#and the spin down occupation
states = [ [1,1,0,1,0,0], [1,1,0,0,1,0], [1,1,0,0,0,1],
          [1,0,1,1,0,0], [1,0,1,0,1,0], [1,0,1,0,0,1],
          [0,1,1,1,0,0], [0,1,1,0,1,0], [0,1,1,0,0,1] ]

print(len(states))
H = np.zeros((len(states),len(states)) )
print(H[0,4])
print(states[0])

for i in range(len(states)):
    psi_i = states[i]
    for j in range(len(states)):
        psi_j = states[j]
        #Check over sites
        #Check rest of state for hopping or double occupation
        for l in range(1,N-1):
            #Check edges
            
            if psi_i[l]==1 and (psi_j[l+1]==1 and psi_i[l+1]==0) or (psi_j[l-1]==1 and psi_i[l-1]==0):
                H[i,j] = -t/2.
                break
            if psi_i[l+N]==1 and (psi_j[l+1+N]==1 and psi_i[l+1+N]==0) or (psi_j[l-1+N]==1 and psi_i[l-1+N]==0):
                H[i,j] = -t/2.
                break
            if psi_i==psi_j:
                for l in range(N):
                    if psi_i[l]==1 and psi_j[l+N]==1:
                        H[i,j] = U/4.
                        break
                        
print(H)
tst = [[0,1],[2,3]]
print(tst[1][1])

psi_i = [[1,1,0],[1,0,0]]
psi_j = [[1,1,0],[0,1,0]]

print(psi_j[0])
print(psi_j[1], ' 1st: ',psi_j[1][0], ' 2nd: ',psi_j[1][1], ' 3rd: ',psi_j[1][2])

for l in range(N):
    print('l: ',l)
    if psi_j[1][l]==0 and psi_j[1][l-1]==1:
        psi_j[1][l]=1
        psi_j[1][l-1]=0
        print('1st')
        print(psi_j)
        break
    if psi_j[1][l-1]==0 and psi_j[1][l]==1:
        psi_j[1][l-1]=1
        psi_j[1][l]=0
        print('2nd')
        print(psi_j)
        break
    if psi_j[1][l]==0 and psi_j[1][l+1]==1:
        psi_j[1][l]=1
        psi_j[1][l+1]=0
        print('3rd: l=',l,' l+1=',l+1)
        print(psi_j)
        break
    if psi_j[1][l+1]==0 and psi_j[1][l]==1:
        psi_j[1][l+1]=1
        psi_j[1][l]=0
        print('4th')
        print(psi_j)
        break
def hoptst(l,m,spin,state):
    #Spin is either 0 or 1 which corresponds to which array w/n state we're examining
    if (state[spin][l]==0 and state[spin][m]==1):
        state[spin][l]=1
        state[spin][m]=0
        return state
    elif (state[spin][m]==0 and state[spin][l]==1):
        state[spin][m]=1
        state[spin][l]=0
        return state
    else:
        return []

#Try using hoptst:
print(hoptst(0,1,1,psi_j))

### Function to permutate a given list

# Python function to print permutations of a given list 
def permutation(lst): 
  
    # If lst is empty then there are no permutations 
    if len(lst) == 0: 
        return [] 
  
    # If there is only one element in lst then, only 
    # one permuatation is possible 
    if len(lst) == 1: 
        return [lst] 
  
    # Find the permutations for lst if there are 
    # more than 1 characters 
  
    l = [] # empty list that will store current permutation 
  
    # Iterate the input(lst) and calculate the permutation 
    for i in range(len(lst)): 
        m = lst[i] 
  
       # Extract lst[i] or m from the list.  remLst is 
       # remaining list 
        remLst = lst[:i] + lst[i+1:] 
  
       # Generating all permutations where m is first 
       # element 
        for p in permutation(remLst): 
            l.append([m] + p) 
    return l 
  
  
# Driver program to test above function 
data = list('1000') 
for p in permutation(data): 
    print(p)

    

import itertools
items = [1,0,0]
perms = itertools.permutations

tst = 15.2
tst = np.full(3, tst)
print(tst)

H = np.array([[1, 1], [1, -1]])
H_2 = np.tensordot(H,H, 0)
print(H_2)
print('==============================')
M = np.array([[0,1,1,0],[1,1,0,0],[1,0,1,0],[0,0,0,1]])
print(M)
print(np.dot(np.conj(M.T),M))

