import warnings

warnings.filterwarnings("ignore")

from matplotlib import pyplot as plt
import numpy as np

from qiskit.circuit.library import TwoLocal
from qiskit_nature.drivers import UnitsType, Molecule
from qiskit_nature.drivers.second_quantization import PySCFDriver
from qiskit_nature.problems.second_quantization import ElectronicStructureProblem
from qiskit_nature.mappers.second_quantization import JordanWignerMapper
from qiskit_nature.converters.second_quantization import QubitConverter
from qiskit import Aer

import sys

sys.path.append("../../")
from entanglement_forging import EntanglementForgedVQE
from entanglement_forging import EntanglementForgedGroundStateSolver
from entanglement_forging import EntanglementForgedConfig

distances = np.arange(0.3, 1.5, 0.1)

molecules = []

for dist in distances:
    molecule = Molecule(
        geometry=[["H", [0.0, 0.0, 0.0]], ["H", [0.0, 0.0, dist]]],
        charge=0,
        multiplicity=1,
    )
    molecules = molecules + [molecule]

bitstrings = [[1, 0], [0, 1]]
ansatz = TwoLocal(2, [], "cry", [[0, 1], [1, 0]], reps=1)

backend = Aer.get_backend("statevector_simulator")
converter = QubitConverter(JordanWignerMapper())

config = EntanglementForgedConfig(backend=backend, maxiter=100)
calc = EntanglementForgedGroundStateSolver(
    qubit_converter=converter, ansatz=ansatz, bitstrings_u=bitstrings, config=config
)

energies = []

for molecule, distance in zip(molecules, distances):
    driver = PySCFDriver.from_molecule(molecule, basis="sto3g")
    problem = ElectronicStructureProblem(driver)
    problem.second_q_ops()
    res = calc.solve(problem)
    energies_history = res.get_energies_history()
    energy = [None]
    if len(energies_history) > 0:
        energy = res.get_energies_history()[-1]
    energies = energies + energy
    print(f"Distance = {distance}, Energy = {energy[0]}")

plt.plot(distances, energies)
plt.xlabel("Atomic distance (Angstrom)")
plt.ylabel("Energy")
plt.show()

import qiskit.tools.jupyter

%qiskit_version_table
%qiskit_copyright