# import common packages

import numpy as np
import matplotlib.pyplot as plt

from qiskit import Aer, IBMQ, QuantumRegister
from qiskit.providers.ibmq import least_busy
from qiskit.providers.aer import noise

# lib from Qiskit Aqua

from qiskit.aqua import Operator, QuantumInstance
from qiskit.aqua.algorithms import VQE, ExactEigensolver
from qiskit.aqua.components.optimizers import COBYLA, SPSA, L_BFGS_B
from qiskit.aqua.components.variational_forms import RY, RYRZ, SwapRZ


# lib from Qiskit Aqua Chemistry

from qiskit.chemistry import QiskitChemistry
from qiskit.chemistry import FermionicOperator
from qiskit.chemistry.drivers import PySCFDriver, UnitsType
from qiskit.chemistry.aqua_extensions.components.variational_forms import UCCSD
from qiskit.chemistry.aqua_extensions.components.initial_states import HartreeFock

# setup qiskit.chemistry logging

import logging
from qiskit.chemistry import set_qiskit_chemistry_logging

# function that allows us to create multiple figures in single notebook cell
def show_figure(fig):
    dummy_fig = plt.figure()
    new_dummy_fig_manager = dummy_fig.canvas.manager
    new_dummy_fig_manager.canvas.figure = fig
    fig.set_canvas(new_dummy_fig_manager.canvas)

set_qiskit_chemistry_logging(logging.INFO) # choose among DEBUG, INFO, WARNING, ERROR, CRITICAL and NOTSET

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

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))

print("Contents of molecule object: \n{}".format(sorted(molecule.__dir__())))

# get the integrals from molecule object
h1 = molecule.one_body_integrals
h2 = molecule.two_body_integrals

# create a fermionic operator
ferOp = FermionicOperator(h1=h1, h2=h2)

map_type = 'parity' # 'jordan_wigner', 'parity', 'bravyi_kitaev', 'bksf'

# if 'parity' mapping is selected, reduce the number of qubits (-2 qubits)
qubit_reduction = True if map_type == 'parity' else False

qubitOp = ferOp.mapping(map_type=map_type, threshold=10**-10) # paulis with coefs below threshold are not considered
qubitOp = qubitOp.two_qubit_reduced_operator(num_particles) if qubit_reduction else qubitOp
qubitOp.chop(10**-10) # cut out the pauli strings with prefacator smaller then

print(qubitOp)
qubitOp.paulis

# to see all the contents of given class
# qubitOp.__dir__()
# one can group the Paulis
# qubitOp.to_grouped_paulis()
# qubitOp.paulis

for pauli in qubitOp.paulis:
    print(pauli[0],str(pauli[1]))

max_iter = 200 # maximum number of iterations

# optimizer = COBYLA(maxiter=max_iter)
# optimizer = SPSA(maxiter=max_iter)
optimizer = L_BFGS_B(maxiter=max_iter)

HF_state = HartreeFock(qubitOp.num_qubits, num_spin_orbitals, num_particles, map_type, 
                       qubit_reduction)
HF_state_vec = HF_state.construct_circuit(mode='vector')

# we can use qiskit's vizualisation module to plot the state
from qiskit.visualization import plot_state_city
plot_state_city(HF_state_vec)


# you might want to visualise on the Bloch sphere for a non-superposition state
from qiskit.visualization import plot_state_city, plot_bloch_multivector
plot_bloch_multivector(HF_state_vec)

var_form_uccsd = UCCSD(qubitOp.num_qubits, depth=1, 
                   num_orbitals=num_spin_orbitals, num_particles=num_particles, 
                   active_occupied=None, active_unoccupied=None, # you could modify the active space here 
                   initial_state=HF_state, qubit_mapping=map_type,
                   two_qubit_reduction=qubit_reduction, num_time_slices=1)

var_form_swaprz = SwapRZ(qubitOp.num_qubits, depth=1, 
                   initial_state=HF_state)

var_form_ryrz = RYRZ(qubitOp.num_qubits, depth=1, 
                   initial_state=HF_state)

# var_form_ry = RY(qubitOp.num_qubits, depth=1, 
#                    initial_state=HF_state)

simultation_type = 'state_vector'# qasm, qasm_cloud, real_device, state_vector

