%matplotlib widget

from monodromy.depthPass import MonodromyDepth
from qiskit.circuit.library import XXPlusYYGate
import numpy as np
from qiskit import QuantumCircuit
from qiskit.quantum_info import Operator
from qiskit.extensions import UnitaryGate
from weylchamber import canonical_gate


gate = XXPlusYYGate(np.pi / 2, np.pi / 6)

qc = QuantumCircuit(2)
qc.append(gate, [0, 1])
qc.swap(0, 1)
gate = UnitaryGate(Operator(qc))


gate = UnitaryGate(canonical_gate(0.5, 0.25, 0.25).full())
MonodromyDepth(basis_gate=gate).coverage_set

from slam.utils.polytopes.coverage_plot import plot_coverage_set

plot_coverage_set(MonodromyDepth(basis_gate=gate).coverage_set)

from qiskit import QuantumCircuit
from weylchamber import c1c2c3, WeylChamber
from qiskit.quantum_info import Operator
import numpy as np
from qiskit import transpile

# observable = SparsePauliOp("XZ")
# qubits = observable.num_qubits
# ansatz = RealAmplitudes(observable.num_qubits, reps=3)
# display(ansatz.decompose().draw('mpl'))
# backend = service.get_backend('ibm_lagos')
# transpiled_ansatz = transpile(ansatz, backend=backend,
#                               basis_gates=['u1', 'u2', 'u3', 'rx', 'ry', 'rz', 'cx', 'cz'],
#                               seed_transpiler=42,
#                               optimization_level=3)
# display(transpiled_ansatz.decompose().draw('mpl'))

# from qiskit.circuit.random import random_circuit
# from qiskit.quantum_info import SparsePauliOp
# from qiskit.circuit.library import RealAmplitudes
# from qiskit_ibm_runtime import QiskitRuntimeService, Session, Estimator, Options

# observable = SparsePauliOp("XZ")
# qubits = observable.num_qubits
# ansatz = RealAmplitudes(observable.num_qubits, reps=3)
# display(ansatz.decompose().draw('mpl'))
# backend = service.get_backend('simulator_statevector')
# transpiled_ansatz = transpile(ansatz, backend=backend,
#                               basis_gates=['u1', 'u2', 'u3', 'rx', 'ry', 'rz', 'cx', 'cz'],
#                               seed_transpiler=42,
#                               optimization_level=3)

# service = QiskitRuntimeService()
# with Session(service=service, backend="ibmq_qasm_simulator") as session:
#      estimator = Estimator(session=session, options=options)
#      job = estimator.run(circuit, observable)
#      result = job.result()
#      # Close the session only if all jobs are finished, and you don't need to run more in the session
#      session.close() # Closes the session

# display(circuit.draw("mpl"))
# print(f" > Observable: {observable.paulis}")
# print(f" > Expectation value: {result.values[0]}")
# print(f" > Metadata: {result.metadata[0]}")

# from qiskit import QuantumCircuit
# qc = QuantumCircuit(2)
# qc.cx(0, 1)
# qc.swap(0,1)
# qc.draw('mpl', filename='cns.svg')
# # print(qc.draw("latex_source"))

# write circuit for CX decomposed into sqiswaps
from qiskit import QuantumCircuit

qc = QuantumCircuit(2)
qc.cx(0, 1)

from slam.utils.transpiler_pass.weyl_decompose import (
    RootiSwapWeylDecomposition,
    SiSwapGate,
)

# decomposer = RootiSwapWeylDecomposition()
# from qiskit.converters import circuit_to_dag, dag_to_circuit
# dag = decomposer.run(circuit_to_dag(qc))
# qc2 = dag_to_circuit(dag)


from qiskit.transpiler.passmanager import PassManager
from qiskit.transpiler.passes import Optimize1qGates

pm = PassManager()
pm.append(RootiSwapWeylDecomposition())
pm.append(Optimize1qGates())
qc2 = pm.run(qc)
qc2.draw("mpl")

qc3 = QuantumCircuit(2)
from qiskit.circuit.library import iSwapGate
import numpy as np

