import pennylane as qml
import jax
from jax import numpy as np

jax.config.update("jax_platform_name", "cpu")

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

x = np.array([0.1, 0.2, 0.3])

@qml.qnode(dev, diff_method="adjoint")
def circuit(a):
    qml.RX(a[0], wires=0)
    qml.CNOT(wires=(0,1))
    qml.RY(a[1], wires=1)
    qml.RZ(a[2], wires=1)
    return qml.expval(qml.PauliX(wires=1))

n_gates = 4
n_params = 3

ops = [
    qml.RX(x[0], wires=0),
    qml.CNOT(wires=(0,1)),
    qml.RY(x[1], wires=1),
    qml.RZ(x[2], wires=1)
]
M = qml.PauliX(wires=1)

from pennylane.devices.qubit import create_initial_state, apply_operation

state = create_initial_state((0, 1))

for op in ops:
    state = apply_operation(op, state)

print(state)


bra = apply_operation(M, state)
ket = state

M_expval = np.vdot(bra, ket)
print("vdot  : ", M_expval)
print("QNode : ", circuit(x))

bra_n = create_initial_state((0, 1))

for op in ops:
    bra_n = apply_operation(op, bra_n)
bra_n = apply_operation(M, bra_n)
bra_n = apply_operation(qml.adjoint(ops[-1]), bra_n)

ket_n = create_initial_state((0, 1))

for op in ops[:-1]: # don't apply last operation
    ket_n = apply_operation(op, ket_n)

M_expval_n = np.vdot(bra_n, ket_n)
print(M_expval_n)

bra_n_v2 = apply_operation(M, state)
ket_n_v2 = state

adj_op = qml.adjoint(ops[-1])

bra_n_v2 = apply_operation(adj_op, bra_n_v2)
ket_n_v2 = apply_operation(adj_op, ket_n_v2)

M_expval_n_v2 = np.vdot(bra_n_v2, ket_n_v2)
print(M_expval_n_v2)

bra_loop = apply_operation(M, state)
ket_loop = state

for op in reversed(ops):
    adj_op = qml.adjoint(op)
    bra_loop = apply_operation(adj_op, bra_loop)
    ket_loop = apply_operation(adj_op, ket_loop)
    print(np.vdot(bra_loop, ket_loop))

grad_op0 = qml.operation.operation_derivative(ops[0])
print(grad_op0)

bra = apply_operation(M, state)
ket = state

grads = []

for op in reversed(ops):
    adj_op = qml.adjoint(op)
    ket = apply_operation(adj_op, ket)

    # Calculating the derivative
    if op.num_params != 0:
        dU = qml.operation.operation_derivative(op)
        ket_temp = apply_operation(qml.QubitUnitary(dU, op.wires), ket)

        dM = 2 * np.real(np.vdot(bra, ket_temp))
        grads.append(dM)

    bra = apply_operation(adj_op, bra)


# Finally reverse the order of the gradients
# since we calculated them in reverse
grads = grads[::-1]

print("our calculation: ", [float(grad) for grad in grads])

grad_compare = jax.grad(circuit)(x)
print("comparison: ", grad_compare)

dev_lightning = qml.device('lightning.qubit', wires=2)

@qml.qnode(dev_lightning, diff_method="adjoint")
def circuit_adjoint(a):
    qml.RX(a[0], wires=0)
    qml.CNOT(wires=(0,1))
    qml.RY(a[1], wires=1)
    qml.RZ(a[2], wires=1)
    return qml.expval(M)

print(jax.grad(circuit_adjoint)(x))