import numpy as np

from qiskit.providers.basicaer import BasicAer
from qiskit.utils import QuantumInstance
from qiskit.algorithms import VQE
from qiskit.algorithms.optimizers import L_BFGS_B
from qiskit.opflow.gradients import Gradient
from qiskit_nature.algorithms import AdaptVQE, VQEUCCFactory
from qiskit_nature.circuit.library import HartreeFock, UCC
from qiskit_nature.drivers import UnitsType
from qiskit_nature.drivers.second_quantization import PySCFDriver, HDF5Driver
from qiskit_nature.mappers.second_quantization import ParityMapper
from qiskit_nature.converters.second_quantization import QubitConverter
from qiskit_nature.problems.second_quantization import ElectronicStructureProblem
from qiskit_nature.properties.second_quantization.electronic import (
    ElectronicEnergy,
    ParticleNumber,
)
from qiskit_nature.transformers.second_quantization.electronic import ActiveSpaceTransformer
from qiskit_nature.properties.second_quantization.electronic.bases import ElectronicBasis
from qiskit_nature.properties.second_quantization.electronic.integrals import (
    OneBodyElectronicIntegrals,
    TwoBodyElectronicIntegrals,)

#OLD CODE
class AdaptVQE(GroundStateEigensolver):
    """A ground state calculation employing the AdaptVQE algorithm."""
    def __init__(
        self,
        qubit_converter: QubitConverter,
        solver: MinimumEigensolverFactory,
        threshold: float = 1e-5,
        delta: float = 1,
        max_iterations: Optional[int] = None,
    ) -> None:

#NEW CODE
class AdaptVQE(GroundStateEigensolver):
    """A ground state calculation employing the AdaptVQE algorithm."""
    def __init__(
        self,
        qubit_converter: QubitConverter,
        solver: MinimumEigensolverFactory,
        threshold: float = 1e-5,
        delta: float = 1,  #this is going to be deprecated
        max_iterations: Optional[int] = None,
        gradient: Optional[Gradient] = None,
    ) -> None:

 def _compute_gradients(
        self,
        theta: List[float],
        vqe: VQE,
    ) -> List[Tuple[float, PauliSumOp]]:
        res = []
        for exc in self._excitation_pool:
            # add next excitation to ansatz
            self._ansatz.operators = self._excitation_list + [exc]
            # set the current ansatz
            vqe.ansatz = self._ansatz
            # evaluate energies
            parameter_sets = theta + [-self._delta] + theta + [self._delta]
            energy_evaluation = vqe.get_energy_evaluation(self._main_operator)
            energy_results = energy_evaluation(np.asarray(parameter_sets))
            # compute gradient
            gradient = (energy_results[0] - energy_results[1]) / (2 * self._delta)
            res.append((np.abs(gradient), exc))

        return res

 def _compute_gradients(
        self,
        theta: List[float],
        vqe: VQE,
    ) -> List[Tuple[float, PauliSumOp]]:
        res=[]
        for exc in self._excitation_pool:
            self._ansatz.operators = self._excitation_list + [exc]
            if self.gradient.grad_method.analytic:
                vqe.ansatz = self._ansatz
            else:
                vqe.ansatz = self._ansatz.decompose()
            param_sets = vqe._ansatz_params 
            op = vqe.construct_expectation(theta, self._main_operator)
            state_grad = self.gradient.convert(operator=op, params=param_sets)
            # Assign the parameters and evaluate the gradient
            value_dict = {param_sets[-1]:0.0}
            state_grad_result = state_grad.assign_parameters(value_dict).eval()
            logger.info("State gradient computed with parameter shift", state_grad_result)
            res.append((np.abs(state_grad_result[-1]), exc))
        return res

driver = PySCFDriver(
atom="H .0 .0 .0; H .0 .0 0.735", unit=UnitsType.ANGSTROM, basis="sto3g"
        )
problem = ElectronicStructureProblem(driver)

expected = -1.85727503

qubit_converter = QubitConverter(ParityMapper())
solver=VQEUCCFactory(QuantumInstance(BasicAer.get_backend("statevector_simulator")))

%%time
#logging.basicConfig(level = logging.DEBUG)
grad=Gradient(grad_method="fin_diff")

calc= AdaptVQE(qubit_converter,solver,gradient=grad)
res = calc.solve(problem)
print("Energy calculated: ",res.electronic_energies[0])
print("Expected energy:   ", expected)

%%time
grad=Gradient(grad_method="param_shift")
calc= AdaptVQE(qubit_converter,solver,gradient=grad)
res = calc.solve(problem)
print("Energy calculated: ",res.electronic_energies[0])
print("Expected energy:   ", expected)

qubit_converter = QubitConverter(ParityMapper())
solver=VQEUCCFactory(QuantumInstance(BasicAer.get_backend("statevector_simulator")))
inter_dist = 1.6
driver1 = PySCFDriver(
    atom="Li .0 .0 .0; H .0 .0 " + str(inter_dist),
    unit=UnitsType.ANGSTROM,
    charge=0,
    spin=0,
    basis="sto3g",
)
transformer = ActiveSpaceTransformer(
    num_electrons=2,
    num_molecular_orbitals=3,
)
problem1 = ElectronicStructureProblem(driver1, [transformer])
solver = VQEUCCFactory(QuantumInstance(BasicAer.get_backend("statevector_simulator")))

%%time
grad=Gradient(grad_method="fin_diff")
calc = AdaptVQE(qubit_converter, solver,gradient=grad)
res = calc.solve(problem1)
print("Energy calculated: ",res.electronic_energies[0])

%%time
grad=Gradient(grad_method="param_shift")
calc = AdaptVQE(qubit_converter, solver,gradient=grad)
res = calc.solve(problem1)
print("Energy calculated: ",res.electronic_energies[0])