"""
Created on Jun 25, 2016

@author: pab
"""

from typing import Any

import numpy as np
from numpy import pi, r_, sqrt
from numpy.typing import ArrayLike
from scipy import constants, linalg, optimize

from numdifftools._typing import Array, ArrayOrScalar, Differentiator, EstimateResult
from numdifftools.multicomplex import c_abs


class ClassicalHamiltonian:
    """
    Hamiltonian

    Parameters
    ----------
    n : scalar
        number of ions in the chain
    w : scalar
        angular trap frequency
    C : scalar
        Coulomb constant times the electronic charge in SI units.
    m : scalar
        the mass of a single trapped ion in the chain
    """

    n: int
    w: float
    C: float
    m: float

    def __init__(self) -> None:
        self.n = 2
        f = 1000000  # f is a scalar, it's the trap frequency
        self.w = 2 * pi * f
        self.C = (4 * pi * constants.epsilon_0) ** (-1) * constants.e**2
        # C is a scalar, it's the I
        self.m = 39.96 * 1.66e-27

    def potential(self, positionvector: ArrayLike) -> Any:
        """
        Return potential

        Parameters
        ----------
        positionvector: 1-d array (vector) of length n
            positions of the n ions
        """
        x = np.asarray(positionvector)
        w = self.w
        C = self.C
        m = self.m

        # First we consider the potential of the harmonic oscillator
        v_x = 0.5 * m * (w**2) * sum(x**2)
        # then we add the coulomb interaction:
        for i, xi in enumerate(x):
            for xj in x[i + 1 :]:
                v_x += C / (c_abs(xi - xj))
        return v_x

    def initialposition(self) -> Array:
        """Defines initial position as an estimate for the minimize process."""
        n = self.n
        x_0: Array = r_[-(n - 1) / 2 : (n - 1) / 2 : n * 1j]
        return x_0

    def normal_modes(self, eigenvalues: ArrayLike) -> Array:
        """Return normal modes

        Computed eigenvalues of the matrix Vx are of the form
            (normal_modes)**2*m.
        """
        m = self.m
        normal_modes: Array = sqrt(np.asarray(eigenvalues) / m)
        return normal_modes


def run_hamiltonian(
    hessian: Differentiator,
    verbose: bool = True,
    full_output: bool = True,
) -> tuple[ArrayOrScalar, ArrayOrScalar, Array]:
    c = ClassicalHamiltonian()

    xopt: Array = optimize.fmin(c.potential, c.initialposition(), xtol=1e-10)

    hessian.fun = c.potential

    if hasattr(hessian, "full_output"):
        hessian.full_output = full_output

    true_h: Array = np.array([[5.23748385e-12, -2.61873829e-12], [-2.61873829e-12, 5.23748385e-12]])

    result = hessian(xopt)
    if isinstance(result, EstimateResult):
        h: ArrayOrScalar = result.estimate
        error: ArrayOrScalar = result.error_estimate
    else:
        h = result
        error = np.abs(h - true_h)

    eigenvalues = linalg.eigvals(h)
    normal_modes = c.normal_modes(eigenvalues)

    if verbose:
        print(c.potential([-0.5, 0.5]))
        print(c.potential([-0.5, 0.0]))
        print(c.potential([0.0, 0.0]))
        print(xopt)
        print("h", h)
        print("h-true_h", np.abs(h - true_h))
        print("error_estimate", error)

        print("eigenvalues", eigenvalues)
        print("normal_modes", normal_modes)

    return h, error, true_h
