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 (
    EntanglementForgedConfig,
    EntanglementForgedDriver,
    EntanglementForgedGroundStateSolver,
)

# Coefficients that define the one-body terms of the Hamiltonain

hcore = np.array([[-1.12421758, -0.9652574], [-0.9652574, -1.12421758]])

# Coefficients that define the two-body terms of the Hamiltonian
mo_coeff = np.array([[0.54830202, 1.21832731], [0.54830202, -1.21832731]])

# Coefficients for the molecular orbitals

eri = np.array(
    [
        [
            [[0.77460594, 0.44744572], [0.44744572, 0.57187698]],
            [[0.44744572, 0.3009177], [0.3009177, 0.44744572]],
        ],
        [
            [[0.44744572, 0.3009177], [0.3009177, 0.44744572]],
            [[0.57187698, 0.44744572], [0.44744572, 0.77460594]],
        ],
    ]
)

driver = EntanglementForgedDriver(
    hcore=hcore,
    mo_coeff=mo_coeff,
    eri=eri,
    num_alpha=1,
    num_beta=1,
    nuclear_repulsion_energy=0.7199689944489797,
)
problem = ElectronicStructureProblem(driver)
problem.second_q_ops()

converter = QubitConverter(JordanWignerMapper())

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

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

ansatz.draw()

from qiskit import Aer

backend = Aer.get_backend("statevector_simulator")

config = EntanglementForgedConfig(
    backend=backend, maxiter=200, initial_params=[0, 0.5 * np.pi]
)

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.plot([e[1] for e in res.get_energies_history()[0:-1]])
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.plot([p[1] for p in res.get_parameters_history()[0:-1]])
plt.show()

import qiskit.tools.jupyter

%qiskit_version_table
%qiskit_copyright