%load_ext autoreload
%autoreload 2

!pip install git+https://github.com/qiskit-community/qiskit-sat-synthesis

import numpy as np

# Qiskit includes
from qiskit.circuit import QuantumCircuit
from qiskit.circuit.library import LinearFunction, PermutationGate
from qiskit.quantum_info import Clifford
from qiskit.transpiler import CouplingMap
from qiskit.compiler import transpile
from qiskit.transpiler.passes.synthesis.high_level_synthesis import HLSConfig
from qiskit.transpiler.passes import HighLevelSynthesis

from qiskit.transpiler.passes.synthesis import high_level_synthesis_plugin_names
print(f'Clifford plugin names: {high_level_synthesis_plugin_names("clifford")}')
print(f'LinearFunction plugin names: {high_level_synthesis_plugin_names("linear_function")}')
print(f'PermutationGate plugin names: {high_level_synthesis_plugin_names("permutation")}')

mat = np.array(
    [
        [1, 1, 0, 1, 1],
        [0, 1, 1, 1, 1],
        [0, 0, 1, 1, 0],
        [0, 0, 0, 1, 1],
        [1, 0, 1, 1, 0],
    ]
)
mat

qc = QuantumCircuit(5)
qc.append(LinearFunction(mat), [0, 1, 2, 3, 4])
qc.draw("mpl")

# Transpile (default)

tqc = transpile(qc)
print(f"ops: {tqc.count_ops()}, depth: {tqc.depth()}")
tqc.draw("mpl")

hls_config = HLSConfig(linear_function=[("sat_count", {})])
tqc = transpile(qc, hls_config=hls_config)
print(f"ops: {tqc.count_ops()}, depth: {tqc.depth()}")
tqc.draw("mpl")

hls_config = HLSConfig(linear_function=[("sat_depth", {})])
tqc = transpile(qc, hls_config=hls_config)
print(f"ops: {tqc.count_ops()}, depth: {tqc.depth()}")
tqc.draw("mpl")

hls_config = HLSConfig(
    linear_function=[("sat_depth", {"optimize_1q_gates": True, "verbosity": 1})]
)
tqc = transpile(qc, hls_config=hls_config)
print(f"ops: {tqc.count_ops()}, depth: {tqc.depth()}")
tqc.draw("mpl")

hls_config = HLSConfig(linear_function=[("sat_depth", {"allow_final_permutation": True})])
tqc = transpile(qc, hls_config=hls_config)
print(f"ops: {tqc.count_ops()}, depth: {tqc.depth()}")
tqc.draw("mpl")

hls_config = HLSConfig(linear_function=[("sat_depth", {"allow_layout_permutation": True})])
tqc = transpile(qc, hls_config=hls_config)
print(f"ops: {tqc.count_ops()}, depth: {tqc.depth()}")
tqc.draw("mpl")

hls_config = HLSConfig(linear_function=[("sat_depth", {})])
coupling_map = CouplingMap.from_line(5)
tqc = HighLevelSynthesis(hls_config=hls_config, coupling_map=coupling_map, use_qubit_indices=True)(qc)
print(f"ops: {tqc.count_ops()}, depth: {tqc.depth()}")
tqc.draw("mpl")

# This is cyclic shift: 1->0, 2->1, 0->2

pattern = [1, 2, 3, 4, 5, 6, 7, 0]
qc = QuantumCircuit(8)
qc.append(PermutationGate(pattern), [0, 1, 2, 3, 4, 5, 6, 7])
qc.draw(output='mpl')

hls_config = HLSConfig(permutation=[("sat_depth", {"optimize_2q_gates": True})])
coupling_map = CouplingMap.from_full(8)
tqc = HighLevelSynthesis(hls_config=hls_config, coupling_map=coupling_map, use_qubit_indices=True)(qc)
print(f"ops: {tqc.count_ops()}, depth: {tqc.depth()}")
tqc.draw("mpl")

hls_config = HLSConfig(permutation=[("sat_depth", {"optimize_2q_gates": True})])
coupling_map = CouplingMap.from_line(8)
tqc = HighLevelSynthesis(hls_config=hls_config, coupling_map=coupling_map, use_qubit_indices=True)(qc)
print(f"ops: {tqc.count_ops()}, depth: {tqc.depth()}")
tqc.draw("mpl")

print(mat)

from qiskit_sat_synthesis.synthesize_linear import synthesize_linear_depth, synthesize_linear_count

line5 = list(CouplingMap.from_line(5).get_edges())  # line connectivity
full5 = list(CouplingMap.from_full(5).get_edges())  # all-to-all connectivity

res = synthesize_linear_depth(mat, coupling_map=full5, verbosity=0, allow_final_permutation=True)
res.circuit.draw("mpl")

res.circuit_with_permutations.draw(output='mpl')

print(res.solutions[0].final_permutation)

LinearFunction(res.circuit_with_permutations) == LinearFunction(mat)

