%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 FullClassicalHubbardEvolutionChain as fhc

import random as rand
import scipy.linalg as la



#System parameters
numsites = 3
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)
states_list = fhc.get_states(numsites)
#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, wfks = fhc.sys_evolve(states_list, wfk_full, t, U, classical_total_time, classical_time_step)
print(wfks[10])



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()


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)


#==== Create all Possible States for given system size ===#
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)


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
        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]:
        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


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)
    wavefunctions = []
    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
    wavefunctions.append(np.ndarray.tolist(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.
        for j in mapping[i]:
            wfk_sum += evolve[0][j]
            norm += evolve[0][j]
        if norm == 0.:
            norm = 1.
        mode_evolve[0][i] = wfk_sum #/ norm
        
        #print('wfk_sum: ',wfk_sum,'    norm: ',norm)
        energies[0] = wfk_energy(wfk, hamiltonian)
        #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)
        wavefunctions.append(np.ndarray.tolist(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]
                norm += evolve[t][j]
            #print('wfk_sum: ',wfk_sum,'    norm: ',norm)
            if norm == 0.:
                norm = 1.
            mode_evolve[t][i] = wfk_sum #/ norm
        #print(mode_evolve[t])

    #Return time evolution
    return mode_evolve, energies, wavefunctions

#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, wfks = sys_evolve(states_list, wfk_full, t, U, classical_total_time, classical_time_step)
print(wfks[10])

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()

