### Install Qiskit and relevant packages, if needed
### IMPORTANT: Make sure you are using python 3.10 or 3.11 for compatibility of the required packages
%pip install qiskit[visualization]==1.0.2
%pip install qiskit-ibm-runtime
%pip install qiskit-aer
%pip install qiskit-transpiler-service
%pip install graphviz
%pip install circuit-knitting-toolbox
%pip install qiskit-serverless -U
%pip install git+https://github.com/qiskit-community/Quantum-Challenge-Grader.git -U

### Save API Token, if needed

%set_env QXToken=c5a303670922a83ac5699dc2fea334193919a9e2fe6943debde1ac52815df5e39c6a137e05136088f3a7715e8d6dbd10097baf193df1fa22206a85931166946f


# Make sure there is no space between the equal sign
# and the beginning of your token

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

qc_grader.__version__

# Imports
import numpy as np
import matplotlib.pyplot as plt

from qiskit import QuantumCircuit
from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager
from qiskit.visualization import plot_gate_map
from qiskit.quantum_info import SparsePauliOp

from qiskit_aer import AerSimulator
from qiskit_ibm_runtime import QiskitRuntimeService, SamplerV2 as Sampler
from circuit_knitting.cutting import generate_cutting_experiments, cut_gates

# Setup the grader

from qc_grader.challenges.iqc_2024 import grade_lab3_ckt_ex1, grade_lab3_ckt_ex2

# create a bell pair
bell_state = QuantumCircuit(2)
bell_state.h(0)
bell_state.cx(0,1)
bell_state.draw("mpl")

## If this is your first time accessing the backend
## remove # and fill your API key, and run the code
#service = QiskitRuntimeService(
#    channel='ibm_quantum',
#    instance='ibm-q/open/main',
#    token='<IBM Quantum API key>'
#)

service = QiskitRuntimeService(channel="ibm_quantum")

# Specify a system to use for transpilation, DO NOT change
backend = service.backend("ibm_kyoto")

layout=[122, 126]
qubit_color = []
for i in range(127):
    if i in layout:
        qubit_color.append("#ff0066")
    else:
        qubit_color.append("#6600cc")
plot_gate_map(backend, qubit_color=qubit_color, qubit_size=60, font_size=25, figsize=(8,8))

# transpile the circuit
pm = generate_preset_pass_manager(backend=backend, optimization_level=1, initial_layout=layout, seed_transpiler=0)
isa_qc = pm.run(bell_state)

# original circuit depth
isa_qc_depth = isa_qc.depth()
print(f"Transpiled circuit depth: ", isa_qc_depth)
isa_qc.draw("mpl", scale=0.6, idle_wires=False, fold=False)

# Find the indices of the distant gates
cut_indices = [
    i
    for i, instruction in enumerate(bell_state.data)
    if {bell_state.find_bit(q)[0] for q in instruction.qubits} == {0, 1}
]

# Decompose distant CNOTs into TwoQubitQPDGate instances
qpd_circuit, bases = cut_gates(bell_state, cut_indices)

qpd_circuit.draw("mpl", scale=0.6)

observable = SparsePauliOp(["ZI"])

# Generate the sub-experiments and sampling coefficients
sub_experiments, coefficients = generate_cutting_experiments(
    circuits=qpd_circuit, observables=observable.paulis, num_samples=np.inf
)

# Transpile the circuit
pm = generate_preset_pass_manager(backend=backend, optimization_level=1, initial_layout=layout, seed_transpiler=0)
isa_qpd_circuit = pm.run(sub_experiments[5])

# depth using circuit cutting
isa_qpd_depth = isa_qpd_circuit.depth()

print(f"Original circuit depth after transpile: ", isa_qc_depth)
print(f"QPD sub-experiment depth after transpile: ", isa_qpd_depth)
print(f"Number of sub-experiments:", len(sub_experiments))
isa_qpd_circuit.draw("mpl", scale=0.6, idle_wires=False, fold=False)

x = np.array([c.depth() for c in pm.run(sub_experiments)])
print(x)

toffoli_layout = [122, 124, 126]

toffoli = QuantumCircuit(3)
toffoli.ccx(0, 1, 2)
toffoli.draw("mpl")

layout=[122,124, 126]
qubit_color = []
for i in range(127):
    if i in layout:
        qubit_color.append("#ff0066")
    else:
        qubit_color.append("#6600cc")
plot_gate_map(backend, qubit_color=qubit_color, qubit_size=60, font_size=25, figsize=(8,8))

# To know the original circuit depth

### Write your code below here ###
# Transpile the circuit
pm = generate_preset_pass_manager(backend=backend, optimization_level=1, initial_layout=layout, seed_transpiler=0)
isa_qc = pm.run(toffoli)



# Calculate original circuit depth
isa_toffoli_depth =  isa_qc.depth()

### Don't change any code past this line ###
print(f"Transpiled circuit depth: ", isa_toffoli_depth)
isa_qc.draw("mpl", scale=0.6, idle_wires=False, fold=False)

# To know the depth using circuit cutting

# Decompose the toffoli circuit
toffoli_ = toffoli.decompose()

### Write your code below here ###
# Find the indices of the distant gates
gates_connecting_to_cut = {0,2} # Hint: Expected type: set {int, int}. Docs: https://docs.python.org/3/tutorial/datastructures.html#sets

cut_indices = [
    i
    for i, instruction in enumerate(toffoli_.data)
    if {toffoli_.find_bit(q)[0] for q in instruction.qubits} == gates_connecting_to_cut
]

# Decompose distant CNOTs into TwoQubitQPDGate instances

qpd_circuit, bases = cut_gates(toffoli_, cut_indices)


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

