from qiskit import BasicAer, Aer, IBMQ
from qiskit.aqua import QuantumInstance, aqua_globals
from qiskit.aqua.algorithms import VQE, NumPyEigensolver
from qiskit.aqua.components.initial_states import Zero
from qiskit.aqua.components.optimizers import COBYLA, L_BFGS_B, SLSQP, SPSA
from qiskit.aqua.components.variational_forms import RY, RYRZ, SwapRZ
from qiskit.aqua.operators import WeightedPauliOperator, Z2Symmetries
from qiskit.chemistry import FermionicOperator
from qiskit.chemistry.drivers import PySCFDriver, UnitsType
from qiskit.chemistry.components.variational_forms import UCCSD
from qiskit.chemistry.components.initial_states import HartreeFock
from qiskit.circuit.library import EfficientSU2


from qiskit.providers.aer import QasmSimulator
from qiskit.providers.aer.noise import NoiseModel
from qiskit.providers.aer.noise.errors import QuantumError, ReadoutError
from qiskit.providers.aer.noise.errors import pauli_error
from qiskit.providers.aer.noise.errors import depolarizing_error
from qiskit.providers.aer.noise.errors import thermal_relaxation_error

from qiskit.providers.aer import noise
provider = IBMQ.load_account()

import numpy as np
import matplotlib.pyplot as plt
from functools import partial

import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)

# Classically solve for the lowest eigenvalue
# This is used just to compare how well you VQE approximation is performing
def exact_solver(qubitOp, shift):
    ee = NumPyEigensolver(qubitOp)
    result = ee.run()
    ref = result['eigenvalues']
    ref+= shift
    print('Reference value: {}'.format(ref))
    return ref

# Define your function for computing the qubit operations of LiH
def compute_LiH_qubitOp(map_type, inter_dist, basis='sto3g'):
    
    # Specify details of our molecule
    driver = PySCFDriver(atom='Li .0 .0 .0; H .0 .0 ' + str(inter_dist), unit=UnitsType.ANGSTROM, charge=0, spin=0, basis=basis)

    # Compute relevant 1 and 2 body integrals.
    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))

    # Please be aware that the idx here with respective to original idx
    freeze_list = [0]
    remove_list = [-3, -2] # negative number denotes the reverse order
    
    # Prepare full idx of freeze_list and remove_list
    # Convert all negative idx to positive
    remove_list = [x % molecule.num_orbitals for x in remove_list]
    freeze_list = [x % molecule.num_orbitals for x in freeze_list]
    
    # Update the idx in remove_list of the idx after frozen, since the idx of orbitals are changed after freezing
    remove_list = [x - len(freeze_list) for x in remove_list]
    remove_list += [x + molecule.num_orbitals - len(freeze_list)  for x in remove_list]
    freeze_list += [x + molecule.num_orbitals for x in freeze_list]

    # Prepare fermionic hamiltonian with orbital freezing and eliminating, and then map to qubit hamiltonian
    # and if PARITY mapping is selected, reduction qubits
    energy_shift = 0.0
    qubit_reduction = True if map_type == 'parity' else False

    ferOp = FermionicOperator(h1=h1, h2=h2)
    if len(freeze_list) > 0:
        ferOp, energy_shift = ferOp.fermion_mode_freezing(freeze_list)
        num_spin_orbitals -= len(freeze_list)
        num_particles -= len(freeze_list)
    if len(remove_list) > 0:
        ferOp = ferOp.fermion_mode_elimination(remove_list)
        num_spin_orbitals -= len(remove_list)

    qubitOp = ferOp.mapping(map_type=map_type)
    qubitOp = Z2Symmetries.two_qubit_reduction(qubitOp, num_particles) if qubit_reduction else qubitOp
    qubitOp.chop(10**-10)
    
    #Overall energy shift:
    shift = energy_shift + nuclear_repulsion_energy 

    return qubitOp, num_spin_orbitals, num_particles, qubit_reduction, shift

map_type = 'parity'
inter_dist = 1.6
qubitOp, num_spin_orbitals, num_particles, qubit_reduction, shift = compute_LiH_qubitOp(map_type , inter_dist)
# Classically solve for the exact solution and use that as your reference value
ref = exact_solver(qubitOp , shift)

# Specify your initial state
init_state = HartreeFock(num_spin_orbitals,num_particles, qubit_mapping=map_type) 

# Select a state preparation ansatz
# Equivalently, choose a parameterization for our trial wave function.
var_form = UCCSD(num_orbitals=num_spin_orbitals, num_particles=num_particles, qubit_mapping=map_type)

# Choose where to run/simulate our circuit
quantum_instance = Aer.get_backend('statevector_simulator')

# Choose the classical optimizer
optimizer =  SPSA(max_trials = 1000)

# Run your VQE instance
vqe = VQE(qubitOp, var_form, optimizer)
vqe_results = vqe.run(quantum_instance)

# Calculating ground ststae energy from vqe results:
ground_state_energy = vqe_results['eigenvalue'] + shift
print(ground_state_energy)

