import sys

sys.path.append("..")

# the following is from Qiskits two_qubit_decomp, (doesn't work as import since need to adjust parameters and return vals)
import cmath
from qiskit.quantum_info.synthesis.two_qubit_decompose import *
import scipy.linalg as la

_ipx = np.array([[0, 1j], [1j, 0]], dtype=complex)
_ipy = np.array([[0, 1], [-1, 0]], dtype=complex)
_ipz = np.array([[1j, 0], [0, -1j]], dtype=complex)
_id = np.array([[1, 0], [0, 1]], dtype=complex)


def KAKDecomp(unitary_matrix, *, fidelity=(1.0 - 1.0e-9)):
    """Perform the Weyl chamber decomposition, and optionally choose a specialized subclass.

    The flip into the Weyl Chamber is described in B. Kraus and J. I. Cirac, Phys. Rev. A 63,
    062309 (2001).

    FIXME: There's a cleaner-seeming method based on choosing branch cuts carefully, in Andrew
    M. Childs, Henry L. Haselgrove, and Michael A. Nielsen, Phys. Rev. A 68, 052311, but I
    wasn't able to get that to work.

    The overall decomposition scheme is taken from Drury and Love, arXiv:0806.4015 [quant-ph].
    """
    pi = np.pi
    pi2 = np.pi / 2
    pi4 = np.pi / 4

    # Make U be in SU(4)
    U = np.array(unitary_matrix, dtype=complex, copy=True)
    detU = la.det(U)
    U *= detU ** (-0.25)
    global_phase = cmath.phase(detU) / 4

    Up = transform_to_magic_basis(U, reverse=True)
    M2 = Up.T.dot(Up)

    # M2 is a symmetric complex matrix. We need to decompose it as M2 = P D P^T where
    # P ∈ SO(4), D is diagonal with unit-magnitude elements.
    #
    # We can't use raw `eig` directly because it isn't guaranteed to give us real or othogonal
    # eigenvectors.  Instead, since `M2` is complex-symmetric,
    #   M2 = A + iB
    # for real-symmetric `A` and `B`, and as
    #   M2^+ @ M2 = A^2 + B^2 + i [A, B] = 1
    # we must have `A` and `B` commute, and consequently they are simultaneously diagonalizable.
    # Mixing them together _should_ account for any degeneracy problems, but it's not
    # guaranteed, so we repeat it a little bit.  The fixed seed is to make failures
    # deterministic; the value is not important.
    state = np.random.default_rng(2020)
    for _ in range(100):  # FIXME: this randomized algorithm is horrendous
        M2real = state.normal() * M2.real + state.normal() * M2.imag
        _, P = np.linalg.eigh(M2real)
        D = P.T.dot(M2).dot(P).diagonal()
        if np.allclose(P.dot(np.diag(D)).dot(P.T), M2, rtol=0, atol=1.0e-13):
            break
    else:
        raise ValueError

    d = -np.angle(D) / 2
    d[3] = -d[0] - d[1] - d[2]
    cs = np.mod((d[:3] + d[3]) / 2, 2 * np.pi)

    # Reorder the eigenvalues to get in the Weyl chamber
    cstemp = np.mod(cs, pi2)
    np.minimum(cstemp, pi2 - cstemp, cstemp)
    order = np.argsort(cstemp)[[1, 2, 0]]
    cs = cs[order]
    d[:3] = d[order]
    P[:, :3] = P[:, order]

    # Fix the sign of P to be in SO(4)
    if np.real(la.det(P)) < 0:
        P[:, -1] = -P[:, -1]

    # Find K1, K2 so that U = K1.A.K2, with K being product of single-qubit unitaries
    K1 = transform_to_magic_basis(Up @ P @ np.diag(np.exp(1j * d)))
    K2 = transform_to_magic_basis(P.T)

    K1l, K1r, phase_l = decompose_two_qubit_product_gate(K1)
    K2l, K2r, phase_r = decompose_two_qubit_product_gate(K2)
    global_phase += phase_l + phase_r

    K1l = K1l.copy()

    # Flip into Weyl chamber
    if cs[0] > pi2:
        cs[0] -= 3 * pi2
        K1l = K1l.dot(_ipy)
        K1r = K1r.dot(_ipy)
        global_phase += pi2
    if cs[1] > pi2:
        cs[1] -= 3 * pi2
        K1l = K1l.dot(_ipx)
        K1r = K1r.dot(_ipx)
        global_phase += pi2
    conjs = 0
    if cs[0] > pi4:
        cs[0] = pi2 - cs[0]
        K1l = K1l.dot(_ipy)
        K2r = _ipy.dot(K2r)
        conjs += 1
        global_phase -= pi2
    if cs[1] > pi4:
        cs[1] = pi2 - cs[1]
        K1l = K1l.dot(_ipx)
        K2r = _ipx.dot(K2r)
        conjs += 1
        global_phase += pi2
        if conjs == 1:
            global_phase -= pi
    if cs[2] > pi2:
        cs[2] -= 3 * pi2
        K1l = K1l.dot(_ipz)
        K1r = K1r.dot(_ipz)
        global_phase += pi2
        if conjs == 1:
            global_phase -= pi
    if conjs == 1:
        cs[2] = pi2 - cs[2]
        K1l = K1l.dot(_ipz)
        K2r = _ipz.dot(K2r)
        global_phase += pi2
    if cs[2] > pi4:
        cs[2] -= pi2
        K1l = K1l.dot(_ipz)
        K1r = K1r.dot(_ipz)
        global_phase -= pi2

    a, b, c = cs[1], cs[0], cs[2]
    return global_phase, (a, b, c), K1l, K1r, K2l, K2r

