import numpy as np
import matplotlib.pyplot as plt

# Importing standard Qiskit libraries
from qiskit import QuantumCircuit, QuantumRegister, transpile, Aer, IBMQ, execute
from qiskit.tools.jupyter import *
from qiskit.visualization import *
from qiskit.providers.aer import QasmSimulator
from tqdm.notebook import tqdm

from qiskit.providers.aer import QasmSimulator
from qiskit.tools.monitor import job_monitor
from qiskit.circuit import Parameter
import qiskit.quantum_info as qi

from qc_grader.challenges.spring_2022.helpers import generate_disordered_tb_instruction

# Suppress warnings
import warnings
warnings.filterwarnings('ignore')

import pylatexenc
import IPython

pop = [1,0]
Ne=0  # Sum of even populations.
No=0 # Sum of odd populations.
for j in range(0, len(pop), 2):
    Ne += pop[j]
for j in range(1, len(pop), 2):
    No += pop[j]
    
I = (Ne-No)/(Ne+No)
print(f"Ne={Ne}, No={No}, I={I}")

def get_imbalance(state):
    ###EDIT CODE BELOW
    ### HINT: MAKE SURE TO SKIP CALCULATING IMBALANCE OF THE |00...0> STATE
    imbalance_val=0
    
    def get_basis(i):
        num_qubits = int(np.log2(len(state.probabilities())))
        basis = []
        for i in list(f'{i:0{num_qubits}b}'):
          basis.append(int(i))
        return basis # A list such as [1,0,0,0]

    def imbalance(basis):   
        Ne=0  # Sum of even populations.
        No=0 # Sum of odd populations.
        for j in range(0, len(basis), 2):
            Ne += basis[j]
        for j in range(1, len(basis), 2):
            No += basis[j] 
        return (Ne-No)/(Ne+No)    
    
    pop = state.probabilities() # List of probabilities of the state |\psi> for measuring all qubits.
    for i in range(1, len(pop)): # Skipping i=0 corresponding to the |00...0> basis.
        basis = get_basis(i)
        prob_amp = pop[i] # float
        imbalance_val += imbalance(basis)*prob_amp
            
    ###DO NOT EDIT BELOW
    
    return imbalance_val

st1 = qi.Statevector.from_label('0111') # -1/3
st2 = np.sqrt(2/3)*qi.Statevector.from_label('0111') + np.sqrt(1/3)*qi.Statevector.from_label('1011') # -1/9
st3 = (qi.Statevector.from_label('1000')+qi.Statevector.from_label('1001'))*(1/np.sqrt(2)) # 0.5

for i in [st1, st2, st3]:
    print(get_imbalance(i))

## Grade and submit your solution
from qc_grader.challenges.spring_2022 import grade_ex3a

grade_ex3a(get_imbalance) 


bell_state = qi.Statevector(np.array([0,1,1,0])/np.sqrt(2))

rho_0 = qi.partial_trace(bell_state,[1]) # We trace out qubit 1
rho_1 = qi.partial_trace(bell_state,[0]) # We trace out qubit 0

print('QB0 vn entropy: ', qi.entropy(rho_0, base=np.exp(1)))
print('QB1 vn entropy: ', qi.entropy(rho_1, base=np.exp(1)))

t = Parameter('t')

num_qubits=12
deltas=[Parameter('delta_{:d}'.format(idx)) for idx in range(num_qubits)]
disorder_trot_step=generate_disordered_tb_instruction(t, deltas, num_qubits)

# Here we define the disorder pattern

beta=(np.sqrt(5)-1)/2 # DO NOT EDIT
AA_pattern=np.cos(2*np.pi*beta*np.arange(num_qubits)) # DO NOT EDIT

delta_t=0.1
time_steps=np.arange(0,21,2)

circuits={}
Ws=[1,4,10]

for W in Ws:
    disorders=W*AA_pattern
    
    circuits[W]=[]

    for n_steps in time_steps:

        qr = QuantumRegister(num_qubits)
        qc = QuantumCircuit(qr)

        ###EDIT CODE BELOW
        
        qc.x([0,4,8])
        
        ###DO NOT EDIT BELOW     

        for _ in range(n_steps):
            qc.append(disorder_trot_step, [i for i in range(num_qubits)])
        
        if n_steps!=0:
            qc = qc.bind_parameters({t: delta_t})
            qc = qc.bind_parameters({deltas[idx]: disorders[idx] for idx in range(num_qubits)})

        circuits[W].append(qc)

from qiskit import transpile

# Use Aer's statevector simulator
from qiskit import Aer

# Run the quantum circuit on a statevector simulator backend
backend_sim = Aer.get_backend('statevector_simulator')

probability_densities={}
state_vector_imbalances={}
vn_entropies={}

for W in tqdm(Ws):
    probability_densities[W]=[]
    state_vector_imbalances[W]=[]
    vn_entropies[W]=[]
    
    for circ in circuits[W]:

        transpiled_circ=transpile(circ, backend_sim, optimization_level=3)

        job_sim = backend_sim.run(transpiled_circ)

        # Grab the results from the job.
        result_sim = job_sim.result()
        outputstate = result_sim.get_statevector(transpiled_circ, decimals=6)
        ps=[]
        for idx in range(num_qubits):
            ps.append(np.abs(qi.partial_trace(outputstate,[i for i in range(num_qubits) if i!=idx]))[1,1]**2)
        
        entropy=0
        ### EDIT CODE BELOW (extract the density matrix of qubit 0 by tracing out all other qubits)
        rho0=qi.partial_trace(outputstate, [i for i in range(num_qubits) if i!=0]) # trace out qubit 0.       
        entropy=qi.entropy(rho0, base=np.exp(1))
        
        ###DO NOT EDIT BELOW 
        
        imbalance=get_imbalance(outputstate)
        ### EDIT CODE BELOW 
        
        
        ###DO NOT EDIT BELOW
        
        
        vn_entropies[W].append(entropy)
        probability_densities[W].append(ps)
        state_vector_imbalances[W].append(imbalance)

fig, axs = plt.subplots(1,3,figsize=(15,5), facecolor='white', sharey=True)

for i,W in enumerate(Ws):
    ax=axs[i]
    ax.pcolormesh(np.arange(0,num_qubits,1), time_steps*delta_t ,probability_densities[W])
    ax.set_xlabel('Qubit index')
    ax.set_xticks(np.arange(1,num_qubits+1,1))

axs[0].set_ylabel('Time (1/J)')

plt.show()

for W in Ws:
    plt.plot(time_steps*delta_t,vn_entropies[W], '--o', label='W={:d}'.format(W))

plt.xlabel(r'Time (1/J)')
plt.ylabel(r'$\mathcal{S}_{\rm vn}(\rho_0)$')
plt.legend()
plt.show()

## Grade and submit your solution
from qc_grader.challenges.spring_2022 import grade_ex3b

grade_ex3b(vn_entropies)


for W in Ws:
    plt.plot(time_steps*delta_t,state_vector_imbalances[W], '--o', label='W={:d}'.format(W))

plt.xlabel(r'Time (1/J)')
plt.ylabel(r'$\mathcal{I}$')
plt.legend()
plt.show()

## Grade and submit your solution
from qc_grader.challenges.spring_2022 import grade_ex3c

grade_ex3c(state_vector_imbalances)
