import sys

sys.path.append("..")

from qiskit import transpile
from qiskit.converters import circuit_to_dag, dag_to_circuit
from circuit_suite import circuits

circuit_label = "QFT"
qsize = 3
qc = transpile(circuits[circuit_label].circuit_lambda(qsize)).decompose()
qc.draw(output="mpl")

# from qiskit import transpile
# from qiskit.converters import circuit_to_dag, dag_to_circuit
# from circuit_suite import circuits
# circuit_label="QFT"
# c = transpile(circuits[circuit_label].circuit_lambda(qsize), basis_gates=['u', 'cx'])
# c.draw(output='mpl')

dag = circuit_to_dag(qc)


def weight_fn(source_node, target_node, weight):
    if dag._multi_graph[source_node].name is None:
        return 0
    else:
        return 1


import retworkx

longest_path = retworkx.dag_longest_path(dag._multi_graph, weight_fn=weight_fn)
dlp = [dag._multi_graph[index] for index in longest_path]
for node in dag.op_nodes():
    if not node in dlp:
        dag.remove_op_node(node)
qc_t = dag_to_circuit(dag)
qc_t.draw(output="mpl")

x = [qsize]
y = [
    circuit_to_dag(
        transpile(circuits[circuit_label].circuit_lambda(xi), basis_gates=["u", "cx"])
    ).count_ops_longest_path()["cx"]
    for xi in x
]
z = [
    circuit_to_dag(
        transpile(circuits[circuit_label].circuit_lambda(xi), basis_gates=["u", "cx"])
    ).count_ops()["cx"]
    for xi in x
]
print(y, z)

c = transpile(circuits[circuit_label].circuit_lambda(qsize), basis_gates=["u", "cx"])
d = circuit_to_dag(c)
w = d._wires
for wi in w:
    print([node.name for node in d.nodes_on_wire(wi, only_ops=True)])

c.draw(output="mpl")

import numpy as np
from numpy import pi

# importing Qiskit
from qiskit import QuantumCircuit, transpile, assemble, Aer, IBMQ
from qiskit.providers.ibmq import least_busy
from qiskit.tools.monitor import job_monitor
from qiskit.visualization import plot_histogram, plot_bloch_multivector

from qiskit.circuit.library.basis_change import QFT

inverse_qft_lambda = lambda q: QFT(q, inverse=True)
qc = inverse_qft_lambda(qsize)
qc.decompose().draw(output="mpl")

# replace every R_z with a Riswap?
# can e made a transformation pass if wanted to
from qiskit.circuit.library.standard_gates import iSwapGate
from clonk.utils.riswap_gates.riswap import RiSwapGate

dag = circuit_to_dag(qc.decompose())
for node in dag.op_nodes():
    # if we hit a Rz gate then replace it
    if node.op.name in ["cp"]:
        # get the rotation angle
        angle = node.op.params[0]

        replacement = QuantumCircuit(2)
        # replacement.append(RiSwapGate(angle/np.pi), [0,1])
        replacement.append(iSwapGate().power(angle / np.pi), [0, 1])

        # replacement.draw(output='mpl')
        dag.substitute_node_with_dag(node, circuit_to_dag(replacement))

        # print(angle)
qc_modified = dag_to_circuit(dag)
qc_modified.draw(output="mpl")

nqubits = 3
number = 3
prep = QuantumCircuit(qsize)
for qubit in range(qsize):
    prep.h(qubit)
prep.p(number * pi / 4, 0)
prep.p(number * pi / 2, 1)
prep.p(number * pi, 2)

prep.draw()

qc = prep + qc_modified
qc.draw(output="mpl")

# now simulate the qft to get a histogram
backend = sim = Aer.get_backend("aer_simulator")
qc.measure_all()
transpiled_qc = transpile(qc, backend, optimization_level=3, basis_gates=["u", "cx"])
job = backend.run(transpiled_qc, shots=1000)
job_monitor(job)

counts = job.result().get_counts()
plot_histogram(counts)