import pennylane as qml
from pennylane import numpy as np
import matplotlib.pyplot as plt
from enum import Enum

qml.drawer.use_style('sketch')

class ReOrIm(Enum):
    real = 0
    imag = 1

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

@qml.qnode(dev)
def hadamard_test_circuit(theta, phi, re_or_im=ReOrIm.real):
    # Prepare the superposition state |psi> = (|0>|a> + |1>|b>)/sqrt(2)
    # 1. Apply the Hadamard gate to the first qubit
    qml.Hadamard(wires=0)

    # 2. Apply S^dagger to the first qubit if we are measuring the imaginary part
    if re_or_im == ReOrIm.imag:
        qml.adjoint(qml.S)(wires=0)
        
    # 3. Apply the controlled unitary U where U|0> = |a> and U|1> = |b>
    qml.RY(2 * theta, wires=1) # |a> = Ry(theta)|0>
    qml.CRY(2 * (phi-theta), wires=[0, 1])

    # 4. Apply the Hadamard gate to the first qubit
    qml.Hadamard(wires=0)
    return qml.probs(wires=0)

qml.draw_mpl(hadamard_test_circuit)(np.pi/4, np.pi/3)
plt.show()

qml.draw_mpl(hadamard_test_circuit)(np.pi/4, np.pi/3, re_or_im=ReOrIm.imag)
plt.show()

p_0_real = hadamard_test_circuit(np.pi/4, np.pi/3, re_or_im=ReOrIm.real)[0]
p_0_imag = hadamard_test_circuit(np.pi/4, np.pi/3, re_or_im=ReOrIm.imag)[0]

print(f"Probability of measuring |0> when measuring the real part: {p_0_real:.4f}")
print(f"Probability of measuring |0> when measuring the imaginary part: {p_0_imag:.4f}")

inner_prod_real = 2 * p_0_real - 1
inner_prod_imag = 1 - 2 * p_0_imag

print(f"Real part of the inner product: {inner_prod_real:.4f}")
print(f"Imaginary part of the inner product: {inner_prod_imag:.4f}")

# The inner product of the two states is the sum of the real and imaginary parts
inner_prod = inner_prod_real + 1j * inner_prod_imag
print(f"Inner product of the two states: {inner_prod:.4f}")