import numpy as np
import matplotlib.pyplot as plt
from copy import deepcopy

# Importing standard Qiskit libraries
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, 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_XX, generate_YY, generate_disordered_tb_instruction

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

import pylatexenc
import IPython

t = Parameter('t')

XX = generate_XX(t)
YY = generate_YY(t)

from scipy.special import jv

t_test = 20
r = np.linspace(-50,50,101)
gaussian = np.exp(-r**2/t_test) / np.sum(np.exp(-r**2/t_test))
bessel = np.abs(jv(r,2*t_test))**2 / np.sum(np.abs(jv(r,2*t_test))**2)

plt.figure(facecolor='white')
plt.plot(r, gaussian, label=r'Gaussian function')
plt.plot(r, bessel, label=r'Bessel function')

plt.xlabel('Position')
plt.ylabel('Probability density')
plt.legend()
plt.show()

num_qubits = 5 ## DO NOT EDIT

Trot_tb_qr = QuantumRegister(num_qubits)
Trot_tb_qc = QuantumCircuit(Trot_tb_qr, name='Trot')

###EDIT CODE BELOW
for j in [0,2]:
    Trot_tb_qc.append(YY, [Trot_tb_qr[j], Trot_tb_qr[j+1]])
    Trot_tb_qc.append(XX, [Trot_tb_qr[j], Trot_tb_qr[j+1]]) 
    
for i in [1,3]:
    Trot_tb_qc.append(YY, [Trot_tb_qr[i], Trot_tb_qr[i+1]])
    Trot_tb_qc.append(XX, [Trot_tb_qr[i], Trot_tb_qr[i+1]])


###DO NOT EDIT BELOW

Trot_tb_gate = Trot_tb_qc.to_instruction()

Trot_tb_qc.draw(output='mpl')

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

grade_ex2a(Trot_tb_qc)


delta_t=0.15 # DO NOT EDIT
time_steps=np.arange(1,20,1) # DO NOT EDIT

circuits=[]

for n_steps in time_steps:
    
    qr = QuantumRegister(num_qubits)
    cr = ClassicalRegister(num_qubits)
    qc = QuantumCircuit(qr,cr)

    ###EDIT CODE BELOW
    qc.x([0])
    ###DO NOT EDIT BELOW
    
    for _ in range(n_steps):
        qc.append(Trot_tb_gate, [i for i in range(num_qubits)])
        
    qc = qc.bind_parameters({t: delta_t})
    
    circuits.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_density=[]
for circ in tqdm(circuits):

    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=5)
    
    ps=[]
    
    ###EDIT CODE BELOW (Extract the probability of finding the excitation on each qubit)
    probs=[]
    for i in range(num_qubits):
        prob=outputstate.probabilities([i]) # type(..) >> list
        probs.append(prob)
        
    for elem in probs:
        excited=elem[1]
        ps.append(excited)
        
    
    ###DO NOT EDIT BELOW
    
    probability_density.append(ps)
    
probability_density=np.array(probability_density)

plt.figure(figsize=(3,5), facecolor='white')
plt.pcolormesh(np.arange(0,num_qubits,1), time_steps*delta_t, probability_density)
plt.xlabel('Qubit index')
plt.ylabel('Time (1/J)')

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

grade_ex2b(probability_density)


from qiskit.tools.monitor import job_monitor
from qiskit import IBMQ

# Loading your IBM Quantum account
IBMQ.load_account()
IBMQ.providers() # see a list of providers you have access to

# Get the special provider assigned to you using information from the output above
hub_name = 'qc-spring-22-4' # e.g. 'qc-spring-22-1'
group_name = 'group-2' # e.g. 'group-2'
project_name = 'recrrGhzxg3VULGQS' # e.g. 'recPdHH04wfWiNHjG'
provider = IBMQ.get_provider(hub=hub_name, group=group_name, project=project_name)

# Get `ibm_nairobi` backend from the provider
backend_name = 'ibm_nairobi'
backend = provider.get_backend(backend_name)
backend # See details of the `ibm_nairobi` quantum system

initial_layout=[0 , 1 , 3 , 5 , 4]

hardware_transpiled_circuits=[]
for circ in circuits:
    
    hardware_circ=deepcopy(circ)
    hardware_circ.barrier()
    hardware_circ.measure(range(num_qubits), range(num_qubits))
    hardware_transpiled_circuits.append(transpile(hardware_circ, backend, initial_layout=initial_layout, optimization_level=3))

shots=1024
job = execute(hardware_transpiled_circuits, backend=backend, shots=shots)
job_monitor(job)
print('Job ID', job.job_id())

qc = QuantumCircuit(5,1)
qc.x([0])
qc.measure(0,0) # Measuring only qubit 0 (q_0) and storing the measurement result in classical bit (a.k.a register) 0.
jobs = execute(qc,backend,shots=1024) # Sending job to ibm_nairobi.
counts = jobs.result().get_counts(qc)
probs = {key:value/shots for key, value in counts.items()}

print(qc)
print(counts)
print(probs)

experiment_results=backend.retrieve_job('629c26b701885c41f6f358c2').result()

i=0
for output in experiment_results.get_counts(): # Doing a for through each measurement (each new experiment).
    print(f'output: {type(output)}\n')
    print(output, '\n')
    #print(output.keys()) # Type >> <class 'dict_keys'>
    print(f"Number of Basis states/prob. amp.: {len(output)}.\n")
    i+=1
    if i == 2:
        break

type(experiment_results.get_counts()), len(experiment_results.get_counts()) 

# There are 19 dictionaries. Each dict. corresponds to an average measurement trial.

