#import Aer here, before calling qiskit_ionq_provider
from qiskit import Aer

from qiskit_ionq import IonQProvider 

#Call provider and set token value
ionq_provider = IonQProvider(token='your token')

ionq_provider.backends()

import numpy as np

from qiskit import execute
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit.utils import QuantumInstance
from qiskit.algorithms import VQE
from qiskit.algorithms.optimizers import COBYLA, ADAM, L_BFGS_B, SLSQP, SPSA
from qiskit import Aer
from qiskit.circuit import Parameter
from qiskit.circuit.parametervector import ParameterVectorElement

from qiskit_nature.drivers import PySCFDriver, UnitsType, Molecule
from qiskit_nature.circuit.library import HartreeFock
from qiskit_nature.problems.second_quantization.electronic import ElectronicStructureProblem
from qiskit_nature.mappers.second_quantization import JordanWignerMapper
from qiskit_nature.converters.second_quantization import QubitConverter

from qiskit.algorithms import NumPyMinimumEigensolver
from qiskit_nature.algorithms.ground_state_solvers import GroundStateEigensolver

from qiskit_nature.runtime import VQEProgram

from matplotlib import pyplot as plt
%matplotlib inline 

class UCCAnsatz():

    def __init__(self, params, num_particles, num_spin_orbitals):

        # number of qubits 
        self.num_of_qubits = num_spin_orbitals

        # number of parameters
        self.num_params = len(params)

        # parameters
        self.params = []
        for index in range(self.num_params):
            p = Parameter("t"+str(index))
            self.params.append(ParameterVectorElement(p, index))

    def UCC1(self):

        qc = QuantumCircuit(4)

        # basis rotation
        qc.rx(np.pi / 2, 0)
        qc.h(1)
        qc.h(2)
        qc.h(3)

        # parameter theta_0
        qc.cx(0, 1)
        qc.cx(1, 2)
        qc.cx(2, 3)
        qc.rz(self.params[0], 3)
        qc.cx(2, 3)
        qc.cx(1, 2)
        qc.cx(0, 1)

        # basis rotation
        qc.rx(-np.pi / 2, 0)
        qc.h(1)
        qc.h(2)
        qc.h(3)


        return qc

################## Hamiltonian Definition #######################################

def GetHamiltonians(mol):

    # construct the driver
    driver = PySCFDriver(molecule=mol, unit=UnitsType.ANGSTROM, basis='sto6g')

    # the electronic structure problem
    problem = ElectronicStructureProblem(driver)

    # get quantum molecule
    q_molecule = driver.run()

    # classical eigenstate
    np_solver = NumPyMinimumEigensolver()
    np_groundstate_solver = GroundStateEigensolver(QubitConverter(JordanWignerMapper()), np_solver)

    np_result = np_groundstate_solver.solve(problem)
    print(f"Classical results is {np_result.eigenenergies}")

    # generate the second-quantized operators
    second_q_ops = problem.second_q_ops()

    # construct a qubit converter
    qubit_converter = QubitConverter(JordanWignerMapper())

    # qubit Operations 
    qubit_op = qubit_converter.convert(second_q_ops[0])

    # return the qubit operations
    return qubit_op


def constructAnsatz(params):

    # initialize the HF state
    hf_state = HartreeFock(4, [1, 1], QubitConverter(JordanWignerMapper()))

    # VQE circuit
    ansatz = UCCAnsatz(params, 2, 4).UCC1()

    # add initial state
    ansatz.compose(hf_state, front=True, inplace=True)

    # return the circuit
    return ansatz