if simultation_type == 'qasm':
    
    sv_mode = False
    qasm_simulator = Aer.get_backend('qasm_simulator') # qasm_simulator, ibmq_qasm_simulator

    # uncomment this and put your token from the IBMQ website.
    #IBMQ.save_account('Your IBMQ API-token', overwrite=True)
    
    IBMQ.load_accounts() # connect with your token to IBMQ
    provider = IBMQ.get_provider(hub='ibm-q', group='open', project='main')
    print(provider.backends()) # lists available backends
    
    # set the noise model
    device = provider.get_backend('ibmq_16_melbourne') # ibmq_16_melbourne 
    properties = device.properties()
    coupling_map = device.configuration().coupling_map
    noise_model = noise.device.basic_device_noise_model(properties)
    basis_gates = noise_model.basis_gates
    shots = 8192 # choose the maximal number of shots (8192 is max on the ibmq_qasm_simulator)
    num_qubits = qubitOp.num_qubits
    
    # run on qasm_simulator 
    quantum_instance = QuantumInstance(qasm_simulator, shots=shots, basis_gates=basis_gates, coupling_map=coupling_map, noise_model=noise_model)

    # all the details of noise model
    print(noise_model.to_dict())
    
elif simultation_type == 'qasm_cloud':
    
    sv_mode = False
    # uncomment this and put your token from the IBMQ website.
    #IBMQ.save_account('Your IBMQ API-token', overwrite=True)
    IBMQ.load_accounts() # connect with your token to IBMQ
    provider = IBMQ.get_provider(hub='ibm-q', group='open', project='main')
    print(provider.backends()) # lists available backends
    
    qasm_simulator_cloud = provider.get_backend('ibmq_qasm_simulator') # ibmq_16_melbourne , ibmq_qasm_simulator
    
    # set the qasm cloud simulator 
    shots = 8192
    quantum_instance = QuantumInstance(qasm_simulator_cloud, shots=shots)

elif simultation_type == 'real_device':
    
    sv_mode = False
    # uncomment this and put your token from the IBMQ website.
    #IBMQ.save_account('Your IBMQ API-token', overwrite=True)
    
    IBMQ.load_accounts() # connect with your token to IBMQ
    provider = IBMQ.get_provider(hub='ibm-q', group='open', project='main')
    print(provider.backends()) # lists available backends
    
    # set the noise model
    device = provider.get_backend('ibmq_16_melbourne') # ibmq_16_melbourne 
    
    # run on qasm_simulator 
    shots = 8192
    quantum_instance = QuantumInstance(device, shots=shots)
    
elif simultation_type == 'state_vector':

    sv_mode = True
    num_qubits = qubitOp.num_qubits

    sv_simulator = Aer.get_backend('statevector_simulator')
    quantum_instance = QuantumInstance(sv_simulator) 
else:
    raise Exception('Please select for simultation_type == qasm, qasm_cloud, real_device, state_vector')


vqe_uccsd = VQE(qubitOp, var_form_uccsd, optimizer)
vqe_ryrz = VQE(qubitOp, var_form_ryrz, optimizer)
vqe_swaprz = VQE(qubitOp, var_form_swaprz, optimizer)

for pauli in qubitOp.paulis:
    print('Circuit for:', pauli[1])

circuit_uccsd = vqe_uccsd.construct_circuit(statevector_mode=sv_mode, parameter=np.ones(var_form_uccsd._num_parameters))
for i in range(len(circuit_uccsd)):
    fig = circuit_uccsd[i].draw(output='mpl')
    show_figure(fig)
    fig.show()

circuit_swaprz = vqe_swaprz.construct_circuit(statevector_mode=sv_mode, parameter=np.ones(var_form_swaprz._num_parameters))
circuit_swaprz[1].draw(output='mpl')

circuit_ryrz = vqe_ryrz.construct_circuit(statevector_mode=sv_mode, parameter=np.ones(var_form_ryrz._num_parameters))
circuit_ryrz[1].draw(output='mpl')

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

evolution_time = 100 # how many time steps
num_time_slices = 4 # number of Trotter steps
qr = QuantumRegister(num_qubits,'q')
initial_circuit = HF_state.construct_circuit(mode='circuit')

evolution_circuit = qubitOp.evolve(state_in=initial_circuit, evo_time=evolution_time, evo_mode='circuit',num_time_slices=num_time_slices, expansion_mode='trotter')
evolution_circuit.draw(output='mpl')

max_time=evolution_time
statevector_of_t = np.zeros([max_time,2*num_qubits])
for evo_t in range(max_time):
    evolution_circuit = qubitOp.evolve(state_in=initial_circuit, evo_time=evo_t, evo_mode='circuit',num_time_slices=num_time_slices, expansion_mode='trotter')
    result = quantum_instance.execute(evolution_circuit)
    statevector_of_t[evo_t,:] = np.asarray(result.get_statevector(evolution_circuit))

time_min = 0
time_max = 4
for i in range(time_min, time_max):
    fig = plot_state_city(statevector_of_t[i,:])
    show_figure(fig)
    fig.show()

# First, we use classical eigendecomposition to get ground state energy (including nuclear repulsion energy) as reference.
qiskit_chemistry_dict = {
    'driver': {'name': 'HDF5'},
    'HDF5': {'hdf5_input': 'H2/H2_equilibrium_0.735_sto-3g.hdf5'},
    'operator': {'name':'hamiltonian', 
                 'qubit_mapping': 'parity', 
                 'two_qubit_reduction': True},
    'algorithm': {'name': 'ExactEigensolver'}
}
solver = QiskitChemistry()
result_classical = solver.run(qiskit_chemistry_dict)

