%matplotlib inline
# Importing standard Qiskit libraries and configuring account
from qiskit import QuantumCircuit, execute, Aer, IBMQ, BasicAer, QuantumRegister, ClassicalRegister
from qiskit.compiler import transpile, assemble
from qiskit.quantum_info import Operator
from qiskit.tools.monitor import job_monitor
from qiskit.tools.jupyter import *
from qiskit.visualization import *
import random as rand
import scipy.linalg as la

provider = IBMQ.load_account()

import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import numpy as np
from matplotlib import rcParams
rcParams['text.usetex'] = True

def reverse_list(s):
    temp_list = list(s)
    temp_list.reverse()
    return ''.join(temp_list)


#Useful tool for converting an integer to a binary bit string
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)
    #return ''.join(sup)+''.join(sdn)




'''The task here is now to define a function which will either update a given circuit with a time-step
or return a single gate which contains all the necessary components of a time-step'''




#==========Needed Functions=============#
#Function to apply a full set of time evolution gates to a given circuit
def qc_evolve(qc, numsite, time, hop, U, trotter_steps):
    #Compute angles for the onsite and hopping gates
    # based on the model parameters t, U, and dt
    theta = hop*time/(2*trotter_slices) 
    phi = U*time/(trotter_slices)
    numq = 2*numsite
    y_hop = Operator([[np.cos(theta), 0, 0, -1j*np.sin(theta)],
                [0, np.cos(theta), 1j*np.sin(theta), 0],
                [0, 1j*np.sin(theta), np.cos(theta), 0],
                [-1j*np.sin(theta), 0, 0, np.cos(theta)]])
    x_hop = Operator([[np.cos(theta), 0, 0, 1j*np.sin(theta)],
                [0, np.cos(theta), 1j*np.sin(theta), 0],
                [0, 1j*np.sin(theta), np.cos(theta), 0],
                [1j*np.sin(theta), 0, 0, np.cos(theta)]])
    z_onsite = Operator([[1, 0, 0, 0],
               [0, 1, 0, 0],
               [0, 0, 1, 0],
               [0, 0, 0, np.exp(1j*phi)]])
    
    #Loop over each time step needed and apply onsite and hopping gates
    for trot in range(trotter_steps):
        #Onsite Terms
        for i in range(0, numsite):
            qc.unitary(z_onsite, [i,i+numsite], label="Z_Onsite")
        
            #Add barrier to separate onsite from hopping terms    
            qc.barrier()

        #Hopping terms
        for i in range(0,numsite-1):
            #Spin-up chain
            qc.unitary(y_hop, [i,i+1], label="YHop")
            qc.unitary(x_hop, [i,i+1], label="Xhop")
            #Spin-down chain
            qc.unitary(y_hop, [i+numsite, i+1+numsite], label="Xhop")
            qc.unitary(x_hop, [i+numsite, i+1+numsite], label="Xhop")

            #Add barrier after finishing the time step
            qc.barrier()
    #Measure the circuit
    for i in range(numq):
        qc.measure(i, i)
        


#Function to run the circuit and store the counts for an evolution with
# num_steps number of time steps.
def sys_evolve(nsites, excitations, total_time, dt, hop, U, trotter_steps):
    #Check for correct data types
    if not isinstance(nsites, int):
        raise TypeError("Number of sites should be int")
    if np.isscalar(excitations):
        raise TypeError("Initial state should be list or numpy array")
    if not np.isscalar(total_time):
        raise TypeError("Evolution time should be scalar")
    if not np.isscalar(dt):
        raise TypeError("Time step should be scalar")
    if not np.isscalar(hop):
        raise TypeError("Hopping term should be scalar")
    if not np.isscalar(U):
        raise TypeError("Repulsion term should be scalar")
    if not isinstance(trotter_steps, int):
        raise TypeError("Number of trotter slices should be int")
        
    numq = 2*nsites
    num_steps = int(total_time/dt)
    print('Num Steps: ',num_steps)
    print('Total Time: ', total_time)
    data = np.zeros((2**numq, num_steps))
    
    for t_step in range(0, num_steps):
        #Create circuit with t_step number of steps
        q = QuantumRegister(numq)
        c = ClassicalRegister(numq)
        qcirc = QuantumCircuit(q,c)

        #=========USE THIS REGION TO SET YOUR INITIAL STATE==============
        #Initialize circuit by setting the occupation to 
        # a spin up and down electron in the middle site
        #qcirc.x(int(nsites/2))
        #qcirc.x(nsites+int(nsites/2))
        
        for flip in excitations:
            qcirc.x(flip)
        #if nsites==3:
            #Half-filling
            #qcirc.x(1)
        #    qcirc.x(4)
        #    qcirc.x(0)
        #    qcirc.x(2)
            #1 electron
        #    qcirc.x(1)        
        #===============================================================
    
        qcirc.barrier()
        #Append circuit with Trotter steps needed
        qc_evolve(qcirc, nsites, t_step*dt, hop, U, trotter_steps)
        
    
    #Choose provider and backend
        provider = IBMQ.get_provider()
        #backend = Aer.get_backend('statevector_simulator')
        backend = Aer.get_backend('qasm_simulator')
        #backend = provider.get_backend('ibmq_qasm_simulator')
        #backend = provider.get_backend('ibmqx4')
        #backend = provider.get_backend('ibmqx2')
        #backend = provider.get_backend('ibmq_16_melbourne')

        shots = 8192
        max_credits = 10 #Max number of credits to spend on execution
        job_exp = execute(qcirc, backend=backend, shots=shots, max_credits=max_credits)
        job_monitor(job_exp)
        result = job_exp.result()
        counts = result.get_counts(qcirc)
        print(result.get_counts(qcirc))
        print("Job: ",t_step+1, " of ", num_steps," complete.")
    
    #Store results in data array and normalize them
        for i in range(2**numq):
            if counts.get(get_bin(i,numq)) is None:
                dat = 0
            else:
                dat = counts.get(get_bin(i,numq))
            data[i,t_step] = dat/shots
    return data