from qiskit.circuit import Parameter


class ParamIter:
    def __iter__(self):
        self.index = 0
        return self

    def __next__(self):
        x = self.index
        self.index += 1
        return Parameter(f"p{x}")


# print(next(myiter))
# print(next(myiter))

from clonk.utils.riswap_gates.riswap import RiSwapGate

"""
Efficient compiler consturction of parameterized circuits
https://qiskit.org/documentation/tutorials/circuits_advanced/01_advanced_circuits.html

"""


def template_circuit(param_list, base=None):
    myiter = iter(ParamIter())
    if base is None:
        qc = QuantumCircuit(2)
        qc.u(*[next(myiter) for _ in range(3)], 0)
        qc.u(*[next(myiter) for _ in range(3)], 1)
    else:
        qc = base.copy()
    for param in param_list:
        qc.append(RiSwapGate(1 / param), [0, 1])
        qc.u(*[next(myiter) for _ in range(3)], 0)
        qc.u(*[next(myiter) for _ in range(3)], 1)
    return qc


qc = template_circuit([2])
qc.draw()

from qiskit.quantum_info import random_unitary, Operator
from qiskit import QuantumCircuit

target_qc = QuantumCircuit(2)
target_qc.append(random_unitary(dims=(2, 2)), [0, 1])
global_phase, (a, b, c), K1l, K1r, K2l, K2r = KAKDecomp(Operator(target_qc).data)
target_coordinates = (a, b, c)
print(target_coordinates)

def unitary_distance_function(A, B):
    # return (1 - np.abs(np.sum(np.multiply(B,np.conj(np.transpose(A))))) / 4)
    # return (1 - (np.abs(np.sum(np.multiply(B,np.conj(A)))))**2+4 / 20) # quantum volume paper
    return 1 - np.abs(np.sum(np.multiply(B, np.conj(A)))) / 4

%matplotlib inline
from scipy.optimize import minimize


