import warnings

warnings.filterwarnings("ignore")

import numpy as np
from matplotlib import pyplot as plt

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
from qiskit import Aer

import sys

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

radius_1 = 0.958  # position for the first H atom
radius_2 = 0.958  # position for the second H atom
thetas_in_deg = 104.478  # bond angles.

H1_x = radius_1
H2_x = radius_2 * np.cos(np.pi / 180 * thetas_in_deg)
H2_y = radius_2 * np.sin(np.pi / 180 * thetas_in_deg)

molecule = Molecule(
    geometry=[
        ["O", [0.0, 0.0, 0.0]],
        ["H", [H1_x, 0.0, 0.0]],
        ["H", [H2_x, H2_y, 0.0]],
    ],
    charge=0,
    multiplicity=1,
)
driver = PySCFDriver.from_molecule(molecule=molecule, basis="sto6g")
problem = ElectronicStructureProblem(driver)

converter = QubitConverter(JordanWignerMapper())

from qiskit_nature.algorithms.ground_state_solvers import (
    GroundStateEigensolver,
    NumPyMinimumEigensolverFactory,
)

solver = GroundStateEigensolver(
    converter, NumPyMinimumEigensolverFactory(use_default_filter_criterion=False)
)

result = solver.solve(problem)

print("Classical energy = ", result.total_energies[0])

orbitals_to_reduce = [0, 3]

from entanglement_forging import reduce_bitstrings

bitstrings = [[1, 1, 1, 1, 1, 0, 0], [1, 0, 1, 1, 1, 0, 1], [1, 0, 1, 1, 1, 1, 0]]
reduced_bitstrings = reduce_bitstrings(bitstrings, orbitals_to_reduce)

print(f"Bitstrings: {bitstrings}")
print(f"Bitstrings after orbital reduction: {reduced_bitstrings}")

from qiskit.circuit import Parameter, QuantumCircuit

theta = Parameter("θ")

hop_gate = QuantumCircuit(2, name="Hop gate")
hop_gate.h(0)
hop_gate.cx(1, 0)
hop_gate.cx(0, 1)
hop_gate.ry(-theta, 0)
hop_gate.ry(-theta, 1)
hop_gate.cx(0, 1)
hop_gate.h(0)

hop_gate.draw()

theta_1, theta_2, theta_3, theta_4 = (
    Parameter("θ1"),
    Parameter("θ2"),
    Parameter("θ3"),
    Parameter("θ4"),
)

ansatz = QuantumCircuit(5)
ansatz.append(hop_gate.to_gate({theta: theta_1}), [0, 1])
ansatz.append(hop_gate.to_gate({theta: theta_2}), [3, 4])
ansatz.append(hop_gate.to_gate({theta: 0}), [1, 4])
ansatz.append(hop_gate.to_gate({theta: theta_3}), [0, 2])
ansatz.append(hop_gate.to_gate({theta: theta_4}), [3, 4])

ansatz.draw("text", justify="right", fold=-1)

from entanglement_forging import Log

Log.VERBOSE = False

backend = Aer.get_backend("statevector_simulator")

config = EntanglementForgedConfig(
    backend=backend, maxiter=350, spsa_c0=20 * np.pi, initial_params=[0, 0, 0, 0]
)

calc = EntanglementForgedGroundStateSolver(
    qubit_converter=converter,
    ansatz=ansatz,
    bitstrings_u=reduced_bitstrings,
    config=config,
    orbitals_to_reduce=orbitals_to_reduce,
)
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()

import qiskit.tools.jupyter

%qiskit_version_table
%qiskit_copyright