import qiskit
import numpy as np
from qiskit.aqua.operators import PauliExpectation, CircuitSampler, ExpectationFactory
from qiskit.aqua.operators import CircuitStateFn, StateFn, ListOp

import sys
sys.path.append('../../')
from typing import Union
import qiskit
import numpy as np
from qiskit import Aer, execute
from qiskit.aqua.utils.backend_utils import is_aer_provider
from qiskit.aqua.operators import (PauliExpectation, CircuitSampler, ExpectationFactory,
                                   CircuitStateFn, StateFn, ListOp)

import sys
sys.path.append('../')

def sample_hamiltonian(hamiltonian: qiskit.aqua.operators.OperatorBase,
                       backend: Union[qiskit.providers.BaseBackend, qiskit.aqua.QuantumInstance],
                       ansatz: qiskit.QuantumCircuit):
    
    if qiskit.aqua.quantum_instance.QuantumInstance == type(backend):
        sampler = CircuitSampler(backend, param_qobj=is_aer_provider(backend.backend))
    else:
        sampler = CircuitSampler(backend)
        
    expectation = ExpectationFactory.build(operator=hamiltonian,backend=backend)
    observable_meas = expectation.convert(StateFn(hamiltonian, is_measurement=True))
    
    ansatz_circuit_op = CircuitStateFn(ansatz)
    
    expect_op = observable_meas.compose(ansatz_circuit_op).reduce()
    sampled_expect_op = sampler.convert(expect_op)
    
    return np.real(sampled_expect_op.eval())

from VOLTA.Ansatz import get_num_parameters, get_var_form
from qiskit.aqua.operators import I, X, Y, Z
from qiskit.aqua import QuantumInstance

n_qubits = 2
n_params = get_num_parameters(n_qubits)
params = [0.]*n_params
wavefunction = get_var_form(params, n_qubits)

hamiltonian = 1/2*((Z^Z) + (Z^I) + (X^X))
hamiltonian = X^X
backend = QuantumInstance(backend=qiskit.BasicAer.get_backend('qasm_simulator'),
                          shots=10000)
sample_hamiltonian(hamiltonian, backend, wavefunction)

# Variational Form

backend = qiskit.BasicAer.get_backend('qasm_simulator')
sampler = CircuitSampler(backend=backend)

expectation = ExpectationFactory.build(operator=hamiltonian,backend=backend)
observable_meas = expectation.convert(StateFn(hamiltonian, is_measurement=True))


print(observable_meas)

ansatz_circuit_op = CircuitStateFn(wavefunction)
expect_op = observable_meas.compose(ansatz_circuit_op).reduce()
sampled_expect_op = sampler.convert(expect_op)
means = np.real(sampled_expect_op.eval())

means

