# imports
import numpy as np
from typing import List, Callable
from scipy.optimize import minimize
from scipy.optimize._optimize import OptimizeResult
import matplotlib.pyplot as plt

from qiskit import QuantumCircuit
from qiskit.quantum_info import Statevector, Operator, SparsePauliOp
from qiskit.primitives import StatevectorSampler, PrimitiveJob
from qiskit.circuit.library import TwoLocal
from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager
from qiskit.visualization import plot_histogram
from qiskit_ibm_runtime.fake_provider import FakeSherbrooke
from qiskit_ibm_runtime import Session, EstimatorV2 as Estimator
from qiskit_aer import AerSimulator

# Setup the grader
from qc_grader.challenges.iqc_2024 import (
    grade_lab1_ex1,
    grade_lab1_ex2,
    grade_lab1_ex3,
    grade_lab1_ex4,
    grade_lab1_ex5,
    grade_lab1_ex6,
    grade_lab1_ex7,
)

# Build a circuit to form a psi-minus Bell state
# Apply gates to the provided QuantumCircuit, qc
qc = QuantumCircuit(2)
### Write your code below here ###
qc.h(0)
qc.cx(0,1)
qc.z(1)
qc.x(1)
### Don't change any code past this line ###
qc.measure_all()
qc.draw('mpl')

# Submit your answer using following code

grade_lab1_ex1(qc) # Expected result type: QuantumCircuit

qc.measure_all()

### Write your code below here ###
sampler = StatevectorSampler() #Add your code here
pub = (qc)#Add your code here
job_sampler = sampler.run([pub])#Add your code here
### Don't change any code past this line ###
result_sampler = job_sampler.result()
counts_sampler = result_sampler[0].data.meas.get_counts()
print(counts_sampler)

# Submit your answer using following code

grade_lab1_ex2(job_sampler) # Expected result type: PrimitiveJob

plot_histogram(counts_sampler)

# Step 1
qc = QuantumCircuit(3) #your_code_here

# Step 2 (provided)
qc.ry(1.91063324, 0)

# Add steps 3-6 below
qc.ch(0,1)
qc.cx(1,2)
qc.cx(0,1)
qc.x(0)

### Don't change any code past this line ###
qc.measure_all()
qc.draw('mpl')

# Submit your answer using following code

grade_lab1_ex3(qc) # Expected result type: # Expected result type: QuantumCircuit

sampler = StatevectorSampler()
pub = (qc)
job_sampler = sampler.run([pub], shots=10000)

result_sampler = job_sampler.result()
counts_sampler = result_sampler[0].data.meas.get_counts()

print(counts_sampler)

plot_histogram(counts_sampler)

pauli_op = SparsePauliOp(['ZII', 'IZI', 'IIZ'])
print(pauli_op.to_matrix())

num_qubits = 3#Add your code here
rotation_blocks = ['ry','rz'] #Add your code here
entanglement_blocks = 'cz' #Add your code here
entanglement = 'full' #Add your code here
reps= 1
insert_barriers=True

ansatz = TwoLocal(num_qubits= num_qubits, rotation_blocks= rotation_blocks, entanglement_blocks= entanglement_blocks, entanglement= entanglement, reps=reps ,  insert_barriers= insert_barriers) #Add your code here


### Don't change any code past this line ###
ansatz.decompose().draw('mpl')

# Submit your answer using following code

grade_lab1_ex4(num_qubits, rotation_blocks, entanglement_blocks, entanglement) # Expected result type: int, List[str], str, str

num_params = ansatz.num_parameters
num_params

backend_answer = FakeSherbrooke() #Add your code
optimization_level_answer = 0 #Add your code
pm = generate_preset_pass_manager(backend=backend_answer,optimization_level=optimization_level_answer)
isa_circuit = pm.run(ansatz) # Add your code

# Submit your answer using following code

grade_lab1_ex5(isa_circuit) # Expected result type: QuantumCircuit

isa_circuit.draw('mpl', idle_wires=False,)

# Define our Hamiltonian
hamiltonian_isa = pauli_op.apply_layout(layout=isa_circuit.layout)

def cost_func(params, ansatz, hamiltonian, estimator, callback_dict):
    """Return estimate of energy from estimator

    Parameters:
        params (ndarray): Array of ansatz parameters
        ansatz (QuantumCircuit): Parameterized ansatz circuit
        hamiltonian (SparsePauliOp): Operator representation of Hamiltonian
        estimator (EstimatorV2): Estimator primitive instance

    Returns:
        float: Energy estimate
    """
    pub = (ansatz, [hamiltonian], [params]) #Add your code
    result = estimator.run(pubs=[pub]).result() #Add your code
    energy = result[0].data.evs[0] #Add your code

    callback_dict["iters"] += 1 #Add your code
    callback_dict["prev_vector"] = params #Add your code
    callback_dict["cost_history"].append(energy) #Add your code


### Don't change any code past this line ###
    print(energy)
    return energy, result

# Submit your answer using following code

grade_lab1_ex6(cost_func) # Expected result type: Callable

callback_dict = {
    "prev_vector": None,
    "iters": 0,
    "cost_history": [],
}

x0 = 2 * np.pi * np.random.random(num_params)
x0

### Select a Backend
## Use FakeSherbrooke to simulate with noise that matches closer to the real experiment. This will run slower.
## Use AerSimulator to simulate without noise to quickly iterate. This will run faster.

backend = FakeSherbrooke()
# backend = AerSimulator()

# ### Don't change any code past this line ###

# Here we have updated the cost function to return only the energy to be compatible with recent scipy versions (>=1.10)
def cost_func_2(*args, **kwargs):
    energy, result = cost_func(*args, **kwargs)
    return energy

with Session(backend=backend) as session:
    estimator = Estimator(session=session)

    res = minimize(
        cost_func_2,
        x0,
        args=(isa_circuit, hamiltonian_isa, estimator, callback_dict),
        method="cobyla",
        options={'maxiter': 30})

# Submit your answer using following code

grade_lab1_ex7(res) # Expected result type: OptimizeResult

fig, ax = plt.subplots()
plt.plot(range(callback_dict["iters"]), callback_dict["cost_history"],'--',color = 'k')
plt.xlabel("Energy")
plt.ylabel("Cost")
plt.draw()