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 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

import sys

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

# We start by setting up the chemical problem.

molecule = Molecule(
    geometry=[["H", [0.0, 0.0, 0.0]], ["H", [0.0, 0.0, 0.735]]],
    charge=0,
    multiplicity=1,
)
driver = PySCFDriver.from_molecule(molecule, basis="sto3g")
problem = ElectronicStructureProblem(driver)
problem.second_q_ops()

converter = QubitConverter(JordanWignerMapper())

# Prepare the bitstrings and the ansatz

bitstrings = [[1, 0], [0, 1]]

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

from qiskit import Aer

backend = Aer.get_backend("qasm_simulator")
config = EntanglementForgedConfig(backend=backend, copysample_job_size=100)

calc = EntanglementForgedGroundStateSolver(
    qubit_converter=converter, ansatz=ansatz, bitstrings_u=bitstrings, config=config
)
res = calc.solve(problem)

res

print("Energies (from only one paramset in each iteration):")
plt.plot([e[0] for e in res.get_energies_history()])
plt.show()
print("Schmidts (from only one paramset in each iteration):")
plt.plot([s[0] for s in res.get_schmidts_history()])
plt.show()
print("Parameters (from only one paramset in each iteration):")
plt.plot([p[0] for p in res.get_parameters_history()])
plt.show()

res

import qiskit.tools.jupyter

%qiskit_version_table
%qiskit_copyright