%load_ext autoreload
%autoreload 2

# imports

from qiskit.circuit import QuantumCircuit
from qiskit.quantum_info import Clifford, Operator, StabilizerState, Statevector
from qiskit.transpiler import CouplingMap

from synthesis.sat_problem_clifford import SatProblemClifford
from synthesis.utils import line_coupling_map
from synthesis.synthesize import synthesize_optimal
from synthesis.synthesize_clifford import synthesize_clifford_depth2q

# naive preparation of the state of interest (does not adhere to LNN connectivity)

nq = 13  # number of qubits
k = 3  # number of consequtive ones (0: GHZ, 1: checkerboard, n/2: hartree-fock)
ones = [i for m in range(1, int(nq / k), 2) for i in range(m * k, (m + 1) * k)]
qc = QuantumCircuit(nq)
qc.h(nq - 1)
for target in ones:
    qc.cx(nq - 1, target)

display(Statevector(qc).draw("latex"))
qc.draw("mpl")

# prepackaged clifford synthesis routine

cliff = Clifford(qc)
res = synthesize_clifford_depth2q(
    cliff, coupling_map=line_coupling_map(cliff.num_qubits), state_preparation_mode=True
)
res.circuit.draw("mpl", fold=-1)

# Create a custom problem: we want to use a single H-gate in the first layer and only CX-gates after that


def create_problem_custom(cliff, depth):
    nq = cliff.num_qubits
    sat_problem = SatProblemClifford(nq, verbosity=1)
    sat_problem.set_init_matrix_to_identity(nq)
    sat_problem.set_final_clifford(cliff)
    sat_problem.set_state_preparation_mode(True)
    sat_problem.add_layer(gates=["H"], coupling_maps=[])
    for _ in range(depth):
        sat_problem.add_layer(gates=["CX"], coupling_maps=[line_coupling_map(nq)])
    return sat_problem


def synthesize_clifford_custom(cliff):
    res = synthesize_optimal(
        target_obj=cliff, create_sat_problem_fn=create_problem_custom, verbosity=1
    )
    return res.circuit

cliff = Clifford(qc)
circuit = synthesize_clifford_custom(cliff)
circuit.draw("mpl", fold=-1)

# let's say we also want at most 2 unique layers


def create_problem_custom_u(cliff, depth):
    nq = cliff.num_qubits
    sat_problem = SatProblemClifford(nq, verbosity=1)
    sat_problem.set_init_matrix_to_identity(nq)
    sat_problem.set_final_clifford(cliff)
    sat_problem.set_state_preparation_mode(True)
    sat_problem.add_layer(gates=["H"], coupling_maps=[])
    two_qubit_layers = []
    for _ in range(depth):
        layer = sat_problem.add_layer(
            gates=["CX"], coupling_maps=[line_coupling_map(nq)]
        )
        two_qubit_layers.append(layer)
    sat_problem.add_max_unique_layers_constraint(two_qubit_layers, 2)
    return sat_problem


def synthesize_clifford_custom_u(cliff):
    res = synthesize_optimal(
        target_obj=cliff, create_sat_problem_fn=create_problem_custom_u, verbosity=1
    )
    return res.circuit


circuit = synthesize_clifford_custom_u(cliff)
circuit.draw("mpl", fold=-1)

