from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit.transpiler import CouplingMap
from qiskit.converters import dag_to_circuit
from qiskit.transpiler.layout import Layout
import numpy as np
from qiskit.transpiler.passes import Unroller
from qiskit.transpiler.passes import (
    TrivialLayout,
    BasicSwap,
    ApplyLayout,
    SetLayout,
    SabreLayout,
    FullAncillaAllocation,
    EnlargeWithAncilla,
)
from qiskit.transpiler.passmanager import PassManager
from virtual_swap.deprecated.cns_sabre import CNS_SabreSwap

# load a dummy circuit
# generate network of random cnots
n, M = 3, 8
N = n**2
register = QuantumRegister(N)
qc = QuantumCircuit(register)

for _ in range(M):
    qc.cx(*np.random.choice(N, 2, replace=False))
qc.decompose().draw("mpl")

# load a dummy coupling map and layout
coupling = CouplingMap.from_line(N)
layout = Layout({q: i for q, i in zip(register, range(N))})

# print DAG
from qiskit.converters import circuit_to_dag

dag = circuit_to_dag(qc)
dag.draw("mpl")

set_pass = SetLayout(layout)
dag_p = set_pass.run(dag)
apply_pass = ApplyLayout()
apply_pass.property_set = set_pass.property_set
dag_p = apply_pass.run(dag_p)
dag_p.draw("mpl")

# get layers of the DAG

"""A layer is a circuit whose gates act on disjoint qubits, i.e., a layer has depth 1. The total number of layers equals the circuit depth d. The layers are indexed from 0 to d-1 with the earliest layer at index 0. The layers are constructed using a greedy algorithm. Each returned layer is a dict containing {“graph”: circuit graph, “partition”: list of qubit lists}.

The returned layer contains new (but semantically equivalent) DAGOpNodes, DAGInNodes, and DAGOutNodes. These are not the same as nodes of the original dag, but are equivalent via DAGNode.semantic_eq(node1, node2)."""

# convert generator to list
layers = list(dag_p.layers())
for layer in layers:
    print(layer)

# # get serial layers
# """A serial layer is a circuit with one gate. The layers have the same structure as in layers()."""
# serial_layers = dag.serial_layers()
# for layer in serial_layers:
#     print(layer)

# # get multigraph layers
# multi_layers = dag.multigraph_layers()
# for layer in multi_layers:
#     print(layer)

# want to get topological distance from output F and inputs to F'
cns_sub = False
F = layers[0]["graph"]
F_prime = layers[1]["graph"]

F.draw()

# F.output_map

from virtual_swap.cns_transform import cns_transform
from qiskit.circuit.library.standard_gates import CXGate

# get cx from F
cns_sub = True
cns_cx = F.op_nodes(CXGate)[0]
F_cns = cns_transform(F, cns_cx)
F_cns.draw()

# def evaluate_layer(F, F_prime):
#     # cost from F:output -> F':input

# work on getting output permutations from a single layer
# this includes virtual swaps + slack swaps that don't affect the depth
v_to_p = layout.get_virtual_bits()

# FIXME, SABRE does meet in middle to increase parallelism
# we need to implement our own here,
# tests both CNS subs and slack swaps
swapper = BasicSwap(coupling)
if cns_sub:
    F_swapped = swapper.run(F_cns)
else:
    F_swapped = swapper.run(F)

F_swapped.draw()

# swapper.property_set
# XXX this needs to be manually set following a swap pass
f_layout = swapper.property_set["final_layout"]
if cns_sub:
    temp = f_layout[cns_cx.qargs[0]]
    f_layout[cns_cx.qargs[0]] = f_layout[cns_cx.qargs[1]]
    f_layout[cns_cx.qargs[1]] = temp
print(f_layout)

f_v_to_p = f_layout.get_virtual_bits()
print(f_v_to_p)
# for x in F_swapped.output_map.items():
#     print(x)

F_prime.draw()

trial_cost = 0
for x in F_prime.two_qubit_ops():
    print(f_v_to_p[x.qargs[0]], f_v_to_p[x.qargs[1]])
    print(coupling.distance(f_v_to_p[x.qargs[0]], f_v_to_p[x.qargs[1]]))
    trial_cost += coupling.distance(f_v_to_p[x.qargs[0]], f_v_to_p[x.qargs[1]])

    # FIXME do we want to return sum or max?
print(f"Trial cost", trial_cost)

# we need to define a cost function, intermediate layer to front layer
# the intermediate layer are gates that are physical on the topology
# the intermediate layer considers CNS subs and placement of SWAP gates
# once a permutation of the intermediate layer is found that minimizes cost to the front layer
# move the intermediate layer to the mapped dag, and front layer to intermediate layer
# once the front layer is empty, we are done

# open question, do we resolve all dependencies between intermediate and front layer at once?
# maybe we find a route for one gate, then pop the next gate in
# if we have a dynamic layer, makes reasoning about cost and when to move gates harder
# but that is how SABRE works, just by iteratively checking what is physical
# repeat pushing things over until we can't push anymore
# idk, I think conceptually treating it as a set of layers that need to be resolved is easier

# we define the minimum about of gates that add to the depth of the circuit for data movement
# types of SWAP paths we could take

# if 2 qubits can meet, they can SWAP together to meet at midpoint between them
# however, we might prefer to meet biased towards one qubit (for sake of lookahead)
# then, we incur a higher depth cost on this layer, but hopefully reduces lookahead cost

# so define a depth cost on the layer,

# maybe we just do the thing that tells us which decomp to use first