qpd_circuit.draw("mpl", scale=0.6)

# set the observables
observable = SparsePauliOp(["ZZZ"])

### Write your code below here ###
# Generate the sub-experiments and sampling coefficients
sub_experiments, coefficients = generate_cutting_experiments(
    circuits=qpd_circuit, observables=observable.paulis, num_samples=np.inf

)

# Transpile the circuit
# Note: Use optimization_level=1 and seed_transpiler=0

pm = generate_preset_pass_manager(backend=backend, optimization_level=1, initial_layout=layout, seed_transpiler=0)
isa_qpd_circuit = pm.run(sub_experiments[5])
# Depth using circuit cutting
isa_qpd_toffoli_depth = isa_qpd_circuit.depth()

### Don't change any code past this line ###
print(f"Transpiled circuit depth: ", isa_toffoli_depth)
print(f"QPD sub-experiment depth after transpile: ", isa_qpd_toffoli_depth)
print(f"Number of sub-experiments:", len(sub_experiments))
isa_qpd_circuit.draw("mpl", scale=0.6, idle_wires=False, fold=False)

### Write your code below here ###
# mean of the depth of all sub-experiments
depth_list = [c.depth() for c in pm.run(sub_experiments)]
isa_qpd_toffoli_depth_mean = np.mean(depth_list)

### Don't change any code past this line ###
print(isa_qpd_toffoli_depth_mean)

# Submit your answer using following code

grade_lab3_ckt_ex1(gates_connecting_to_cut, isa_toffoli_depth, depth_list) 
# Expected result type: set, int, numpy.ndarray

### Write your code below here ###
# Find the indices of the distant gates
gates_connecting_to_cut_1 = {1,2}# Hint: Expected type: set {int, int}
gates_connecting_to_cut_2 = {0,2}# Hint: Expected type: set {int, int} 

cut_indices =  [
    i
    for i, instruction in enumerate(toffoli_.data)
    if {toffoli_.find_bit(q)[0] for q in instruction.qubits} == gates_connecting_to_cut_1 or {toffoli_.find_bit(q)[0] for q in instruction.qubits} == gates_connecting_to_cut_2
]

# Decompose distant CNOTs into TwoQubitQPDGate instances
qpd_circuit_2, bases = cut_gates(toffoli_, cut_indices)

### Don't change any code past this line ###
qpd_circuit_2.draw("mpl", scale=0.6)

# set the observables
observable = SparsePauliOp(["ZZZ"])

### Write your code below here ###
# Generate the sub-experiments and sampling coefficients

sub_experiments_2, coefficients = generate_cutting_experiments( 
    circuits=qpd_circuit_2, observables=observable.paulis, num_samples=np.inf

)

# Transpile the circuit
# Note: Use optimization_level=1 and seed_transpiler=0

pm = generate_preset_pass_manager(backend=backend, optimization_level=1, initial_layout=layout, seed_transpiler=0)
isa_qpd_circuit_2 = pm.run(sub_experiments_2[5])
# Depth using circuit cutting
isa_qpd_toffoli_depth_2 = pm.run(sub_experiments_2[5]).depth()


### Don't change any code past this line ###
print(f"QPD sub-experiment depth after transpile: ", isa_qpd_toffoli_depth_2)
print(f"Number of sub-experiments:", len(sub_experiments_2))
isa_qpd_circuit_2.draw("mpl", scale=0.6, idle_wires=False, fold=False)

# Submit your answer using following code

grade_lab3_ckt_ex2(gates_connecting_to_cut_1, gates_connecting_to_cut_2, sub_experiments_2) 
# Expected result type: set, set, list[QuantumCircuit]

### Write your code below here ###
# mean of the depth of all sub-experiments
depth_list_2 = [c.depth() for c in pm.run(sub_experiments_2)]
isa_qpd_toffoli_depth_2_mean = np.mean(depth_list_2)

### Don't change any code past this line ###
print(isa_qpd_toffoli_depth_2_mean)

# Number of sub-experiments
num_sub_experiments_1_cut = len(sub_experiments)
num_sub_experiments_2_cut = len(sub_experiments_2)

# Data for plotting
categories = ['Before Cutting', 'After 1 Cut', 'After 2 Cuts']
depth_values = [isa_toffoli_depth, isa_qpd_toffoli_depth_mean, isa_qpd_toffoli_depth_2_mean]
num_sub_experiments = [1, num_sub_experiments_1_cut, num_sub_experiments_2_cut]

# Create figure and axis
fig, ax1 = plt.subplots()

# Plot depth values
color = 'tab:blue'
ax1.set_xlabel('Number of Cuts')
ax1.set_ylabel('Circuit Depth', color=color)
bars = ax1.bar(categories, depth_values, color=color, alpha=0.6)
ax1.tick_params(axis='y', labelcolor=color)

# Add value labels on bars
for bar in bars:
    yval = bar.get_height()
    ax1.text(bar.get_x() + bar.get_width() / 2, yval + 1, round(yval, 2), ha='center', color=color, fontsize=10)

# Create a second y-axis to plot the number of subexperiments
ax2 = ax1.twinx()
color = 'tab:green'
ax2.set_ylabel('Number of Sub-experiments', color=color)
ax2.plot(categories, num_sub_experiments, color=color, marker='o')
ax2.tick_params(axis='y', labelcolor=color)

# Add value labels on points
for i, num in enumerate(num_sub_experiments):
    ax2.text(i, num + 0.1, num, ha='center', color=color, fontsize=10)

# Add titles and labels
plt.title('Circuit Knitting Toolbox Results')
fig.tight_layout()  # Adjust layout to make room for both y-axes

# Show plot
plt.show()