import sys; sys.path.append('..')
from qiskit import transpile
from qiskit.circuit import QuantumCircuit
from qiskit.transpiler import PassManager
from qiskit.qasm2 import dumps
from qiskit.qasm3 import dumps as dumps3

import pyzx as zx

from zxpass import ZXPass
pass_manager = PassManager(ZXPass())

# Select the benchmark circuit to examine.
subdir = "medium"
circuit_name = "qft_n18"

# Output the original circuit.
with open(f"QASMBench/{subdir}/{circuit_name}/{circuit_name}.qasm", "r") as file:
    qasm_str = file.read()

print(f"Benchmark circuit: {circuit_name}")
qc = QuantumCircuit.from_qasm_str(qasm_str)
# qc.draw(output = "mpl")

# Output the ZX diagram for the original circuit.
g = zx.Circuit.from_qasm(qasm_str).to_graph()
zx.draw(g)

# Output the circuit produced by Qiskit optimisation.
opt_qc = transpile(qc, basis_gates=['u3', 'cx'], optimization_level=3)
# opt_qc.draw(output = "mpl")

# Output the ZX diagram for the Qiskit-optimised circuit.
opt_g = zx.Circuit.from_qasm(dumps3(opt_qc)).to_graph()
zx.draw(opt_g)

# Output the ZX-optimised circuit.
zx_qc = pass_manager.run(qc)
# zx_qc.draw(output = "mpl")

# Re-create the graph from the circuit. 
zx_g = zx.Circuit.from_qasm(dumps(zx_qc)).to_graph()
zx.draw(zx_g)