distances = np.arange(0.2 , 5 , 0.1)
energies = []
for interdist in distances:
    qubitOp, num_spin_orbitals, num_particles, qubit_reduction, shift = compute_LiH_qubitOp(map_type , interdist)
    # Classically solve for the exact solution and use that as your reference value
    ref = exact_solver(qubitOp , shift)

    # Specify your initial state
    init_state = HartreeFock(num_spin_orbitals,num_particles, qubit_mapping=map_type) 

    # Select a state preparation ansatz
    # Equivalently, choose a parameterization for our trial wave function.
    var_form = UCCSD(num_orbitals=num_spin_orbitals, num_particles=num_particles, qubit_mapping=map_type)

    # Choose where to run/simulate our circuit
    quantum_instance = Aer.get_backend('statevector_simulator')

    # Choose the classical optimizer
    optimizer =  SPSA(max_trials = 1000)

    # Run your VQE instance
    vqe = VQE(qubitOp, var_form, optimizer)
    vqe_results = vqe.run(quantum_instance)

    # Calculating ground state energy from vqe results:
    ground_state_energy = vqe_results['eigenvalue'] + shift
    energies.append(ground_state_energy)

plt.plot(distances , energies)
plt.title('Bond energy v. interatomic distance')
plt.xlabel('Interatomic distance/ Angstroms')
plt.ylabel('Energy / Hartrees')

def store_intermediate_result(eval_count, parameters, mean, std):
            counts.append(eval_count)
            values.append(mean)
            params.append(parameters)
            deviation.append(std)

map_type = 'parity'
inter_dist = 1.6
# Dictionary of optimizers:
opt_dict = {'SPSA' , 'SLSQP' , 'COBYLA' , 'L_BFGS_B'}

for opt in opt_dict:
    print('Testing', str(opt) , 'optimizer')
    qubitOp, num_spin_orbitals, num_particles, qubit_reduction, shift = compute_LiH_qubitOp(map_type , interdist)
    # Classically solve for the exact solution and use that as your reference value
    ref = exact_solver(qubitOp , shift)

    # Specify your initial state
    init_state = HartreeFock(num_spin_orbitals,num_particles, qubit_mapping=map_type) 

    # Select a state preparation ansatz
    # Equivalently, choose a parameterization for our trial wave function.
    var_form = UCCSD(num_orbitals=num_spin_orbitals, num_particles=num_particles, qubit_mapping=map_type)

    # Choose where to run/simulate our circuit
    quantum_instance = Aer.get_backend('statevector_simulator')

    # Choose the classical optimizer
    if opt == 'SPSA':
        optimizer = SPSA(max_trials = 500)
    elif opt == 'SLSQP':
        optimizer = SLSQP(maxiter = 1000)
    elif opt == 'L_BFGS_B':
        optimizer = L_BFGS_B(maxfun = 1000 ,  maxiter = 1000)
    elif opt == 'COBYLA':
        optimizer = COBYLA(maxiter = 1000)
    
    counts =[]
    values =[]
    params =[]
    deviation =[]

    # Run your VQE instance
    vqe = VQE(qubitOp, var_form, optimizer , callback = store_intermediate_result)
    vqe_results = vqe.run(quantum_instance)
    
    #Printing error in final value:
    ground_state_energy = vqe_results['eigenvalue'] + shift
    energy_error_ground = np.abs(np.real(ref) - ground_state_energy)
    print('Error:', str(energy_error_ground))
    # Calculating energy error
    vqe_energies = np.real(values) + shift
    energy_error = np.abs(np.real(ref) - vqe_energies)
    plt.plot(counts , energy_error , label=str(opt))
plt.legend()
plt.xlabel('Counts')
plt.ylabel('Energy Error/ Hartree')
plt.title('Energy Convergence of VQE: UCCSD Ansatz')

map_type = 'parity'
inter_dist = 1.6
# Dictionary of optimizers:
opt_dict = {'SPSA' , 'SLSQP' , 'COBYLA' , 'L_BFGS_B'}

for opt in opt_dict:
    print('Testing', str(opt) , 'optimizer')
    qubitOp, num_spin_orbitals, num_particles, qubit_reduction, shift = compute_LiH_qubitOp(map_type , interdist)
    # Classically solve for the exact solution and use that as your reference value
    ref = exact_solver(qubitOp , shift)

    # Specify your initial state
    init_state = HartreeFock(num_spin_orbitals,num_particles, qubit_mapping=map_type) 

    # Select a state preparation ansatz
    # Equivalently, choose a parameterization for our trial wave function.
    var_form = EfficientSU2(qubitOp.num_qubits , entanglement='full')

    # Choose where to run/simulate our circuit
    quantum_instance = Aer.get_backend('statevector_simulator')

    # Choose the classical optimizer
    if opt == 'SPSA':
        optimizer = SPSA(max_trials = 1000)
    elif opt == 'SLSQP':
        optimizer = SLSQP(maxiter = 1000)
    elif opt == 'L_BFGS_B':
        optimizer = L_BFGS_B(maxfun = 1000 ,  maxiter = 1000)
    elif opt == 'COBYLA':
        optimizer = COBYLA(maxiter = 1000)
    
    counts =[]
    values =[]
    params =[]
    deviation =[]

    # Run your VQE instance
    vqe = VQE(qubitOp, var_form, optimizer , callback = store_intermediate_result)
    vqe_results = vqe.run(quantum_instance)

    #Printing error in final value:
    ground_state_energy = vqe_results['eigenvalue'] + shift
    energy_error_ground = np.abs(np.real(ref) - ground_state_energy)
    print('Error:', str(energy_error_ground))
    
    # Calculating energy error
    vqe_energies = np.real(values) + shift
    energy_error = np.abs(np.real(ref) - vqe_energies)
    plt.plot(counts , energy_error , label=str(opt))
plt.legend()
plt.xlabel('Counts')
plt.ylabel('Energy Error/ Hartree')
plt.title('Energy Convergence of VQE: RyRz Ansatz')



