from mpqp.execution.vqa import *

from mpqp import QCircuit
from mpqp.gates import *
from sympy import symbols

x, y, z = symbols("x y z")
circuit = QCircuit([Rx(x, 0), Ry(y, 1), Rz(z,0), Rz(z,1), CNOT(0,1)])
print(circuit)

from mpqp.measures import Observable, ExpectationMeasure
from mpqp.execution import IBMDevice, run

import numpy as np
matrix = np.array([[4,  2,  3, 8],
                   [2, -3,  1, 0],
                   [3,  1, -1, 5],
                   [8,  0,  5, 2]])
hamiltonian = Observable(matrix)

circuit.add(ExpectationMeasure([0, 1], observable=hamiltonian, shots=0))

minimize(circuit, Optimizer.COBYLA, IBMDevice.AER_SIMULATOR, optimizer_options={"maxiter":200})

circuit2 = circuit.without_measurements()
def cost_function(params):
    r1 = run(circuit2, IBMDevice.AER_SIMULATOR_STATEVECTOR, {x: params[0], y: params[1], z: params[2]})
    r2 = run(circuit, IBMDevice.AER_SIMULATOR, {x: params[0], y: params[1], z: params[2]})
    return abs(r1.amplitudes[0]) - np.sqrt(r2.expectation_value**3)

minimize(cost_function, Optimizer.COBYLA, nb_params=3, optimizer_options={"maxiter":200})