# qc-grader should be 0.18.11 (or higher)
import qc_grader

qc_grader.__version__

# Import all in one cell

import numpy as np
import matplotlib.pyplot as plt
from timeit import default_timer as timer
import warnings

from qiskit import QuantumCircuit
from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager
from qiskit.circuit.random import random_circuit
from qiskit.quantum_info import SparsePauliOp
from qiskit.circuit.library import TwoLocal, EfficientSU2
from qiskit_ibm_runtime import QiskitRuntimeService, Estimator, Session, Options
from qiskit_serverless import QiskitFunction, save_result, get_arguments, save_result, distribute_task, distribute_qiskit_function, IBMServerlessClient, QiskitFunction
from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager
from qiskit_transpiler_service.transpiler_service import TranspilerService
from qiskit_aer import AerSimulator

# Import for grader
from qc_grader.challenges.iqc_2024 import grade_lab3_qs_ex1, grade_lab3_qs_ex2

service = QiskitRuntimeService(channel="ibm_quantum")
 
# Specify a system to use for transpilation
real_backend = service.backend("ibm_brisbane")

# Qiskit Pattern Step 1: Map quantum circuits and operators (Define Ansatz and operators)
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

# Define Ansatz
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

# Define parameters
num_params = ansatz.num_parameters

# Qiskit Pattern Step 2: Optimize the circuit for quantum execution
optimization_level = 2
pm = generate_preset_pass_manager(backend=real_backend, optimization_level=optimization_level)
isa_circuit = pm.run(ansatz) # Add your code here

# Define Hamiltonian for VQE
pauli_op = SparsePauliOp(['ZII', 'IZI', 'IIZ'])
hamiltonian_isa = pauli_op.apply_layout(layout=isa_circuit.layout)

# Setup Qiskit Serverless Client and Qiskit Runtime client
client = IBMServerlessClient("hidden") # Add in your IBM Quantum Token to QiskitServerless Client

# For the challenge, we will be using QiskitRuntime Local testing mode. Change to True only if you wish to use real backend.
USE_RUNTIME_SERVICE = False

if USE_RUNTIME_SERVICE:
    service = QiskitRuntimeService(
        channel='ibm_quantum', 
        verify=False
    )
else:
    service = None

# Define the Qiskit Function
if USE_RUNTIME_SERVICE:
    function = QiskitFunction(title= "vqe", entrypoint="vqe.py", working_dir="./vqe")
else:
    function = QiskitFunction(title= "vqe" , entrypoint="vqe.py", working_dir="./vqe",  dependencies=["qiskit_aer"])


# Upload the Qiskit Function using IBMServerlessClient
client.upload(function)

# Define input_arguments
input_arguments = {
    "ansatz": isa_circuit, # Replace with your transpiled ansatz
    "operator": hamiltonian_isa, # Replace with the hamiltonian operator
    "method": "COBYLA", # Using COBYLA method for the optimizer
    "service": service, # Add your code here
}

# Qiskit Pattern Step 3: Run the payload on backend
job = client.run("vqe", arguments= input_arguments) # Pass the arguments dict here)

# Submit your answer using following code

grade_lab3_qs_ex1(function, input_arguments, job)
# Expected result type: QiskitFunction, dict, Job

# Return jobid
job

# Check job completion status
job.status()

# Monitor log
logs = job.logs()

for log in logs.splitlines():
    print(log)

# Return result from QiskitFunction job
job.result()

# Qiskit Pattern Step 4: Postprocess and analyze the Estimator V2 results
result = job.result()

plt.style.use('ggplot')
fig, ax = plt.subplots()
plt.plot(range(result["iters"]), result["cost_history"], ls='--', c='k')
plt.xlabel("Energy")
plt.ylabel("Cost")
plt.draw()

# Setup 3 circuits with Efficient SU2
num_qubits = [41, 51, 61]
circuits = [EfficientSU2(nq, su2_gates=["rz","ry"], entanglement="circular", reps=1).decompose() for nq in num_qubits]

# Setup Qiskit Runtime Service backend

# QiskitRuntimeService.save_account(
#     channel="ibm_quantum",
#     token="YOUR_IBM_QUANTUM_TOKEN",
#     set_as_default=True,
#     # Use 'overwrite=True' if you're updating your token.
#     overwrite=True,
# )

service = QiskitRuntimeService(channel="ibm_quantum")
backend = service.backend("ibm_brisbane")

# Define Configs
optimization_levels = [1,2,3] # Add your code here
pass_managers = [{'pass_manager': generate_preset_pass_manager(optimization_level=level, backend=backend), 'optimization_level': level} for level in optimization_levels]

transpiler_services = [ 
        {'service': TranspilerService(backend_name="ibm_brisbane", ai="false", optimization_level=3,), 'ai': False, 'optimization_level': 3},
        {'service': TranspilerService(backend_name="ibm_brisbane", ai="true", optimization_level=3,), 'ai': True, 'optimization_level': 3}
    ]

configs = pass_managers + transpiler_services

# Local transpilation setup
def transpile_parallel_local(circuit: QuantumCircuit, config):
    """Transpilation for an abstract circuit into an ISA circuit for a given config."""
    transpiled_circuit = config.run(circuit)
    return transpiled_circuit

# Run local transpilation
warnings.filterwarnings("ignore")

start = timer()

# Run transpilations locally for baseline
results = []
for circuit in circuits:
    for config in configs:
        if 'pass_manager' in config:
            results.append(transpile_parallel_local(circuit, config['pass_manager']))
        else:
            results.append(transpile_parallel_local(circuit, config['service']))
            
end = timer()

# Record local execution time
execution_time_local = end - start
print("Execution time locally: ", execution_time_local)

# Authenticate to the remote cluster and submit the pattern for remote execution if not done in previous exercise
serverless = IBMServerlessClient("hidden")

transpile_parallel_function = QiskitFunction(
    title="transpile_parallel",
    entrypoint="transpile_parallel.py",
    working_dir="./transpile_parallel",
    dependencies=["qiskit-transpiler-service"]
)

serverless.upload(transpile_parallel_function) # Add your code here

# Get list of functions
serverless.list()

# Fetch the specific function titled "transpile_parallel"
transpile_parallel_serverless = serverless.get("transpile_parallel") # Add your code here

# Run the "transpile_parallel" function in the serverless environment
job = transpile_parallel_serverless.run(
    circuits= circuits, # Add your code here,
    backend_name= backend # Add your code here
                                       )

# Submit your answer using following code

grade_lab3_qs_ex2(optimization_levels, transpiler_services, transpile_parallel_function, transpile_parallel_serverless, job)
# Expected result type: list, list, QiskitFunction, QiskitFunction, Job

job.status()

logs = job.logs()

for log in logs.splitlines():
    print(log)

result = job.result()

result_transpiled = result["transpiled_circuits"]
# Compare execution times:
execution_time_serverless = result["execution_time"]

from utils import plot_execution_times
plot_execution_times(execution_time_serverless, execution_time_local)

from utils import process_transpiled_circuits

best_circuits, best_depths, best_methods = process_transpiled_circuits(configs, result_transpiled)

# Display the best circuits, depths, and methods
for i, (circuit, depth, method) in enumerate(zip(best_circuits, best_depths, best_methods)):
    print(f"Best result for circuit {i + 1}:")
    print(f"  Depth: {depth}")
    print(f"  Method: {method}")
    # Display or process the best circuit as needed
    # e.g., circuit.draw(output="mpl")