def sample(param_list, N=1000):
    #
    iter = []
    current_best = []
    nuop_distance = []
    #
    minimum_qc = None
    min_distance = None
    modified_template = None

    for template_index in range(0, len(param_list)):
        modified_params = param_list[: 1 + template_index]
        template_qc = template_circuit(modified_params, base=modified_template)
        initial_guess = np.random.random(len(template_qc.parameters)) * 2 * np.pi
        #
        qc = template_circuit(modified_params, base=modified_template)

        def min_fun(x):
            temp_qc = qc.assign_parameters(
                {parameter: i for parameter, i in zip(qc.parameters, x)}
            )
            _, (a, b, c), _, _, _, _ = KAKDecomp(Operator(temp_qc).data)
            temp_coordinates = (a, b, c)
            return np.linalg.norm(
                np.array(target_coordinates) - np.array(temp_coordinates)
            )

        min_result = minimize(fun=min_fun, x0=initial_guess)
        # #
        # for _ in range(N):
        #     qc = template_circuit(modified_params, base=modified_template)
        #     parameters = qc.parameters
        #     temp_qc = qc.assign_parameters({parameter:np.random.random()*2*np.pi for parameter in parameters})
        #     global_phase, (a, b, c), K1l, K1r, K2l, K2r = KAKDecomp(Operator(temp_qc).data)
        #     temp_coordinates = (a,b,c)
        #     temp_dist = np.linalg.norm(np.array(target_coordinates)-np.array(temp_coordinates))
        #     if min_distance is None or temp_dist < min_distance:
        #         minimum_qc = temp_qc
        #         min_distance = temp_dist
        #     #
        #     current_best.append(min_distance)
        #     nuop_distance.append(unitary_distance_function(Operator(minimum_qc).data, Operator(target_qc).data))
        # #

        # modified_template = template_qc.assign_parameters({parameter:i for parameter,i in zip(template_qc.parameters,min_result.x)})
        min_distance = min_result.fun
        print(min_distance)

    modified_template = template_qc.assign_parameters(
        {parameter: i for parameter, i in zip(template_qc.parameters, min_result.x)}
    )

    # import matplotlib.pyplot as plt
    # plt.plot(range(len(current_best)), current_best, label="coords")
    # plt.plot(range(len(current_best)), nuop_distance, label="nuop")
    # plt.yscale('log')
    return modified_template, min_result.fun


minimum_qc, min_distance = sample([2, 4, 8])
print(min_distance)
minimum_qc.draw(output="mpl")

%matplotlib inline
from scipy.optimize import minimize


def sample(param_list, N=1000):
    #
    iter = []
    current_best = []
    nuop_distance = []
    #
    minimum_qc = None
    min_distance = None
    modified_template = None

    for template_index in range(0, len(param_list)):
        modified_params = param_list[template_index : 1 + template_index]
        template_qc = template_circuit(modified_params, base=modified_template)
        initial_guess = np.random.random(len(template_qc.parameters)) * 2 * np.pi
        #
        qc = template_circuit(modified_params, base=modified_template)

        def min_fun(x):
            temp_qc = qc.assign_parameters(
                {parameter: i for parameter, i in zip(qc.parameters, x)}
            )
            _, (a, b, c), _, _, _, _ = KAKDecomp(Operator(temp_qc).data)
            temp_coordinates = (a, b, c)
            return np.linalg.norm(
                np.array(target_coordinates) - np.array(temp_coordinates)
            )

        min_result = minimize(fun=min_fun, x0=initial_guess)
        # #
        # for _ in range(N):
        #     qc = template_circuit(modified_params, base=modified_template)
        #     parameters = qc.parameters
        #     temp_qc = qc.assign_parameters({parameter:np.random.random()*2*np.pi for parameter in parameters})
        #     global_phase, (a, b, c), K1l, K1r, K2l, K2r = KAKDecomp(Operator(temp_qc).data)
        #     temp_coordinates = (a,b,c)
        #     temp_dist = np.linalg.norm(np.array(target_coordinates)-np.array(temp_coordinates))
        #     if min_distance is None or temp_dist < min_distance:
        #         minimum_qc = temp_qc
        #         min_distance = temp_dist
        #     #
        #     current_best.append(min_distance)
        #     nuop_distance.append(unitary_distance_function(Operator(minimum_qc).data, Operator(target_qc).data))
        # #

        modified_template = template_qc.assign_parameters(
            {parameter: i for parameter, i in zip(template_qc.parameters, min_result.x)}
        )
        min_distance = min_result.fun
        print(min_distance)

    modified_template = template_qc.assign_parameters(
        {parameter: i for parameter, i in zip(template_qc.parameters, min_result.x)}
    )

    # import matplotlib.pyplot as plt
    # plt.plot(range(len(current_best)), current_best, label="coords")
    # plt.plot(range(len(current_best)), nuop_distance, label="nuop")
    # plt.yscale('log')
    return modified_template, min_result.fun