# qc3.u(np.pi/2, -np.pi/2, 0, 0)
qc3.ry(np.pi / 2, 0)
qc3.rz(-np.pi / 2, 0)
# qc3.u(np.pi/2, np.pi, -np.pi/2, 1)
qc3.rz(np.pi / 2, 1)
qc3.ry(-np.pi / 2, 1)
qc3.append(SiSwapGate(), [0, 1])

# qc3.u(np.pi, -np.pi/2, np.pi/2, 0)
qc3.rx(np.pi, 0)
# qc3.u(0, 0, np.pi, 1)
qc3.rz(-np.pi, 1)
qc3.append(SiSwapGate(), [0, 1])

# qc3.u(np.pi/2, -np.pi/2, -3*np.pi/2, 0)
qc3.rx(np.pi / 2, 0)
# qc3.u(0, -np.pi/2, -np.pi, 1)
qc3.rz(np.pi / 2, 1)

qc3.draw("mpl", filename="cx-sqiswap.svg")

# from qiskit.quantum_info import Operator
# Operator(qc4).equiv(qc)
qc3 = QuantumCircuit(2)
qc3.u(np.pi / 2, -np.pi / 2, 0, 0)
# qc3.ry(np.pi/2,  0)
# qc3.rz(-np.pi/2, 0)
qc3.u(np.pi / 2, np.pi, -np.pi / 2, 1)
# qc3.rz(np.pi/2, 1)
# qc3.ry(-np.pi/2, 1)
# qc3.append(SiSwapGate(), [0,1])
qc3.barrier()

qc3.u(np.pi, -np.pi / 2, np.pi / 2, 0)
qc3.u(0, 0, np.pi, 1)
# qc3.append(SiSwapGate(), [0,1])
qc3.barrier()

qc3.u(np.pi / 2, -np.pi / 2, -3 * np.pi / 2, 0)
qc3.u(0, -np.pi / 2, -np.pi, 1)

# pm = PassManager()
# # pm.append(RootiSwapWeylDecomposition())
# pm.append(Optimize1qGates(basis=['u']))
# qc4 = pm.run(qc3)
qc4 = transpile(
    qc3, basis_gates=["x", "y", "z", "h", "t", "s", "sx", "sy", "sz", "rz", "ry"]
)
qc4.draw("mpl", filename="cx_decomp")

from virtual_swap.cns_transform import cx_replace
from qiskit.transpiler.passes import (
    Collect1qRuns,
    ConsolidateBlocks,
    Unroller,
    Optimize1qGates,
)

pm = PassManager()
# pm.append(Collect1qRuns())
# pm.append(ConsolidateBlocks())
# pm.append(Unroller(['u']))
pm.append(Optimize1qGates(basis=["u"]))
# cx_replace.swap(0, 1)
cx_replace2 = pm.run(cx_replace)
cx_replace2.draw("mpl", filename="cx_decomp.svg")

qc = QuantumCircuit(2)
qc.crz(np.pi / 2, 0, 1)
qc = transpile(qc, basis_gates=["u", "cx"])
qc.draw()

# replace the last cx gate with iswap+virtual_swap
from qiskit.converters import circuit_to_dag, dag_to_circuit
from qiskit.circuit.library import CXGate

dag = circuit_to_dag(qc)
cx_node = dag.op_nodes(CXGate)[1]
from virtual_swap.cns_transform import cns_transform

qc2 = dag_to_circuit(cns_transform(dag, cx_node))
qc2.decompose().draw()

Operator(qc2).equiv(qc)

from slam.utils.transpiler_pass.weyl_decompose import RootiSwapWeylDecomposition
from qiskit.transpiler import PassManager
from qiskit.transpiler.passes import ConsolidateBlocks, Collect2qBlocks, Unroller

pm = PassManager()
pm.append(Unroller(["u", "cx", "iswap"]))
pm.append(Collect2qBlocks())
pm.append(ConsolidateBlocks(force_consolidate=True))
pm.append(RootiSwapWeylDecomposition())
pm.run(qc2).draw()