def constructCircuit(params, mol, backend_id):

    # Hamiltonian
    qubit_op = GetHamiltonians(mol)

    # Optimizer
    optimizer = COBYLA()

    # ansatz
    ansatz = constructAnsatz(params)
    print(ansatz)

    # initial parameters 
    init_params = {}
    for index in range(len(params)):
        p = Parameter("t"+str(index))
        init_params[ParameterVectorElement(p, index)] = 0.0

    # backend 
    if (backend_id == "qasm_simulator"):
        backend = Aer.get_backend("qasm_simulator")
    elif (backend_id == "ionq_qpu"):
        backend = ionq_provider.get_backend("ionq_qpu")

    def job_callback(job_id, job_status, queue_position, job) -> None:
        '''
        Printing logs for debugging.
        The function is the parameter job_callback of QuantumInstance
        :param job_id:
        :param job_status:
        :param queue_position:
        :param job:
        :return:
        '''
        print(f"[API JOB-CALLBACK] Job Id:         {job_id}")
        print(f"[API JOB-CALLBACK] Job status:     {job_status}")
        print(f"[API JOB-CALLBACK] Queue position: {queue_position}")
        print(f"[API JOB-CALLBACK] Job:            {job}")

    ionq_quantum_instance = QuantumInstance(backend=backend, job_callback=job_callback)

    counts = []
    values = []
    stds   = []
    def store_intermediate_result(eval_count, parameters, mean, std):
        counts.append(eval_count)
        values.append(mean)
        stds.append(std)

    # VQE algorithm
    algorithm = VQE(ansatz,
                    optimizer=optimizer,
                    initial_point = np.array([0.0]),
                    callback=store_intermediate_result,
                    quantum_instance=ionq_quantum_instance)

    result = algorithm.compute_minimum_eigenvalue(qubit_op)
    print(result)
    print("Optimized VQE Energy ", result.eigenvalue.real)
    print(values)


# Execute program with default parameters
def run(backend_id="qasm_simulator"):

    molecule = Molecule(geometry=[['H',  [0., 0., 0.]],
                                  ['H', [0., 0., 0.75]]],
                                  charge=0, multiplicity=1)
    opt_values = constructCircuit([0.0], molecule, backend_id)

run()

opt_process = [-1.8308337397334875, -1.6045286426151377, -1.3176850880519417, -1.642972785671419, -1.853378862063689, -1.806712185391583, -1.8281518913000796, -1.853965175510973, -1.8546227409607912, -1.8512623145326073, -1.8499200431393457, -1.8490846844773303, -1.8508890835046492, -1.8474757232657733, -1.8492536553877457, -1.8544537700503756, -1.8461599187778577, -1.8476358718744068, -1.849466737807071]
plt.plot(opt_process, linestyle="", marker="o")
plt.show()

run("ionq_qpu")

opt_process_exp = [-1.8112359199652812, -1.4416498480789817, -1.2335035452225904, -1.4989706208229303, -1.6937521890989031, -1.6557546667316045, -1.6898027909394815, -1.740016937435431, -1.7592133266130963, -1.755505060521276, -1.7199835843097526, -1.7476529655074913, -1.7142357090197362, -1.7366022956905651, -1.7307574124517149, -1.7437590702130892, -1.721561556686939]
plt.plot(opt_process, linestyle="", marker="o", color="b", label="simulator")
plt.plot(opt_process_exp, linestyle="", marker="o", color="r", label="ionq_qpu")
plt.legend()
plt.show()

opt_process_exp = [-1.8112359199652812, -1.4416498480789817, -1.2335035452225904, -1.4989706208229303, -1.6937521890989031, -1.6557546667316045, -1.6898027909394815, -1.740016937435431, -1.7592133266130963, -1.755505060521276, -1.7199835843097526, -1.7476529655074913, -1.7142357090197362, -1.7366022956905651, -1.7307574124517149, -1.7437590702130892, -1.721561556686939]
opt_process = [-1.8308337397334875, -1.6045286426151377, -1.3176850880519417, -1.642972785671419, -1.853378862063689, -1.806712185391583, -1.8281518913000796, -1.853965175510973, -1.8546227409607912, -1.8512623145326073, -1.8499200431393457, -1.8490846844773303, -1.8508890835046492, -1.8474757232657733, -1.8492536553877457, -1.8544537700503756, -1.8461599187778577, -1.8476358718744068, -1.849466737807071]
plt.plot(opt_process, linestyle="", marker="o", color="b", label="simulator")
plt.plot(opt_process_exp, linestyle="", marker="o", color="r", label="ionq_qpu")
plt.xlabel("Number of Iteration")
plt.ylabel("Energy (Hartree)")
plt.legend()
plt.show()