import marimo as mo
import numpy as np
from qiskit import QuantumCircuit, transpile
from qiskit_aer import StatevectorSimulator

circuit_selector = mo.ui.dropdown(
    options={
        "Circuit Phi_A": "circuit_a",
        "Circuit Phi_B": "circuit_b",
        "Circuit Phi_C": "circuit_c"
    },
    value="Circuit Phi_A",
    label="Change of basis circuit",
    full_width=True,
)

theta_selector = mo.ui.slider(
    start=0.0,
    stop=3.14159,
    step=0.01,
    value=0.5,
    label="θ"
)

phi_selector = mo.ui.slider(
    start=0.0,
    stop=np.pi,
    step=0.01,
    value=0.0,
    label="φ"
)

psi_selector = mo.ui.slider(
    start=0.0,
    stop=np.pi,
    step=0.01,
    value=0.0,
    label="ψ"
)

mo.vstack([
    circuit_selector,
     mo.vstack([theta_selector, phi_selector, psi_selector]),
], gap="1rem")

def rotate_x(theta):
    return np.array([[1, 0, 0],
                     [0, np.cos(theta), -np.sin(theta)],
                     [0, np.sin(theta), np.cos(theta)]])

def rotate_y(psi):
    return np.array([[np.cos(psi), 0, np.sin(psi)],
                     [0, 1, 0],
                     [-np.sin(psi), 0, np.cos(psi)]])

def rotate_z(phi):
    return np.array([[np.cos(phi), -np.sin(phi), 0],
                     [np.sin(phi), np.cos(phi), 0],
                     [0, 0, 1]])

def rotate_x_z_y(theta, phi, psi):
    return rotate_x(theta) @ rotate_z(phi) @ rotate_y(psi)

def print_result(circuit, theta, phi, psi, index):
    output = ""

    output += "Quantum Circuit to build state |Phi_C^0'(theta, phi, psi)> after change of basis\n"
    output += "Basis states:\n"
    output += "  |Phi_A^0> = |00010011>\n"
    output += "  |Phi_B^0> = |00010101>\n"
    output += "  |Phi_C^0> = |00011001>\n\n"

    output += "Quantum Circuit:\n"
    output += str(circuit)

    simulator = StatevectorSimulator()
    circuit = transpile(circuit, simulator)
    result = simulator.run(circuit).result()
    statevector = result.get_statevector(circuit)
    statevector_dictionnary = statevector.to_dict()
    rotation_matrix = rotate_x_z_y(theta, phi, psi)

    output += "\n\nAmplitude comparison:\n"

    for i, state in enumerate(['00010011', '00010101', '00011001']):
        theoretical_amplitude = rotation_matrix[i, index]
        circuit_amplitude = statevector_dictionnary.get(state, 0.0)

        output += f"\nState: {state}\n"
        output += f"Theoretical amplitude: {theoretical_amplitude}\n"
        output += f"Circuit amplitude: {circuit_amplitude}\n"
        output += f"Absolute error: {abs(theoretical_amplitude-circuit_amplitude):.3e}\n"

    return mo.md(f"```\n{output}\n```")
    

def build_circuit_phi_a_prime(theta, phi, psi): 

    circuit = QuantumCircuit(8)
    circuit.x(0)
    circuit.x(4)
    circuit.ry(2 * psi, 2)
    circuit.cx(2, 1)
    circuit.x(1)
    circuit.cry(2 * phi, 1, 3)
    circuit.cx(3, 1)
    circuit.cx(2, 3)
    circuit.cx(3, 2)
    circuit.cry(2 * theta, 3, 2)
    circuit.cx(2, 3)
    circuit.z(3)

    return print_result(circuit, theta, phi, psi, 0)

def build_circuit_phi_b_prime(theta, phi, psi): 

    circuit = QuantumCircuit(8)
    circuit.x(0)
    circuit.x(4)
    circuit.ry(2 * phi ,1)
    circuit.cx(1, 2)
    circuit.x(2)
    circuit.cry(2 * theta, 2, 3)
    circuit.cx(3, 2)
    circuit.z(1)

    return print_result(circuit, theta, phi, psi, 1)

def build_circuit_phi_c_prime(theta, phi, psi): 


    circuit = QuantumCircuit(8)
    circuit.x(0)
    circuit.x(4)

    circuit.ry(2 * psi, 1)
    circuit.cx(1, 3)
    circuit.x(3)
    circuit.cry(2 * phi, 1, 2)
    circuit.cx(2, 1)
    circuit.cx(3, 2)
    circuit.cry(2 * theta, 2, 3)
    circuit.cx(3,2)
    
    return print_result(circuit, theta, phi, psi, 2)
    

active_circuit = circuit_selector.value
theta_value = theta_selector.value
phi_value = phi_selector.value
psi_value = psi_selector.value

print(psi_value)
if active_circuit == "circuit_a":
      output = build_circuit_phi_a_prime(theta_value, phi_value, psi_value)
elif active_circuit  == "circuit_b":
      output = build_circuit_phi_b_prime(theta_value, phi_value, psi_value)
elif active_circuit  == "circuit_c":
      output = build_circuit_phi_c_prime(theta_value, phi_value, psi_value)


output