import pennylane as qml
from pennylane import numpy as np
import matplotlib.pyplot as plt
import scienceplots

qml.drawer.use_style(style='sketch')

def unitary_matrix_from_hermitian(A):
    eigenvalues, eigenvectors = np.linalg.eigh(A)
    D = np.diag(np.exp(2j * np.pi * eigenvalues))
    U = np.dot(eigenvectors, np.dot(D, eigenvectors.conj().T))
    return U

def MultipleControlledRY(target_wire, control_wires, lamb1, lamb2):
    """
    Apply a multiple controlled RY gate to the target wire, with control on the control wires.

    Args:
        target_wire (int): the target wire the RY gate acts on
        control_wires (list[int]): the control wires
        lamb (float): the eigenvalue. lamb should be converted to the qubit state (e.g. 0.25 (n=2) -> 01)
    """
    n = len(control_wires)
    angle1 = 2 * np.arcsin(lamb1)
    lamb1_bin = format(int(lamb1 * 2**n), f'0{n}b')
    control_values1 = [int(bit) for bit in lamb1_bin]

    angle2 = 2 * np.arcsin(lamb2)
    lamb2_bin = format(int(lamb2 * 2**n), f'0{n}b')
    control_values2 = [int(bit) for bit in lamb2_bin]
    # print(control_values)

    qml.ctrl(qml.RY, control=control_wires, control_values=control_values1)(angle1, wires=target_wire)
    qml.ctrl(qml.RY, control=control_wires, control_values=control_values2)(angle2, wires=target_wire)

dev = qml.device('default.qubit', wires=3)

@qml.qnode(dev)
def aqe_test(lamb1, lamb2):
    qml.PauliX(wires=2)
    MultipleControlledRY(0, [1, 2], lamb1, lamb2)
    return qml.state()

lamb1 = 0.75
lamb2 = 0.25

qml.draw_mpl(aqe_test)(lamb1, lamb2)
plt.show()

state = aqe_test(lamb1, lamb2)
print(state)

@qml.qnode(dev)
def aqe_test(lamb1, lamb2):
    qml.PauliX(wires=1)
    qml.PauliX(wires=2)
    MultipleControlledRY(0, [1, 2], lamb1, lamb2)
    return qml.state()

state = aqe_test(lamb1, lamb2)
print(state)

qml.draw_mpl(aqe_test)(lamb1, lamb2)
plt.show()

A = np.array([[0.5, 0.25], [0.25, 0.5]])
U = unitary_matrix_from_hermitian(A)
x = np.array([1, 0])

dev = qml.device('default.qubit', wires=3)

@qml.qnode(dev)
def qpe_test(U, x):
    # Encode the input vector x into the amplitudes of the state (here, only 2d vectors)
    qml.AmplitudeEmbedding(features=x, wires=2, normalize=True)

    # Apply Hadamard gate to the first n qubits
    qml.Hadamard(wires=0)
    qml.Hadamard(wires=1)

    # Apply controlled unitary operations
    qml.ControlledSequence(qml.QubitUnitary(U, wires=2), control=range(2))

    # Apply inverse QFT
    qml.adjoint(qml.QFT)(wires=range(2))

    # Measure c-register and return the probability distribution
    return qml.state()

qml.draw_mpl(qpe_test)(U, x)
plt.show()

qpe_test(U, x).round(2)

def gen_qmm(n: int, A):
    U = unitary_matrix_from_hermitian(A)
    U_inv = np.linalg.inv(U)
    dev1 = qml.device("default.qubit", wires=n+1) # c-register, encoding x
    dev2 = qml.device("lightning.qubit", wires=n+2, shots=10000, mcmc=True) # ancilla, c-register, encoding x
    
    @qml.qnode(dev1)
    def eigenvalue_circuit(x):
        # Encode the input vector x into the amplitudes of the state (here, only 2d vectors)
        qml.AmplitudeEmbedding(features=x, wires=n, normalize=True)

        # Apply Hadamard gate to the first n qubits
        for i in range(n):
            qml.Hadamard(wires=i)

        # Apply controlled unitary operations
        qml.ControlledSequence(qml.QubitUnitary(U, wires=n), control=range(n))

        # Apply inverse QFT
        qml.adjoint(qml.QFT)(wires=range(n))

        # Measure c-register and return the probability distribution
        return qml.probs(wires=range(n))
    
    @qml.qnode(dev2)
    def final_circuit(x, lamb1, lamb2):
        # Encode the input vector x into the amplitudes of the state (here, only 2d vectors)
        qml.AmplitudeEmbedding(features=x, wires=n+1, normalize=True)

        # Apply Hadamard gate to the c-register
        for i in range(1, n+1):
            qml.Hadamard(wires=i)

        # Apply controlled unitary operations
        qml.ControlledSequence(qml.QubitUnitary(U, wires=n+1), control=range(1,n+1))

        # Apply inverse QFT
        qml.adjoint(qml.QFT)(wires=range(1,n+1))

        # Apply controlled rotations for |lambda> -> |lambda> (sqrt(1-lambda^2) |0> + lambda |1>)
        MultipleControlledRY([0], range(1, n+1), lamb1, lamb2)

        # Apply QFT
        qml.QFT(wires=range(1,n+1))

        # Apply controlled unitary operations
        qml.ControlledSequence(qml.QubitUnitary(U_inv, wires=n+1), control=range(1,n+1))

        # Apply Hadamard gate to the c-register
        for i in range(1, n+1):
            qml.Hadamard(wires=i)
        
        # Return the state of the last qubit and the measurement result
        return qml.sample(wires=[0, n+1])
    
    def circuit(x):
        lamb_prob = eigenvalue_circuit(x)
        # lamb_cand = np.arange(2**n) / 2**n
        sorted_lamb_prob = np.argsort(lamb_prob)[::-1]
        lamb1 = sorted_lamb_prob[0] / 2**(n)
        lamb2 = sorted_lamb_prob[1] / 2**(n)
        print(lamb1, lamb2)
        samples = final_circuit(x, lamb1, lamb2)
        return samples
    
    return circuit, (dev1, eigenvalue_circuit), (dev2, final_circuit)

A = np.array([[0.5, 0.25], [0.25, 0.5]])
x = np.array([1, 2])

A @ x

eigv = np.linalg.eig(A)
print(eigv)

qmm_circuit, (dev1, ec), (dev2, fc) = gen_qmm(2, A)

qml.draw_mpl(ec)(x)
plt.show()

qml.draw_mpl(fc)(x, 0.25, 0.75)
plt.show()

result = qmm_circuit(x)

result

meas_one = result[:,0] == 1
result_zero = result[meas_one & (result[:,1] == 0)]
result_one = result[meas_one & (result[:,1] == 1)]

from sklearn.metrics import confusion_matrix

cm = confusion_matrix(result[:,0], result[:,1])
cm

import matplotlib.pyplot as plt
import seaborn as sns

sns.heatmap(cm, annot=True, fmt='d')
plt.show()

res = np.array([result_zero.shape[0], result_one.shape[0]])
res

np.sqrt(res) / np.linalg.norm(np.sqrt(res))

B = A @ x
B / np.linalg.norm(B)