for output in experiment_results.get_counts():  
    for key, value in output.items():
        print(key, value)
        
    print('')    
    print(output.keys())
    break

Z = np.random.rand(2, 2)
Z

Z.shape, len(Z)

plt.pcolormesh(Z)

Z.tolist()

plt.pcolormesh(Z.tolist())

experiment_results=job.result()
#experiment_results=backend.retrieve_job('629c26b701885c41f6f358c2').result()

probability_density_exp=[]
for output in experiment_results.get_counts(): 
    ps=[]
    
    ###EDIT CODE BELOW (Extract the probability of finding the excitation on each qubit)
    
    for key in ['00001', '00010', '00100', '01000', '10000']:
        try:
            count=float(output.get(key))
            prob=(count/shots)
        except:
            prob=0
        ps.append(prob)
    
    ###DO NOT EDIT BELOW
    
    probability_density_exp.append(ps)

plt.figure(figsize=(3,5), facecolor='white')
plt.pcolormesh(np.arange(0,num_qubits,1), time_steps*delta_t, probability_density_exp)
plt.xlabel('Qubit index')
plt.ylabel('Time (1/J)')

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

plt.plot(np.linspace(-0.5,4.5,101), np.cos(2*np.pi*beta*np.linspace(-0.5,4.5,101)), '--')
plt.plot(np.arange(num_qubits), AA_patern, 'o', label=r'$\epsilon_i$')
plt.xlabel('Position')
plt.ylabel('Energy')
plt.legend()
plt.title('Aubry-Andre potential')
plt.show()

Trot_qr_disorder = QuantumRegister(num_qubits)
Trot_qc_disorder = QuantumCircuit(Trot_qr_disorder, name='Trot disorder')

Trot_qc_disorder.append(Trot_tb_gate,[0,1,2,3,4])
deltas=[Parameter('delta_{:d}'.format(idx)) for idx in range(num_qubits)]

###EDIT CODE BELOW (add a parametric disorder to each qubit)

for i in range(0, num_qubits):
    Trot_qc_disorder.rz(2*deltas[i]*t,i) # (angles, qubits).

###DO NOT EDIT BELOW
    
# Convert custom quantum circuit into a gate
Trot_disorder_gate = Trot_qc_disorder.to_instruction()

Trot_qc_disorder.draw(output='mpl')

delta_t=0.15
time_steps=np.arange(1,20,1)

W=2 # DO NOT EDIT

disorders=W*AA_patern # DO NOT EDIT

disorder_circuits=[]

for n_steps in time_steps:

    qr = QuantumRegister(num_qubits)
    cr = ClassicalRegister(num_qubits)
    qc = QuantumCircuit(qr, cr)

    qc.x(0)
    
    for _ in range(n_steps):
        qc.append(Trot_disorder_gate, [i for i in range(num_qubits)])

    qc = qc.bind_parameters({t: delta_t})
    qc = qc.bind_parameters({deltas[idx]: disorders[idx] for idx in range(num_qubits)})

    disorder_circuits.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_density_localization=[]
for circ in tqdm(disorder_circuits):

    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=5)
    
    ps=[]
    
    ###EDIT CODE BELOW (Extract the probability of finding the excitation on each qubit)
    
    probs=[]
    for i in range(num_qubits):
        prob=outputstate.probabilities([i]) # type(..) >> list
        probs.append(prob)
        
    for elem in probs:
        excited=elem[1]
        ps.append(excited)
    
    ###DO NOT EDIT BELOW
    
    probability_density_localization.append(ps)
    
probability_density_localization=np.array(probability_density_localization)

plt.figure(figsize=(3,5), facecolor='white')
plt.pcolormesh(np.arange(0,num_qubits,1), time_steps*delta_t ,probability_density_localization)
plt.xlabel('Qubit index')
plt.ylabel('Time (1/J)')

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

grade_ex2c(probability_density_localization)


from qiskit.tools.monitor import job_monitor

provider = IBMQ.load_account()
provider = IBMQ.get_provider(hub='qc-spring-22-4', group='group-2', project='recrrGhzxg3VULGQS')
backend = provider.get_backend('ibm_nairobi')

initial_layout=[0 , 1 , 3 , 5 , 4]

hardware_transpiled_circuits_disordered=[]
for circ in disorder_circuits:
    
    hardware_circ=deepcopy(circ)
    hardware_circ.barrier()
    hardware_circ.measure(range(num_qubits), range(num_qubits))
    hardware_transpiled_circuits_disordered.append(transpile(hardware_circ, backend, initial_layout=initial_layout, optimization_level=3))

shots=1024
job_disorder = execute(hardware_transpiled_circuits_disordered, backend=backend, shots=shots)
job_monitor(job_disorder)
print('Job ID', job_disorder.job_id())

disorder_experiment_results=job_disorder.result()
#disorder_experiment_results=backend.retrieve_job('6290b7ce4bb9753402efaa20').result()

disorder_probability_density_exp=[]
for output in disorder_experiment_results.get_counts():
    ps=[]
    
    ###EDIT CODE BELOW (Extract the probability of finding the excitation on each qubit)
    
    for key in ['00001', '00010', '00100', '01000', '10000']:
        try:
            count=float(output.get(key))
            prob=(count/shots)
        except:
            prob=0
        ps.append(prob)
        
    ###DO NOT EDIT BELOW
    
    disorder_probability_density_exp.append(ps)
    
plt.figure(figsize=(3,5), facecolor='white')
plt.pcolormesh(np.arange(0,num_qubits,1), time_steps*delta_t, disorder_probability_density_exp)
plt.xlabel('Qubit index')
plt.ylabel('Time (1/J)')