minimum_qc, min_distance = sample([2, 2, 2])
print(min_distance)
minimum_qc.draw(output="mpl")

def sample(param_list, N=1000):
    # want to characterize the vectors from 2,2, template
    qc = template_circuit(param_list[:1])
    temp_qc = qc.assign_parameters(
        {parameter: np.random.random() * 2 * np.pi for parameter in qc.parameters}
    )
    global_phase, (a, b, c), K1l, K1r, K2l, K2r = KAKDecomp(Operator(temp_qc).data)
    starting_coordinates = (a, b, c)

    list_coordinates = []

    for _ in range(N):
        qc = template_circuit(param_list)
        parameters = qc.parameters
        temp_qc = qc.assign_parameters(
            {parameter: np.random.random() * 2 * np.pi for parameter in parameters}
        )
        global_phase, (a, b, c), K1l, K1r, K2l, K2r = KAKDecomp(Operator(temp_qc).data)
        temp_coordinates = (a, b, c)
        list_coordinates.append(temp_coordinates)
        c1, c2, c3 = weylchamber.c1c2c3(gate)
        w.add_point(c1, c2, c3)
        temp_vector = np.array(temp_coordinates) - np.array(starting_coordinates)
        # print(temp_vector)

    return list_coordinates


w = WeylChamber()
list_coordinates = sample(param_list=[3, 2])
w.plot()
# print(list_coordinates)

import scipy.spatial as ss
import numpy as np

hull = ss.ConvexHull(list_coordinates)
print("volume inside points is: ", hull.volume)

import numpy as np
import qutip
import matplotlib
import matplotlib.pylab as plt
import weylchamber
from weylchamber.visualize import WeylChamber

WeylChamber().plot()

IDENTITY = qutip.identity([2, 2])
CNOT = qutip.qip.operations.cnot()
CPHASE = qutip.qip.operations.cphase(np.pi)
BGATE = qutip.qip.operations.berkeley()
iSWAP = qutip.qip.operations.iswap()
sqrtISWAP = qutip.qip.operations.sqrtiswap()
sqrtSWAP = qutip.qip.operations.sqrtswap()
MGATE = weylchamber.canonical_gate(3 / 4, 1 / 4, 0)

w = WeylChamber()
list_of_gates = [
    ("Identity", IDENTITY),
    ("CNOT", CNOT),
    ("CPHASE", CPHASE),
    ("BGATE", BGATE),
    ("iSWAP", iSWAP),
    ("sqrtISWAP", sqrtISWAP),
    ("sqrtSWAP", sqrtSWAP),
    ("MGATE", MGATE),
]
print("Weyl Chamber Coordinates")
print("----------------------------------")
for name, gate in list_of_gates:
    c1, c2, c3 = weylchamber.c1c2c3(gate)
    print("%10s: \t%.2fπ %.2fπ %.2fπ" % (name, c1, c2, c3))
    w.add_point(c1, c2, c3)
w.plot()

w.scatter(
    *zip(
        *[
            weylchamber.c1c2c3(qutip.qip.operations.cphase(phase))
            for phase in np.linspace(0, 2 * np.pi, 20)
        ]
    )
)
w.plot()

print("Local Invariants")
print("----------------------------------")
for name, gate in list_of_gates:
    g1, g2, g3 = weylchamber.g1g2g3(gate)
    print("%10s: \t%5.2f %5.2f %5.2f" % (name, g1, g2, g3))