import numpy as np
import matplotlib as plt
%matplotlib inline

input_param=[1, 0.5, -0.765, 0.1, 0, -0.654]
exp_output=[0, 0, 0, 0, -0.4553474723, 0]

from qiskit.circuit import QuantumCircuit, ParameterVector

n = 3
qc = QuantumCircuit(n)

param_list = ParameterVector('Parameter', 2*n)

for i in range(len(param_list)//n):
    qc.rx(param_list[3*i], 0)
    qc.ry(param_list[3*i+1], 1)
    qc.rz(param_list[3*i+2], 2)

    qc.cnot(0, 1)
    qc.cnot(1, 2)
    qc.cnot(2, 0)

qc.draw(output='mpl')

from qiskit.opflow import I, X, Y, Z
hamiltonian = Z ^ I ^ Y

from qiskit.opflow import StateFn, AerPauliExpectation

expectation = StateFn(hamiltonian, is_measurement=True) @ StateFn(qc)
pauli_basis = AerPauliExpectation().convert(expectation)

from qiskit import Aer
from qiskit.utils import QuantumInstance
from qiskit.opflow import CircuitSampler

quantum_instance = QuantumInstance(Aer.get_backend('qasm_simulator'),
                                   # we'll set a seed for reproducibility
                                   shots = 32768, seed_simulator = 2718,
                                   seed_transpiler = 2718)
sampler = CircuitSampler(quantum_instance)

def evaluate_expectation(params):
    value_dict = dict(zip(qc.parameters, params))
    result = sampler.convert(pauli_basis, params=value_dict).eval()
    return np.real(result)

def calc_gradient(params):
    gradient = np.zeros_like(params)
    for i in range(len(params)):
        shifted = params.copy()
        shifted[i] += np.pi/2
        forward = evaluate_expectation(shifted)

        shifted[i] -= np.pi
        backward = evaluate_expectation(shifted)

        gradient[i] = 0.5 * (forward - backward)
    return np.round(gradient, 10)

calc_gradient(input_param)

from qiskit.opflow import Gradient

def calc_gradient_with_framework(params):
    state_grad = Gradient(grad_method="param_shift").convert(operator=expectation)
    value_dict = dict(zip(qc.parameters, params))
    result = state_grad.assign_parameters(value_dict).eval()
    return np.round(np.array(result).real, 10)


calc_gradient_with_framework(input_param)

import qiskit.tools.jupyter
%qiskit_version_table