qc3 = QuantumCircuit(2)
qc3.cx(1, 0)
qc3.cx(0, 1)
qc3.draw()

Operator(qc3).equiv(cx_replace)

qc = QuantumCircuit(2)
qc.cx(0, 1)
qc.swap(0, 1)
# transpile into CR gates
transpile(qc, basis_gates=["rzx", "u"], optimization_level=3).draw("mpl")

qc = QuantumCircuit(2)
qc.cx(0, 1)
qc.draw("mpl")

from slam.utils.transpiler_pass.weyl_decompose import RootiSwapWeylDecomposition
from qiskit.transpiler import PassManager

pm = PassManager()
pm.append(RootiSwapWeylDecomposition())
qc = pm.run(qc)
qc.draw("mpl")

qc1 = QuantumCircuit(2)
qc1.rx(np.pi / 2, 1)
qc1.rz(-np.pi / 2, 0)
qc1.rz(np.pi / 2, 1)
qc1.iswap(0, 1)
qc1.rx(np.pi / 2, 0)
qc1.iswap(0, 1)
qc1.rz(np.pi / 2, 1)
qc1.draw("mpl")

from qiskit import QuantumCircuit

qc = QuantumCircuit(2)
qc.cx(1, 0)
qc.h(1)
qc.cx(0, 1)
qc.tdg(1)

qc2 = QuantumCircuit(2)
qc2.s(0)
qc2.cx(1, 0)
qc2.sx(1)
qc2.sdg(0)
qc2.t(1)

qc3 = QuantumCircuit(2)
qc3.cx(1, 0)
qc3.cx(0, 1)

qc4 = QuantumCircuit(2)
qc4.cx(0, 1)
qc4.swap(0, 1)

c1 = c1c2c3(Operator(qc).data)
c2 = c1c2c3(Operator(qc2).data)
c3 = c1c2c3(Operator(qc3).data)
c4 = c1c2c3(Operator(qc4).data)
print(c1, c2, c3, c4)

qc3.draw("mpl")

qc4.draw("mpl")

Operator(qc4).equiv(cx_replace)

from virtual_swap.cns_transform import cx_replace, iswap_replace

cx_replace.draw("mpl")

iswap_replace.draw("mpl")

Operator(qc3).equiv(qc4)

# visualize weyl chamber coordinates
from weylchamber import WeylChamber

w = WeylChamber()
N = 1000
for _ in range(N):
    qc = QuantumCircuit(2)
    # random 1q
    qc.u(
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        0,
    )
    qc.u(
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        1,
    )
    # qc.crz(2 * np.pi * np.random.rand(), 0, 1)
    qc.cx(0, 1)

    # random 1q
    qc.u(
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        0,
    )
    qc.u(
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        1,
    )
    # qc.swap(0, 1)
    # qc.cx(0, 1)
    qc.iswap(0, 1)

    # random 1q
    qc.u(
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        0,
    )
    qc.u(
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        1,
    )

    qc.swap(0, 1)
    op = Operator(qc).data
    w.add_point(*c1c2c3(op))
w.plot()

qc.draw("mpl")

# visualize weyl chamber coordinates
from qiskit import QuantumCircuit
from weylchamber import c1c2c3, WeylChamber
from qiskit.quantum_info import Operator
import numpy as np

w = WeylChamber()
N = 1000
for _ in range(N):
    qc = QuantumCircuit(2)
    # random 1q
    qc.u(
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        0,
    )
    qc.u(
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        1,
    )
    qc.cx(0, 1)

    # random 1q
    qc.u(
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        0,
    )
    qc.u(
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        1,
    )
    qc.iswap(0, 1)

    # random 1q
    qc.u(
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        0,
    )
    qc.u(
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        1,
    )
    qc.swap(0, 1)

    # random 1q
    qc.u(
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        0,
    )
    qc.u(
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        2 * np.pi * np.random.rand(),
        1,
    )
    op = Operator(qc).data
    w.add_point(*c1c2c3(op))
w.plot()

import matplotlib.pyplot as plt

