#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='my 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 runVQE(params, backend_id="qasm_simulator", optimizer_id="COBYLA", noise=None):
    
    molecule = Molecule(geometry=[['H',  [0., 0., 0.]],
                                  ['H', [0., 0., 0.75]]],
                                  charge=0, multiplicity=1)

    # Get Hamiltonian
    qubit_op = GetHamiltonians(molecule)

    # Get classical optimizer
    if (optimizer_id == "COBYLA"):
        optimizer = COBYLA()
    elif (optimizer_id == "SPSA"):
        optimizer = SPSA()
    elif (optimizer_id == "ADAM"):
        optimizer = ADAM()

    # Construct and display ansatz
    ansatz = constructAnsatz(params)
    print('Ansatz :', '\n', 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, noise_model=noise)

    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("Optimized VQE Energy ", result.eigenvalue.real)
    
    return values, stds

# Execute program with default parameters
eng_hist, error_hist = runVQE([0.0])

plt.plot([*range(len(eng_hist))], eng_hist, linestyle="", marker="o")
#plt.plot([*range(len(eng_hist))], eng_hist)
plt.show()

eng_hist_SPSA, error_hist_SPSA = runVQE([0.0],optimizer_id="SPSA")

plt.plot([*range(len(eng_hist))], eng_hist, linestyle="", marker="o", label="COBYLA")
plt.plot([*range(len(eng_hist_SPSA))], eng_hist_SPSA, linestyle="", marker="*", label="SPSA")
plt.legend()
plt.show()

eng_hist_ADAM, error_hist_ADAM = runVQE([0.0],optimizer_id="ADAM")

plt.plot([*range(len(eng_hist))], eng_hist, linestyle="", marker="o", label="COBYLA")
plt.plot([*range(len(eng_hist_SPSA))], eng_hist_SPSA, linestyle="", marker="*", label="SPSA")
plt.plot([*range(len(eng_hist_ADAM))], eng_hist_ADAM, linestyle="", marker="x", label="ADAM")
plt.legend()
plt.show()

# Noise
from qiskit.providers.aer.noise import NoiseModel
from qiskit.providers.aer.noise import depolarizing_error

# default noise model, can be overridden using set_noise_model
noise = NoiseModel()
# Add depolarizing error to all single qubit gates with error rate 0.5%
one_qb_error = 0.005
noise.add_all_qubit_quantum_error(depolarizing_error(one_qb_error, 1), ['u1', 'u2', 'u3'])

# Add depolarizing error to all two qubit gates with error rate 5.0%
two_qb_error = 0.05
noise.add_all_qubit_quantum_error(depolarizing_error(two_qb_error, 2), ['cx'])

eng_hist, error_hist = runVQE([0.0],optimizer_id="COBYLA", noise=noise)
eng_hist_SPSA, error_hist_SPSA = runVQE([0.0],optimizer_id="SPSA", noise=noise)
eng_hist_ADAM, error_hist_ADAM = runVQE([0.0],optimizer_id="ADAM", noise=noise)

plt.plot([*range(len(eng_hist))], eng_hist, linestyle="", marker="o", label="COBYLA")
plt.plot([*range(len(eng_hist_SPSA))], eng_hist_SPSA, linestyle="", marker="*", label="SPSA")
plt.plot([*range(len(eng_hist_ADAM))], eng_hist_ADAM, linestyle="", marker="x", label="ADAM")
plt.legend()
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()

eng_hist, error_hist = runVQE([0.0],optimizer_id="COBYLA", backend_id = "ionq_qpu")
eng_hist_SPSA, error_hist_SPSA = runVQE([0.0],optimizer_id="SPSA", backend_id = "ionq_qpu")
eng_hist_ADAM, error_hist_ADAM = runVQE([0.0],optimizer_id="ADAM", backend_id = "ionq_qpu")

plt.plot([*range(len(eng_hist))], eng_hist, linestyle="", marker="o", label="COBYLA")
plt.plot([*range(len(eng_hist_SPSA))], eng_hist_SPSA, linestyle="", marker="*", label="SPSA")
plt.plot([*range(len(eng_hist_ADAM))], eng_hist_ADAM, linestyle="", marker="x", label="ADAM")
plt.legend()
plt.show()