#==========Set Parameters of the System=============#
dt = 0.25 #Delta t
T = 4.5
time_steps = int(T/dt)
t = 1.0  #Hopping parameter
U = 2.  #On-Site repulsion
#time_steps = 10
nsites = 3
trotter_slices = 5
initial_state = np.array([1,4])

#Run simulation
run_results = sys_evolve(nsites, initial_state, T, dt, t, U, trotter_slices)
#print(True if np.isscalar(initial_state) else False)


#Process and plot data
'''The procedure here is, for each fermionic mode, add the probability of every state containing
that mode (at a given time step), and renormalize the data based on the total occupation of each mode.
Afterwards, plot the data as a function of time step for each mode.'''
proc_data = np.zeros((2*nsites, time_steps))
timesq = np.arange(0.,time_steps*dt, dt)

#Sum over time steps
for t in range(time_steps):
    #Sum over all possible states of computer
    for i in range(2**(2*nsites)):
        #num = get_bin(i, 2*nsite)
        num = ''.join( list( reversed(get_bin(i,2*nsites)) ) )
        #For each state, check which mode(s) it contains and add them
        for mode in range(len(num)):
            if num[mode]=='1':
                proc_data[mode,t] += run_results[i,t]
    
    #Renormalize these sums so that the total occupation of the modes is 1
    norm = 0.0
    for mode in range(len(num)):
        norm += proc_data[mode,t]
    proc_data[:,t] = proc_data[:,t] / norm
    
'''
At this point, proc_data is a 2d array containing the occupation 
of each mode, for every time step
'''
    
#Create plots of the processed data
fig2, ax2 = plt.subplots(figsize=(20,10))
colors = list(mcolors.TABLEAU_COLORS.keys())

for i in range(nsites):
    #Create string label
    strup = "Site "+str(i+1)+r'$\uparrow$'
    strdwn = "Site "+str(i+1)+r'$\downarrow$'
    ax2.plot(timesq, proc_data[i,:], marker="^", color=str(colors[i]), label=strup)
    ax2.plot(timesq, proc_data[i+nsites,:], marker="v", color=str(colors[i]), label=strdwn)


#ax2.set_ylim(0, 0.55)
ax2.set_xlim(0, time_steps*dt+dt/2.)
#ax2.set_xticks(np.arange(0,time_steps*dt+dt, 0.2))
#ax2.set_yticks(np.arange(0,0.55, 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)


#Plot the raw data as a colormap
xticks = np.arange(2**(nsite*2))
xlabels=[]
print("Time Steps: ",time_steps, " Step Size: ",dt)
for i in range(2**(nsite*2)):
    xlabels.append(get_bin(i,6))
    
fig, ax = plt.subplots(figsize=(10,20))
c = ax.pcolor(run_results, cmap='binary')
ax.set_title('Time Evolution of 3 Site One Dimensional Chain', fontsize=22)
plt.yticks(xticks, xlabels, size=18)
ax.set_xlabel('Time Step', fontsize=22)
ax.set_ylabel('State', fontsize=26)
plt.show()

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

t = 1.0
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()

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.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., 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


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

#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.
#'''
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)
    ''' #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())


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


for i in range(nsites):
    #Create string label
    strupq = "Quantum Site "+str(i+1)+r'$\uparrow$'
    strdwnq = "Quantum Site "+str(i+1)+r'$\downarrow$'
    strup = "Numerical Site "+str(i+1)+r'$\uparrow$'
    strdwn = "Numerical Site "+str(i+1)+r'$\downarrow$'
    ax2.scatter(timesq, proc_data[i,:], marker="*", color=str(colors[i]), label=strupq)
    #ax2.scatter(timesq, proc_data[i+nsite,:], marker="v", color=str(colors[i]), label=strdwnq)
    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.plot(times, evolve[:,i], linestyle='-', color=str(colors[i]), linewidth=2, label=strup)
    #ax2.plot(times, evolve[:,i+3], linestyle='--', color=str(colors[i]), linewidth=2, label=strdwn)




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



