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

from qiskit_ionq import IonQProvider 

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

provider.backends()

from qiskit import *
#import qiskit.aqua as aqua
#from qiskit.quantum_info import Pauli
#from qiskit.aqua.operators.primitive_ops import PauliOp
from qiskit.circuit.library import PhaseEstimation
from qiskit import QuantumCircuit
from qiskit.visualization import plot_histogram
%matplotlib inline

import pyscf
from pyscf import gto, scf, cc

from matplotlib import pyplot as plt 
import numpy as np

from qiskit_ionq import IonQProvider
provider = IonQProvider("sUkOPlNLf6vamkFMB3SuHR2qYR5cLaMn")

from qiskit.providers.jobstatus import JobStatus
import time as timer

g = {"I": -0.4804, "Z0": 0.3435, "Z1": -0.4347, "Z0Z1": 0.5716, "X0X1": 0.0910, "Y0Y1": 0.0910}

"""
  This function implements the controlled exp(-iHt)
  Inputs:
    phase: phase to correct
    weights: Pauli weights
    time: Time
    order_index: Pauli operator index
    iter_num: iteration
    trotter_step: trotter steps

  Output:
    qc: Quantum circuit
"""
def buildControlledHam(phase, time, iter_num):

    # initialize the circuit
    qc = QuantumCircuit(3, 3)

    # Hardmard on ancilla
    qc.h(0)

    # initialize to |01>
    qc.x(2)
    
    # Z0 
    qc.crz(g["Z0"] * time * 2 * 2**iter_num, 0, 2)

    # Y0Y1
    qc.rz(np.pi/2, 1)
    qc.rz(np.pi/2, 2)
    qc.h(1)
    qc.h(2)
    qc.cx(1, 2)
    qc.crz(g["Y0Y1"] * time * 2 * 2**iter_num, 0, 2)
    qc.cx(1, 2)
    qc.h(1)
    qc.h(2)
    qc.rz(-np.pi/2, 1)
    qc.rz(-np.pi/2, 2)
            
    # Z1
    qc.crz(g["Z1"] * time * 2 * 2**iter_num, 0, 1)

    # X0X1
    qc.h(1)
    qc.h(2)
    qc.cx(1, 2)
    qc.crz(g["X0X1"] * time * 2 * 2**iter_num, 0, 2)
    qc.cx(1, 2)
    qc.h(1)
    qc.h(2)

    # phase correction
    qc.rz(phase, 0)

    # inverse QFT
    qc.h(0)

    qc.measure([0, 1, 2],[0, 1, 2])

    return qc


"""
  This function implements the controlled exp(-iHt)
  Inputs:
    weights: Pauli weights
    time: Time
    order_index: Pauli operator index
    tot_num_iter: number of iterations
    trotter_step: trotter steps

  Output:
    bits: list of measured bits
"""
def IPEA(time, tot_num_iter, backend_id):

    # get backend
    if (backend_id == "qasm_simulator"):
        backend = Aer.get_backend(backend_id)
    else:
        backend = provider.get_backend("ionq_qpu")

    # bits
    bits = []

    # phase correction
    phase = 0.0

    # loop over iterations 
    for i in range(tot_num_iter-1, -1, -1):

        # construct the circuit
        qc = buildControlledHam(phase, time, i)
        # run the circuit
        job = execute(qc, backend)
        
        if (backend_id == "ionq_qpu"):

            # Check if job is done
            while job.status() is not JobStatus.DONE:
                print("Job status is", job.status() )
                timer.sleep(60)

            # once we break out of that while loop, we know our job is finished
            print("Job status is", job.status() )

        # get result
        result = job.result()

        # get current bit
        this_bit = int(max(result.get_counts(), key=result.get_counts().get)[-1])
        
        bits.append(this_bit)

        # update phase correction
        phase /= 2
        phase -= (2 * np.pi * this_bit / 4.0)

    return bits

"""
  This function that computes eigenvalues from bit list
  Inputs:
    bits: list of measured bits
    time: Time

  Output:
    eig: eigenvalue
"""
def eig_from_bits(bits, time):

    eig = 0.

    m = len(bits)

    # loop over all bits
    for k in range(len(bits)):
        eig += bits[k] / (2**(m-k))

    eig *= -2*np.pi
    eig /= time

    return eig

"""
  This function that performs classical post-processing
  Inputs:
    eig:     eigenvalue
    weights: Pauli operator weights
    R:       Bond distance

  Output:
    energy: total energy
"""
def post_process(eig, weights, R):

    # initialize energy
    energy = eig

    # Z0Z1 contribution
    energy -= weights["Z0Z1"]

    # I contribution
    energy += weights["I"]

    # Nuclear Repulsion ( assume R is in Angstrom )
    energy += 1.0/ (R * 1.88973)

    # return energy 
    return energy


# backend
backend_id = "qasm_simulator"

# Bond Length
R = 0.75

# time 
t = 0.74

# number of iteration
max_num_iter = 8

# get exact energy 
# This function initialises a molecule

mol = pyscf.M(
    atom = 'H 0 0 0; H 0 0 0.75',  # in Angstrom
    basis = 'sto-6g',
    symmetry = False,
)
myhf = mol.RHF().run()

# create an FCI solver based on the SCF object
#
cisolver = pyscf.fci.FCI(myhf)
print('E(FCI) = %.12f' % cisolver.kernel()[0])


error_list = []

# perform IPEA
for i in range(1, max_num_iter + 1):
    bits = IPEA(t, i, backend_id)

    # re-construct phase
    eig = eig_from_bits(bits, t)
    print(eig, 'eig')

    # re-construct energy
    eng = post_process(eig, g, R)
                
    print("Total Energy is %.7f for R = %.2f, t = %.3f with %d iterations" % (eng, R, t, i) )

backend_id = "qpu"
error_list = []

# perform IPEA
for i in range(1, max_num_iter + 1):
    bits = IPEA(t, i, backend_id)

    # re-construct phase
    eig = eig_from_bits(bits, t)

    # re-construct energy
    eng = post_process(eig, g, R)
                
    print("Total Energy is %.7f for R = %.2f, t = %.3f with %d iterations" % (eng, R, t, i) )