w = WeylChamber()
# generate a color gradient
N = 9
color_gradient = [
    plt.cm.tab10(i / (N + 1)) for i in range(N + 1)
]  # using 'plasma' colormap here
w = WeylChamber()

# create two lists for the two types of operations
op1_list = []
op2_list = []

# iterate from I to CX (theta = 0 to np.pi)
for i in range(N + 1):
    qc = QuantumCircuit(2)
    qc.cp(np.pi * i / N, 0, 1)
    op1 = Operator(qc).data
    t = c1c2c3(op1)
    # force to use x < 1/2
    if t[0] > 0.5:
        t = (1 - t[0], t[1], t[2])
    op1_list.append((t, color_gradient[i]))
    qc.swap(0, 1)
    op2 = Operator(qc).data
    op2_list.append((c1c2c3(op2), color_gradient[i]))

# add the operations to the Weyl chamber
for coords, color in op1_list:
    w.scatter(*coords, color=color)
for coords, color in op2_list:
    w.scatter(*coords, color=color)

# define a region using line segments between the following points
# (0,0,0), (1/2, 1/2, 0); (0,0,0), (1,0,0); (1/2, 1/2, 0), (1,0,0); (0,0,0), (1/2, 1/4, 1/4)
lines = [
    [(0, 0, 0), (0.5, 0.5, 0)],
    [(0, 0, 0), (1, 0, 0)],
    [(0.5, 0.5, 0), (1, 0, 0)],
    [(0, 0, 0), (0.5, 0.25, 0.25)],
    [(0.5, 0.5, 0), (0.5, 0.25, 0.25)],
    [(0.5, 0.25, 0.25), (1, 0, 0)],
]
# draw the lines
w.labels = {}
w.plot()
for i in range(len(lines)):
    w._draw_line(
        w.ax, origin=lines[i][0], end=lines[i][1], color="blue", linestyle="--"
    )
w.fig.savefig("cphase_swap.svg", dpi=300)
# fig.savefig("cphase_swap.svg", dpi=300)



w = WeylChamber()

# generate Haar random unitary
from qiskit.quantum_info import random_unitary

op = random_unitary(4).data
qc = QuantumCircuit(2)
qc.unitary(op, [0, 1])
op = Operator(qc).data
w.scatter(*c1c2c3(op), color="blue")
qc.swap(0, 1)
op2 = Operator(qc).data
w.scatter(*c1c2c3(op2), color="red")
w.plot()

import numpy as np
from qiskit.quantum_info import random_unitary
from qiskit.quantum_info import Operator
from qiskit import QuantumCircuit
from weylchamber import WeylChamber, c1c2c3
import matplotlib.pyplot as plt
from matplotlib import cm

w = WeylChamber()


# Generate many random unitaries
N = 1000
coords_initial = []
coords_final = []
for _ in range(N):
    # Generate Haar random unitary
    op = random_unitary(4).data
    qc = QuantumCircuit(2)
    qc.unitary(op, [0, 1])
    op = Operator(qc).data
    coords_initial.append(c1c2c3(op))

    # Apply SWAP
    qc.swap(0, 1)
    op2 = Operator(qc).data
    coords_final.append(c1c2c3(op2))

# Convert to numpy arrays for easier handling
coords_initial = np.array(coords_initial)
coords_final = np.array(coords_final)

# Plot initial points with blue color
w.scatter(
    coords_initial[:, 0], coords_initial[:, 1], coords_initial[:, 2], color="blue"
)

# Plot final points with red color
w.scatter(coords_final[:, 0], coords_final[:, 1], coords_final[:, 2], color="red")

import matplotlib as mpl

# Create a color map with 256 colors
color_map = mpl.cm.get_cmap("viridis", 256)

# Generate colors from the color map
colors = color_map(np.linspace(0, 1, N))

# # Plot trajectories
# for i in range(N):
#     w._draw_line(
#         (coords_initial[i, 0], coords_initial[i, 1], coords_initial[i, 2]),
#         (coords_final[i, 0], coords_final[i, 1], coords_final[i, 2]),
#         color=colors[i],
#     )

w.plot()