res = synthesize_linear_depth(mat, coupling_map=full5, verbosity=1, allow_final_permutation=True, print_solutions=False, max_solutions=5)

for i, solution in enumerate(res.solutions):
    print(f"Solution {i+1}:")
    print(solution.circuit)

from qiskit_sat_synthesis.synthesize_superposed import (
    synthesize_superposed_count,
    synthesize_superposed_depth,
)

coupling_map = [(0, 1), (1, 0), (1, 2), (2, 1), (2, 3), (3, 2), (3, 4), (4, 3), (4, 5), (5, 4), (5, 6), (6, 5), (6, 7), (7, 6), (7, 8), (8, 7)]  # line connectivity

result = synthesize_superposed_depth(
    [1, 0, 1, 0, 1, 0, 1, 0, 1],
    coupling_map=coupling_map,
    optimize_2q_gates=True,
    print_solutions=False,
    verbosity=0,
    max_depth2q=10,
)

result.circuit.draw('mpl')

hor9 = [[0, 1], [3, 4], [6, 7], [1, 0], [4, 3], [7, 6], [1, 2], [4, 5], [7, 8], [2, 1], [5, 4], [8, 7]]
ver9 = [[0, 3], [1, 4], [2, 5], [3, 0], [4, 1], [5, 2], [3, 6], [4, 7], [5, 8], [6, 3], [7, 4], [8, 5]]
all9 = hor9 + ver9

state = [1, 0, 0, 0, 0, 0, 0, 0, 1]

res = synthesize_superposed_depth(
    state,
    coupling_map=all9,
    print_solutions=False,
    check_solutions=True,
    optimize_2q_gates=True,
    verbosity=0,
)
res.circuit.draw(output='mpl')

res = synthesize_superposed_depth(
    state,
    coupling_map_list=[hor9, ver9],
    print_solutions=False,
    check_solutions=True,
    optimize_2q_gates=True,
    verbosity=0,
)
res.circuit.draw(output='mpl')

mat1 = np.array(
    [
        [1, 0, 0, 0],
        [0, 1, 0, 0],
        [0, 0, 1, 0],
        [1, 1, 1, 1],
    ]
)

mat2 = np.array(
    [
        [1, 0, 0, 0],
        [0, 1, 0, 0],
        [1, 0, 1, 0],
        [1, 1, 1, 1],
    ]
)

from qiskit_sat_synthesis.sat_problem_linear import SatProblemLinear
def create_custom_problem(k):
    nq = 4
    coupling_map = [[0, 1], [1, 0], [1, 2], [2, 1], [2, 3], [3, 2]]
    sat_problem = SatProblemLinear(nq, verbosity=0)
    sat_problem.set_init_matrix(mat1)
    for _ in range(k):
        sat_problem.add_layer(
            gates=["CX"],
            coupling_maps=[coupling_map],
        )
    sat_problem.set_final_matrix(mat2)
    sat_problem.set_optimize_2q_gate(True)
    return sat_problem

from qiskit_sat_synthesis.synthesize import SynthesisResult, synthesize_optimal
res = synthesize_optimal(create_sat_problem_fn=create_custom_problem, verbosity=1)

print(res.circuit)

def create_cx_swap_problem(k, mat):
    nq = len(mat)
    coupling_map = list(CouplingMap.from_line(nq).get_edges())
    sat_problem = SatProblemLinear(nq, verbosity=0)
    sat_problem.set_init_matrix_to_identity(nq)
    for _ in range(k):
        sat_problem.add_layer(
            gates=["CX", "SWAP"],
            coupling_maps=[coupling_map],
        )
    sat_problem.set_final_matrix(mat)
    return sat_problem

mat = np.array(
    [
        [1, 1, 0, 1, 1],
        [0, 1, 1, 1, 1],
        [0, 0, 1, 1, 0],
        [0, 0, 0, 1, 1],
        [1, 0, 1, 1, 0],
    ]
)

from functools import partial

res = synthesize_optimal(
    create_sat_problem_fn=partial(create_cx_swap_problem, mat=mat), verbosity=1
)

res.circuit.draw(output='mpl')

print(LinearFunction(res.circuit) == LinearFunction(mat))

def create_cx_swap_problem(k, mat):
    nq = len(mat)
    coupling_map = list(CouplingMap.from_line(nq).get_edges())
    sat_problem = SatProblemLinear(nq, verbosity=0)
    sat_problem.set_init_matrix_to_identity(nq)
    for _ in range(k):
        layer_id = sat_problem.add_layer(gates=["CX", "SWAP"], coupling_maps=[coupling_map])
        sat_problem.add_nonempty_constraint(layer_id)
    sat_problem.set_final_matrix(mat)
    return sat_problem

res = synthesize_optimal(
    create_sat_problem_fn=partial(create_cx_swap_problem, mat=mat), verbosity=1
)

res.circuit.draw(output='mpl')

