import numpy as np
from qiskit import __qiskit_version__
from qiskit.chemistry.drivers import PySCFDriver, PyQuanteDriver, UnitsType
from qiskit.chemistry import FermionicOperator
from qiskit.aqua.algorithms import VQE, ExactEigensolver
from qiskit.chemistry.components.variational_forms import UCCSD
from qiskit.chemistry.components.initial_states import HartreeFock
from qiskit.aqua.components.optimizers import COBYLA
from qiskit.aqua import QuantumInstance
from qiskit import Aer
from qiskit.chemistry.algorithms.q_equation_of_motion.q_equation_of_motion import QEquationOfMotion
from qiskit.visualization import plot_histogram
from qiskit import IBMQ
from qiskit.providers.aer import noise

__qiskit_version__

driver = PySCFDriver(atom='H .0 .0 .0; H .0 .0 0.735', unit=UnitsType.ANGSTROM,
                     charge=0, spin=0, basis='sto3g')
molecule = driver.run()

h1 = molecule.one_body_integrals
h2 = molecule.two_body_integrals

nuclear_repulsion_energy = molecule.nuclear_repulsion_energy

num_particles = molecule.num_alpha + molecule.num_beta
num_spin_orbitals = molecule.num_orbitals * 2

print("HF energy: {}".format(molecule.hf_energy - molecule.nuclear_repulsion_energy))
print("# of electrons: {}".format(num_particles))
print("# of spin orbitals: {}".format(num_spin_orbitals))

ferop = FermionicOperator(h1=h1, h2=h2)

map_type = 'jordan_wigner'
qubitop = ferop.mapping(map_type=map_type)
print(qubitop.print_details())

exact_eigensolver = ExactEigensolver(qubitop, k=1<<qubitop.num_qubits)
ret = exact_eigensolver.run()
print('The ground state energy is: {:.12f}'.format(ret['eigvals'][0].real))
print('The ground state is:')
for i in range(1<<qubitop.num_qubits):
    print(np.binary_repr(i, width=qubitop.num_qubits), ret['eigvecs'][0][i])

qubit_reduction = False
HF_state = HartreeFock(qubitop.num_qubits, num_spin_orbitals, num_particles, map_type, qubit_reduction)

var_form = UCCSD(qubitop.num_qubits, depth=1, 
                   num_orbitals=num_spin_orbitals, num_particles=num_particles, 
                   initial_state=HF_state, qubit_mapping=map_type, 
                   two_qubit_reduction=qubit_reduction)

params = np.random.rand(var_form._num_parameters)
circuit = var_form.construct_circuit(params)
circuit.decompose().draw(output='mpl')

backend = Aer.get_backend('statevector_simulator')
import logging
from qiskit.chemistry import set_qiskit_chemistry_logging
set_qiskit_chemistry_logging(logging.INFO) # choose among DEBUG, INFO, WARNING, ERROR, CRITICAL and NOTSET


max_eval = 200
cobyla = COBYLA(maxiter=max_eval)

vqe = VQE(qubitop, var_form, cobyla)
quantum_instance = QuantumInstance(backend=backend)

results = vqe.run(quantum_instance)
print('The computed ground state energy is: {:.12f}'.format(results['eigvals'][0]))
print("Parameters: {}".format(results['opt_params']))

h1_nop = np.identity(qubitop.num_qubits)

fermionic_nop = FermionicOperator(h1=h1_nop)
qubit_nop = fermionic_nop.mapping(map_type=map_type)
print(qubit_nop.print_details())

wavefn = var_form.construct_circuit(results['opt_params'])
statevector_mode = True
use_simulator_snapshot_mode = True
circuits = qubit_nop.construct_evaluation_circuit(wavefn, statevector_mode, 
                                                  use_simulator_snapshot_mode=use_simulator_snapshot_mode)
results_2 = quantum_instance.execute(circuits)
avg, std = qubit_nop.evaluate_with_result(results_2, statevector_mode, 
                                          use_simulator_snapshot_mode=use_simulator_snapshot_mode)

print(round(np.real(avg),1))

eom =  QEquationOfMotion(qubitop, qubitop.num_qubits, num_particles,
                 qubit_mapping=map_type)
gaps, _ = eom.calculate_excited_states(wavefn, quantum_instance = quantum_instance)

exact_energies = []
for i in range(1<<qubitop.num_qubits):
    state = ret['eigvecs'][i]
    nop_exp = qubit_nop.evaluate_with_statevector(state)
    if round(np.real(nop_exp[0]),1) == 2:
        exact_energies.append(np.real(ret['eigvals'][i]))

qeom_energies = [results['eigvals'][0]]
for i in range(len(gaps)):
    qeom_energies.append(qeom_energies[0]+gaps[i])
    
print('The exact energies are {}'.format(exact_energies))
print('The qEOM energies are {}'.format(qeom_energies))

backend_2 = Aer.get_backend('qasm_simulator')
shots = 1000
quantum_instance_2 = QuantumInstance(backend=backend_2, shots = shots)

results_3 = vqe.run(quantum_instance_2)

print('The computed ground state energy is: {:.12f}'.format(results_3['eigvals'][0]))
print("Parameters: {}".format(results_3['opt_params']))

plot_histogram(results_3['eigvecs'][0], color='midnightblue', title="H2 ground state, 1000 shots")

IBMQ.save_account('679268b65869103438b43782543ef1cb532068e8dbc53d8acc34f160430eeb37655fbaf0351df831623d8bd5cf7496f8e21bd266cab2a963342dd737775465db', overwrite=True)
provider = IBMQ.load_account()

print(provider.backends())

device = provider.get_backend('ibmqx2')
properties = device.properties()
coupling_map = device.configuration().coupling_map
noise_model = noise.device.basic_device_noise_model(properties)
basis_gates = noise_model.basis_gates

quantum_instance_3 = QuantumInstance(backend_2, shots=shots, basis_gates=basis_gates,
                                   coupling_map=coupling_map, noise_model=noise_model)

results_4 = vqe.run(quantum_instance_3)

print('The computed ground state energy is: {:.12f}'.format(results_3['eigvals'][0]))
print("Parameters: {}".format(results_3['opt_params']))

plot_histogram(results_4['eigvecs'][0], color='midnightblue', title="H2 ground state, 1000 shots")