# Second, we use variational quantum eigensolver (VQE)
qiskit_chemistry_dict['algorithm']['name'] = 'VQE'
qiskit_chemistry_dict['optimizer'] = {'name': 'SPSA', 'max_trials': 350}
qiskit_chemistry_dict['variational_form'] = {'name': 'RYRZ', 'depth': 3, 'entanglement':'full'}
backend = Aer.get_backend('statevector_simulator')
solver = QiskitChemistry()
result = solver.run(qiskit_chemistry_dict, backend=backend)

# We print out the results
print('Ground state energy (classical): {:.12f}'.format(result_classical['energy']))
print('Ground state energy (quantum)  : {:.12f}'.format(result['energy']))
print("====================================================")

# You can also print out other info in the field 'printable'
for line in result['printable']:
    print(line)

molecule='H2'
maxiter = 10 # increase to allow optimizer to converge, set low to test
mol_distances = np.arange(0.2, 3.1, 0.2)
mol_distances_vqe = np.arange(0.2, 3.1,1)

# Exact eigensolver
qiskit_chemistry_dict_ee = {
    'driver': {'name': 'HDF5'},
    'HDF5': {'hdf5_input': ''},
    'operator': {'name':'hamiltonian', 
                 'qubit_mapping': 'parity', 
                 'two_qubit_reduction': True},
    'algorithm': {'name': 'ExactEigensolver'}
}

# choose which backend want to use
# backend = Aer.get_backend('qasm_simulator')
backend = Aer.get_backend('statevector_simulator')

# VQE
qiskit_chemistry_dict_vqe = {
    'driver': {'name': 'HDF5'},
    'HDF5': {'hdf5_input': ''},
    'operator': {'name':'hamiltonian', 
                 'qubit_mapping': 'parity', 
                 'two_qubit_reduction': True},
    'algorithm': {'name': 'VQE'},
    'optimizer': {'name': 'L_BFGS_B', 'maxiter': maxiter}, # or SPSA, max_trials=200
    'variational_form': {'name': 'RY', 'depth': 1, 'entanglement':'full'}, # RYRZ, UCCSD (delete entanglement field)
    'backend': {'shots': 1024}
}

algos = ['ExactEigensolver', 'VQE']
energy = np.zeros((len(algos), len(mol_distances)))
energy_vqe = np.zeros((len(algos), len(mol_distances_vqe)))

# Run the algos
for j, algo in enumerate([qiskit_chemistry_dict_ee, qiskit_chemistry_dict_vqe]):
    algo_name = algo['algorithm']['name']
    print("Using {}".format(algo_name))
    if algo_name == 'VQE':
        for i, dis in enumerate(mol_distances_vqe):
            print("Processing atomic distance: {:1.1f} Angstrom".format(dis), end='\r')
            algo['HDF5']['hdf5_input'] = "{}/{:1.1f}_sto-3g.hdf5".format(molecule, dis)
            result = solver.run(algo, backend=backend if algo_name == 'VQE' else None)
            energy_vqe[j][i] = result['energy']
    else:
        for i, dis in enumerate(mol_distances):
            print("Processing atomic distance: {:1.1f} Angstrom".format(dis), end='\r')
            algo['HDF5']['hdf5_input'] = "{}/{:1.1f}_sto-3g.hdf5".format(molecule, dis)
            result = solver.run(algo, backend=backend if algo_name == 'VQE' else None)
            energy[j][i] = result['energy']
        

for i, algo in enumerate(algos):
    if algo == 'VQE':
        plt.plot(mol_distances_vqe, energy_vqe[i], label=algo, linestyle='', marker='x')
    else:
        plt.plot(mol_distances, energy[i], label=algo)
        
plt.xlabel('Atomic distance (Angstrom)')
plt.ylabel('Energy [Ha]')
plt.legend()
plt.show()

mol_distances_errors = []
energy_errors = []

for i, pts_vqe in enumerate(mol_distances_vqe):
    for j, pts in enumerate(mol_distances):
        # find the entries with same bond distances
        if pts_vqe == pts:
            mol_distances_errors.append(pts)
            energy_errors.append(abs(energy[0][j]-energy_vqe[1][i]))
            
# plot the errors
print('Errors |Exact - VQE|:\n',energy_errors)
plt.plot(mol_distances_errors, energy_errors, label='Errors(Exact-VQE)', linestyle='', marker='x')

# chemical accuracy of 1 mHa
plt.hlines(0.001, 0, 3, colors='k', label='Chemical Accuracy')
plt.legend()
plt.show()

import qiskit.tools.jupyter
%qiskit_version_table
%qiskit_copyright



