from qiskit import IBMQ, Aer, transpile, schedule as build_schedule
from qiskit.transpiler import PassManager
from qiskit.transpiler.passes.calibration import RZXCalibrationBuilderNoEcho
from qiskit_nature.drivers import Molecule
from qiskit_nature.drivers.second_quantization import ElectronicStructureDriverType, ElectronicStructureMoleculeDriver
from qiskit_nature.problems.second_quantization import ElectronicStructureProblem
from qiskit_nature.converters.second_quantization import QubitConverter
from qiskit_nature.transformers.second_quantization.electronic import FreezeCoreTransformer
from qiskit_nature.operators.second_quantization import FermionicOp
from qiskit_nature.mappers.second_quantization import ParityMapper
from qiskit_nature.algorithms import GroundStateEigensolver
from qiskit.algorithms import NumPyMinimumEigensolver, VQE
from qiskit.algorithms.optimizers import SPSA
from qiskit.circuit import QuantumCircuit, ParameterVector
from qiskit.utils import QuantumInstance
import matplotlib.pyplot as plt
import numpy as np

IBMQ.load_account()
provider = IBMQ.get_provider()
lagos = provider.get_backend('ibm_lagos')

def HEA_naive(num_q, depth):
    circuit = QuantumCircuit(num_q)
    params = ParameterVector("theta", length=num_q * (3 * depth + 2))
    counter = 0
    for q in range(num_q):
        circuit.rx(params[counter], q)
        counter += 1 
        circuit.rz(params[counter], q)
        counter += 1
    for d in range(depth):
        for q in range(num_q - 1):
            circuit.cx(q, q + 1)
        for q in range(num_q):
            circuit.rz(params[counter], q)
            counter += 1
            circuit.rx(params[counter], q)
            counter += 1 
            circuit.rz(params[counter], q)
            counter += 1
    return circuit, params

def HEA_aware(num_q, depth, hardware):
    circuit = QuantumCircuit(num_q)
    params = ParameterVector("theta", length=num_q * (3 * depth + 2))
    counter = 0
    for q in range(num_q):
        circuit.rx(params[counter], q)
        counter += 1 
        circuit.rz(params[counter], q)
        counter += 1
    for d in range(depth):
        for q in range(num_q - 1):
            gate = QuantumCircuit(num_q)
            gate.rzx(np.pi/2, q, q + 1)
            pass_ = RZXCalibrationBuilderNoEcho(hardware)
            qc_cr = PassManager(pass_).run(gate)
            circuit.compose(qc_cr, inplace=True)
        for q in range(num_q):
            circuit.rz(params[counter], q)
            counter += 1
            circuit.rx(params[counter], q)
            counter += 1 
            circuit.rz(params[counter], q)
            counter += 1
    return circuit, params

depth = 1
qubits = 2
back = lagos
circuit, _ = HEA_naive(qubits, depth)
t = transpile(circuit, back)
schedule = build_schedule(t, back)
schedule.draw()

circuit, _ = HEA_aware(qubits, depth, back)
t = transpile(circuit, back)
schedule = build_schedule(t, back)
schedule.draw()

qubit_converter = QubitConverter(ParityMapper(), two_qubit_reduction=True)
numpy_solver = NumPyMinimumEigensolver()

spsa = SPSA(100)
qi = QuantumInstance(Aer.get_backend('aer_simulator'))
vqe_circuit = VQE(ansatz=circuit, quantum_instance=qi, optimizer=spsa)

# H2 simulation

total_dist = 4
dist = 0.1
incr_early = 0.1
incr_late = 0.3

real_energies = []
vqe_energies = []
dists = []
while dist < total_dist:
    print(dist, total_dist)
    molecule = Molecule(geometry=[['H', [0., 0., 0.]], ['H', [0., 0., dist]]])
    driver = ElectronicStructureMoleculeDriver(molecule, basis='sto3g', \
                                               driver_type=ElectronicStructureDriverType.PYSCF)

    es_problem = ElectronicStructureProblem(driver)
    second_q_ops = es_problem.second_q_ops()    
    
    if dist == 0.1:
        FermionicOp.set_truncation(0)
        print(second_q_ops[0])
        FermionicOp.set_truncation(1)
        print(qubit_converter.convert(second_q_ops[0], num_particles=es_problem.num_particles))
        print(es_problem.grouped_property)    

    calc = GroundStateEigensolver(qubit_converter, numpy_solver)
    res = calc.solve(es_problem)
    real_energies.append(np.real(res.total_energies[0]))

    calc = GroundStateEigensolver(qubit_converter, vqe_circuit)
    res = calc.solve(es_problem)
    vqe_energies.append(np.real(res.total_energies[0]))
    
    
    dists.append(dist)
    if dist > total_dist / 2:
        dist += incr_late
    else:
        dist += incr_early

