from demo_src.graph import generate_demo_graph, draw_graph

demo_graph = generate_demo_graph()
draw_graph(demo_graph)

from demo_src.map import map_graph_to_qubo

qubo = map_graph_to_qubo(demo_graph)
print(qubo.prettyprint())

from demo_src.map import map_qubo_to_ising

cost_hamiltonian, offset = map_qubo_to_ising(qubo)
print("Offset:", offset)
print("Cost Function Hamiltonian:", cost_hamiltonian)

from demo_src.map import map_ising_to_circuit

circuit = map_ising_to_circuit(cost_hamiltonian, num_layers=1)
circuit.draw('mpl')

circuit.parameters

# IBM Quantum devices are named after cities
# For the purpose of the first part of this demo, we will
# use a simulated device from the "fake_provider"
from qiskit_ibm_runtime.fake_provider import FakeVigoV2
backend = FakeVigoV2()

from demo_src.transpile import optimize_circuit
opt_circuit = optimize_circuit(circuit, backend)
opt_circuit.draw('mpl', fold=False)

optimal_gamma = 5.11350346
optimal_beta = 5.52673212

candidate_circuit = opt_circuit.assign_parameters([optimal_gamma, optimal_beta])
candidate_circuit.draw('mpl', fold=False)

# For the purpose of the first part of this demo, we will
# use a simulated device from the "fake_provider"
from qiskit_ibm_runtime.fake_provider import FakeVigoV2
from qiskit.primitives import BackendSampler

backend = FakeVigoV2()
sampler = BackendSampler(backend=backend)

final_distribution = sampler.run(candidate_circuit, shots=int(1e4)).result().quasi_dists[0]
print(final_distribution)

from demo_src.post import sample_most_likely

best_result = sample_most_likely(final_distribution, len(demo_graph))
print("Result bitstring:", best_result)

from demo_src.post import plot_distribution

plot_distribution(final_distribution)

from demo_src.post import plot_result

plot_result(demo_graph, best_result)

with open("data/125node_example.lp", "r") as file:
    problem = file.read()
print(problem)

with open("data/125node_example_ising.txt") as input_file:
    for _ in range(10):
        print(str(next(input_file)).replace("\n", ""))

team = 1  # Fill in your team here, either 1 or 2

from qiskit import qpy  # QPY is the circuit serializer in Qiskit.

# Load the circuits
if team == 1:
    backend_type = "eagle"
elif team == 2:
    backend_type = "heron"
else:
    raise ValueError("team should be 1 or 2.")

# Depth zero-circuit
with open(f"data/125node_{backend_type}_depth_zero.qpy", "rb") as fd:
    depth_zero_circuit = qpy.load(fd)[0]

# Depth one-circuit
with open(f"data/125node_{backend_type}_depth_one.qpy", "rb") as fd:
    depth_one_circuit = qpy.load(fd)[0]

depth_one_circuit.draw('mpl', fold=False)

from qiskit_ibm_runtime import QiskitRuntimeService

# QiskitRuntimeService.save_account(channel="ibm_quantum", token="<MY_IBM_QUANTUM_TOKEN>", overwrite=True, set_as_default=True)

if team == 1:
    # Eagle = Osaka
    service = QiskitRuntimeService(channel='ibm_quantum')
    backend = service.get_backend('ibm_osaka')
elif team == 2:
    # Heron = Torino
    service = QiskitRuntimeService(channel='ibm_quantum')
    backend = service.get_backend('ibm_torino')

from qiskit_ibm_runtime import Sampler, Options

# Since we have already done the transpilation we can skip it.
options = Options()
options.transpilation.skip_transpilation = True

sampler = Sampler(backend=backend, options=options)

depth_one_circuit.parameters

depth_one_bound_circuit = depth_one_circuit.assign_parameters([0.3927,  0.3927], inplace=False)

#sampler_job = sampler.run([depth_one_bound_circuit, depth_zero_circuit])

#sampler_job.job_id()

#sampler_job.status()

from demo_src.run import save_result

# save_result(sampler_job.result(), backend_type, path="sampler_data")

from demo_src.post import load_data, samples_to_objective_values, load_qp, plot_cdf

import matplotlib.pyplot as plt

qp, max_cut, min_cut = load_qp()

depth_one_heron, depth_zero_heron, depth_one_eagle, depth_zero_eagle = load_data(qp)

if team == 1:
    depth_zero = depth_zero_eagle
    depth_one = depth_one_eagle
    
elif team ==2:
    depth_zero = depth_zero_heron
    depth_one = depth_one_heron
    
fig, ax = plt.subplots(1, 1)

plot_cdf(depth_zero, depth_one, max_cut, min_cut, ax, "Approximation ratio:")

fig, ax = plt.subplots(1, 2, figsize=(15, 6))
plot_cdf(depth_zero_eagle, depth_one_eagle, max_cut, min_cut, ax[0], "Approximation ratio Eagle:")
plot_cdf(depth_zero_heron, depth_one_heron, max_cut, min_cut, ax[1], "Approximation ratio Heron:")