plt.plot(dists, real_energies, label='Real', color='red')
plt.scatter(dists, vqe_energies, label='VQE', color='black')
plt.title("H2")
plt.ylim(-1.2, 0.4)
plt.xlim(0, 4)
plt.xlabel('Angstroms')
plt.ylabel('Hartree')
plt.legend()
plt.show()

depth = 2
qubits = 4
circuit, _ = HEA_naive(qubits, depth)
spsa = SPSA(100)
qi = QuantumInstance(Aer.get_backend('aer_simulator'))
vqe_circuit = VQE(ansatz=circuit, quantum_instance=qi, optimizer=spsa)
#print(circuit)

total_dist = 5
dist = 0.3
incr_early = 0.1
incr_late = 0.3

real_energies = []
vqe_energies = []
dists = []
while dist < total_dist:
    print(dist, total_dist)
    molecule = Molecule(geometry=[['Li', [0., 0., 0.]], ['H', [0., 0., dist]]])
    driver = ElectronicStructureMoleculeDriver(molecule, basis='sto3g', \
                                               driver_type=ElectronicStructureDriverType.PYSCF)
    
    transformer = FreezeCoreTransformer(freeze_core=True, remove_orbitals=[3, 4])
    es_problem = ElectronicStructureProblem(driver, transformers=[transformer])
    second_q_ops = es_problem.second_q_ops()    

    if dist == 0.3:
        print(second_q_ops[0])
        print(qubit_converter.convert(second_q_ops[0], num_particles=es_problem.num_particles))
        print(es_problem.grouped_property)    
        print(es_problem.grouped_property_transformed)

    calc = GroundStateEigensolver(qubit_converter, numpy_solver)
    res = calc.solve(es_problem)
    real_energies.append(np.real(res.total_energies[0]))

    calc = GroundStateEigensolver(qubit_converter, vqe_circuit) #VQEClient()
    res = calc.solve(es_problem)
    vqe_energies.append(np.real(res.total_energies[0]))
    
    dists.append(dist)
    if dist > total_dist / 2:
        dist += incr_late
    else:
        dist += incr_early

plt.plot(dists, real_energies, label='Real', color='green')
plt.scatter(dists, vqe_energies, label='VQE', color='black')
plt.title("LiH")
plt.ylim(-8, -6.6)
plt.xlim(0.3, 5)
plt.xlabel('Angstroms')
plt.ylabel('Hartree')
plt.legend()
plt.show()

depth = 2
qubits = 6
circuit, _ = HEA_naive(qubits, depth)
spsa = SPSA(100)
qi = QuantumInstance(Aer.get_backend('aer_simulator'))
vqe_circuit = VQE(ansatz=circuit, quantum_instance=qi, optimizer=spsa)
#print(circuit)

total_dist = 5
dist = 0.4
incr_early = 0.2
incr_late = 0.4

real_energies = []
vqe_energies = []
dists = []
while dist < total_dist:
    print(dist, total_dist)
    molecule = Molecule(geometry=[['H', [0., 0., -dist]], ['Be', [0., 0., 0.]], ['H', [0., 0., dist]]])
    driver = ElectronicStructureMoleculeDriver(molecule, basis='sto3g', \
                                               driver_type=ElectronicStructureDriverType.PYSCF)

    transformer = FreezeCoreTransformer(freeze_core=True, remove_orbitals=[4, 5])
    es_problem = ElectronicStructureProblem(driver, transformers=[transformer])
    second_q_ops = es_problem.second_q_ops()    

    if dist == 0.4:
        print(second_q_ops[0])
        print(qubit_converter.convert(second_q_ops[0], num_particles=es_problem.num_particles))
        print(es_problem.grouped_property)    
        print(es_problem.grouped_property_transformed)

    calc = GroundStateEigensolver(qubit_converter, numpy_solver)
    res = calc.solve(es_problem)
    real_energies.append(np.real(res.total_energies[0]))

    calc = GroundStateEigensolver(qubit_converter, vqe_circuit)
    res = calc.solve(es_problem)
    vqe_energies.append(np.real(res.total_energies[0]))
    
    dists.append(dist)
    if dist > total_dist / 2:
        dist += incr_late
    else:
        dist += incr_early

plt.plot(dists, real_energies, label='Real', color='blue')
plt.scatter(dists, vqe_energies, label='VQE', color='black')
plt.title("BeH2")
plt.ylim(-15.7, -12.0)
plt.xlim(0.48, 5)
plt.xlabel('Angstroms')
plt.ylabel('Hartree')
plt.legend()
plt.show()

from qiskit.tools.jupyter import *

%qiskit_version_table