r"""
hiten.system.orbits.halo
==================

Generation and refinement of halo periodic orbits about the collinear
libration points of the Circular Restricted Three-Body Problem (CRTBP).

The module provides the :pyclass:`HaloOrbit` class which

* synthesises an initial state from the third-order analytic expansion of
  Richardson (1980), yielding a fast and robust first guess for the full
  nonlinear halo orbit;
* refines this guess through a differential-correction procedure that
  enforces the periodicity conditions by solving a reduced :math:`2\times2` boundary
  value problem.

References
----------
Richardson, D. L. (1980). "Analytic construction of periodic orbits about the
collinear libration points".
"""

from typing import TYPE_CHECKING, Literal, Optional, Sequence

import numpy as np
from numpy.typing import NDArray

from hiten.algorithms.poincare.singlehit.backend import _y_plane_crossing
from hiten.algorithms.utils.types import SynodicState
from hiten.system.libration.base import LibrationPoint
from hiten.system.libration.collinear import (CollinearPoint, L1Point, L2Point,
                                              L3Point)
from hiten.system.orbits.base import PeriodicOrbit
from hiten.utils.log_config import logger

if TYPE_CHECKING:
    from hiten.algorithms.continuation.interfaces import \
        _OrbitContinuationConfig
    from hiten.algorithms.corrector.interfaces import _OrbitCorrectionConfig


class HaloOrbit(PeriodicOrbit):
    r"""
    Halo orbit class.

    Parameters
    ----------
    libration_point : CollinearPoint
        Target :pyclass:`hiten.system.libration.collinear.CollinearPoint` around
        which the halo orbit is computed.
    amplitude_z : float, optional
        :math:`z`-amplitude of the halo orbit in the synodic frame. Required if
        *initial_state* is None.
    zenith : {'northern', 'southern'}, optional
        Indicates the symmetry branch with respect to the :math:`x\,y`-plane.
        Required if *initial_state* is None.
    initial_state : Sequence[float] or None, optional
        Six-dimensional state vector
        :math:`[x,\,y,\,z,\,\dot{x},\,\dot{y},\,\dot{z}]` in the rotating
        synodic frame. When *None* an analytical initial guess is generated
        from *amplitude_z* and *zenith*.

    Attributes
    ----------
    amplitude_z : float or None
        :math:`z`-amplitude of the halo orbit in the synodic frame.
    zenith : {'northern', 'southern'} or None
        Indicates the symmetry branch with respect to the :math:`x\,y`-plane.

    Raises
    ------
    ValueError
        If the required amplitude or branch is missing and *initial_state*
        is *None*.
    TypeError
        If *libration_point* is not an instance of
        :pyclass:`CollinearPoint`.
    """
    
    _family = "halo"
    
    amplitude_z: Optional[float] # Amplitude of the halo orbit
    zenith: Optional[Literal["northern", "southern"]]

    def __init__(
            self, 
            libration_point: LibrationPoint, 
            amplitude_z: Optional[float] = None,
            zenith: Optional[Literal["northern", "southern"]] = None,
            initial_state: Optional[Sequence[float]] = None
        ):

        # Validate constructor parameters
        if initial_state is not None and (amplitude_z is not None or zenith is not None):
            raise ValueError("Cannot provide both an initial_state and analytical parameters (amplitude_z, zenith).")

        if not isinstance(libration_point, CollinearPoint):
            err = f"Halo orbits are only defined for CollinearPoint, but got {type(libration_point)}."
            logger.error(err)
            raise TypeError(err)
            
        if initial_state is None:
            if amplitude_z is None or zenith is None:
                err = "Halo orbits require an 'amplitude_z' (z-amplitude) and 'zenith' ('northern'/'southern') parameter when an initial_state is not provided."
                logger.error(err)
                raise ValueError(err)
            if not isinstance(libration_point, (L1Point, L2Point)):
                logger.warning(
                    "The analytical guess for L3 Halo orbits is experimental.\n "
                    "Convergence is not guaranteed and may require more iterations."
                )

        # Store user-supplied amplitude; will be replaced after correction
        self._amplitude_z = amplitude_z

        self.zenith = zenith

        super().__init__(libration_point, initial_state)

        if initial_state is not None:
            # Infer missing zenith
            if self.zenith is None:
                self.zenith = "northern" if self._initial_state[SynodicState.Z] > 0 else "southern"
            # Infer missing amplitude
            if self._amplitude_z is None:
                self._amplitude_z = self._initial_state[SynodicState.Z]

    @property
    def amplitude(self) -> float:
        """(Read-only) Current z-amplitude of the orbit in the synodic frame."""
        if getattr(self, "_initial_state", None) is not None:
            return float(self._initial_state[SynodicState.Z])
        return float(self._amplitude_z)

    @property
    def _correction_config(self) -> "_OrbitCorrectionConfig":
        """Provides the differential correction configuration for halo orbits."""
        from hiten.algorithms.corrector.interfaces import \
            _OrbitCorrectionConfig
        return _OrbitCorrectionConfig(
            event_func=_y_plane_crossing,
            residual_indices=(SynodicState.VX, SynodicState.VZ),
            control_indices=(SynodicState.X, SynodicState.VY),
            extra_jacobian=self._halo_quadratic_term
        )

    @property
    def _continuation_config(self) -> "_OrbitContinuationConfig":
        from hiten.algorithms.continuation.interfaces import \
            _OrbitContinuationConfig
        return _OrbitContinuationConfig(state=SynodicState.Z, amplitude=True)

    def _initial_guess(self) -> NDArray[np.float64]:
        r"""
        Richardson third-order analytical approximation.

        The method evaluates the closed-form expressions published by
        Richardson to obtain an :math:`O(\!\epsilon^{3})` approximation of the halo
        orbit where :math:`\epsilon` is the amplitude ratio.

        Returns
        -------
        numpy.ndarray
            State vector of shape (6,) containing
            :math:`[x,\,y,\,z,\,\dot{x},\,\dot{y},\,\dot{z}]` in the synodic
            frame and normalised CRTBP units.

        Notes
        -----
        The computation follows [Richardson1980]_.

        Examples
        --------
        >>> L1 = L1Point(system)
        >>> orb = HaloOrbit(L1, amplitude_z=0.01, zenith='northern')
        >>> y0 = orb._initial_guess()
        """
        # Determine sign (won) and which "primary" to use

        mu = self.mu
        amplitude_z = self._amplitude_z
        # Get gamma from the libration point instance property
        gamma = self.libration_point.gamma
        
        point_map = {
            L1Point: (+1, 1 - mu),
            L2Point: (-1, 1 - mu),
            L3Point: (+1, -mu)
        }
        
        point_type = type(self.libration_point)
        if point_type in point_map:
            won, primary = point_map[point_type]
        else:
            # This case should ideally not be hit due to __init__ checks, but provides a safeguard.
            raise ValueError(f"Analytical guess for Halo orbits is not supported for {self.libration_point.name} (got {point_type.__name__})")
        
        # Set n for northern/southern family
        n = 1 if self.zenith == "northern" else -1
        
        # Coefficients c(2), c(3), c(4)
        c = [0.0, 0.0, 0.0, 0.0, 0.0]  # just to keep 5 slots: c[2], c[3], c[4]
        
        if isinstance(self.libration_point, L3Point):
            for N in [2, 3, 4]:
                c[N] = (1 / gamma**3) * (
                    (1 - mu) + (-primary * gamma**(N + 1)) / ((1 + gamma)**(N + 1))
                )
        else:
            for N in [2, 3, 4]:
                c[N] = (1 / gamma**3) * (
                    (won**N) * mu 
                    + ((-1)**N)
                    * (primary * gamma**(N + 1) / ((1 + (-won) * gamma)**(N + 1)))
                )

        # Solve for lambda (the in-plane frequency)
        polylambda = [
            1,
            0,
            c[2] - 2,
            0,
            - (c[2] - 1) * (1 + 2 * c[2]),
        ]
        lambda_roots = np.roots(polylambda)

        # Pick the appropriate root based on L_i
        if isinstance(self.libration_point, L3Point):
            lam = abs(lambda_roots[2])  # third element in 0-based indexing
        else:
            lam = abs(lambda_roots[0])  # first element in 0-based indexing

        # Calculate parameters
        k = 2 * lam / (lam**2 + 1 - c[2])
        delta = lam**2 - c[2]

        d1 = (3 * lam**2 / k) * (k * (6 * lam**2 - 1) - 2 * lam)
        d2 = (8 * lam**2 / k) * (k * (11 * lam**2 - 1) - 2 * lam)

        a21 = (3 * c[3] * (k**2 - 2)) / (4 * (1 + 2 * c[2]))
        a22 = (3 * c[3]) / (4 * (1 + 2 * c[2]))
        a23 = - (3 * c[3] * lam / (4 * k * d1)) * (
            3 * k**3 * lam - 6 * k * (k - lam) + 4
        )
        a24 = - (3 * c[3] * lam / (4 * k * d1)) * (2 + 3 * k * lam)

        b21 = - (3 * c[3] * lam / (2 * d1)) * (3 * k * lam - 4)
        b22 = (3 * c[3] * lam) / d1

        d21 = - c[3] / (2 * lam**2)

        a31 = (
            - (9 * lam / (4 * d2)) 
            * (4 * c[3] * (k * a23 - b21) + k * c[4] * (4 + k**2)) 
            + ((9 * lam**2 + 1 - c[2]) / (2 * d2)) 
            * (
                3 * c[3] * (2 * a23 - k * b21) 
                + c[4] * (2 + 3 * k**2)
            )
        )
        a32 = (
            - (1 / d2)
            * (
                (9 * lam / 4) * (4 * c[3] * (k * a24 - b22) + k * c[4]) 
                + 1.5 * (9 * lam**2 + 1 - c[2]) 
                * (c[3] * (k * b22 + d21 - 2 * a24) - c[4])
            )
        )

        b31 = (
            0.375 / d2
            * (
                8 * lam 
                * (3 * c[3] * (k * b21 - 2 * a23) - c[4] * (2 + 3 * k**2))
                + (9 * lam**2 + 1 + 2 * c[2])
                * (4 * c[3] * (k * a23 - b21) + k * c[4] * (4 + k**2))
            )
        )
        b32 = (
            (1 / d2)
            * (
                9 * lam 
                * (c[3] * (k * b22 + d21 - 2 * a24) - c[4])
                + 0.375 * (9 * lam**2 + 1 + 2 * c[2])
                * (4 * c[3] * (k * a24 - b22) + k * c[4])
            )
        )

        d31 = (3 / (64 * lam**2)) * (4 * c[3] * a24 + c[4])
        d32 = (3 / (64 * lam**2)) * (4 * c[3] * (a23 - d21) + c[4] * (4 + k**2))

        s1 = (
            1 
            / (2 * lam * (lam * (1 + k**2) - 2 * k))
            * (
                1.5 * c[3] 
                * (
                    2 * a21 * (k**2 - 2) 
                    - a23 * (k**2 + 2) 
                    - 2 * k * b21
                )
                - 0.375 * c[4] * (3 * k**4 - 8 * k**2 + 8)
            )
        )
        s2 = (
            1 
            / (2 * lam * (lam * (1 + k**2) - 2 * k))
            * (
                1.5 * c[3] 
                * (
                    2 * a22 * (k**2 - 2) 
                    + a24 * (k**2 + 2) 
                    + 2 * k * b22 
                    + 5 * d21
                )
                + 0.375 * c[4] * (12 - k**2)
            )
        )

        a1 = -1.5 * c[3] * (2 * a21 + a23 + 5 * d21) - 0.375 * c[4] * (12 - k**2)
        a2 = 1.5 * c[3] * (a24 - 2 * a22) + 1.125 * c[4]

        l1 = a1 + 2 * lam**2 * s1
        l2 = a2 + 2 * lam**2 * s2

        deltan = -n  # matches the original code's sign usage

        # Solve for amplitude_x from the condition ( -del - l2*amplitude_z^2 ) / l1
        amplitude_x = np.sqrt((-delta - l2 * amplitude_z**2) / l1)

        # Evaluate the expansions at tau1 = 0
        tau1 = 0.0
        
        x = (
            a21 * amplitude_x**2 + a22 * amplitude_z**2
            - amplitude_x * np.cos(tau1)
            + (a23 * amplitude_x**2 - a24 * amplitude_z**2) * np.cos(2 * tau1)
            + (a31 * amplitude_x**3 - a32 * amplitude_x * amplitude_z**2) * np.cos(3 * tau1)
        )
        y = (
            k * amplitude_x * np.sin(tau1)
            + (b21 * amplitude_x**2 - b22 * amplitude_z**2) * np.sin(2 * tau1)
            + (b31 * amplitude_x**3 - b32 * amplitude_x * amplitude_z**2) * np.sin(3 * tau1)
        )
        z = (
            deltan * amplitude_z * np.cos(tau1)
            + deltan * d21 * amplitude_x * amplitude_z * (np.cos(2 * tau1) - 3)
            + deltan * (d32 * amplitude_z * amplitude_x**2 - d31 * amplitude_z**3) * np.cos(3 * tau1)
        )

        xdot = (
            lam * amplitude_x * np.sin(tau1)
            - 2 * lam * (a23 * amplitude_x**2 - a24 * amplitude_z**2) * np.sin(2 * tau1)
            - 3 * lam * (a31 * amplitude_x**3 - a32 * amplitude_x * amplitude_z**2) * np.sin(3 * tau1)
        )
        ydot = (
            lam
            * (
                k * amplitude_x * np.cos(tau1)
                + 2 * (b21 * amplitude_x**2 - b22 * amplitude_z**2) * np.cos(2 * tau1)
                + 3 * (b31 * amplitude_x**3 - b32 * amplitude_x * amplitude_z**2) * np.cos(3 * tau1)
            )
        )
        zdot = (
            - lam * deltan * amplitude_z * np.sin(tau1)
            - 2 * lam * deltan * d21 * amplitude_x * amplitude_z * np.sin(2 * tau1)
            - 3 * lam * deltan * (d32 * amplitude_z * amplitude_x**2 - d31 * amplitude_z**3) * np.sin(3 * tau1)
        )

        # Scale back by gamma using original transformation
        rx = primary + gamma * (-won + x)
        ry = -gamma * y
        rz = gamma * z

        vx = gamma * xdot
        vy = gamma * ydot
        vz = gamma * zdot

        # Return the state vector
        logger.debug(f"Generated initial guess for Halo orbit around {self.libration_point} with amplitude_z={self.amplitude}: {np.array([rx, ry, rz, vx, vy, vz], dtype=np.float64)}")
        return np.array([rx, ry, rz, vx, vy, vz], dtype=np.float64)

    def _halo_quadratic_term(self, X_ev, Phi):
        r"""
        Evaluate the quadratic part of the Jacobian for differential correction.

        Parameters
        ----------
        X_ev : numpy.ndarray, shape (6,)
            State vector at the event time (half-period).
        Phi : numpy.ndarray
            State-transition matrix evaluated at the same event
        Returns
        -------
        numpy.ndarray, shape (2, 2)
            Reduced Jacobian matrix employed by the
            :pyfunc:`hiten.system.orbits.base.PeriodicOrbit.correct`
            solver.
        """
        x, y, z, vx, vy, vz = X_ev
        mu2 = 1 - self.mu
        rho_1 = 1/(((x+self.mu)**2 + y**2 + z**2)**1.5)
        rho_2 = 1/(((x-mu2 )**2 + y**2 + z**2)**1.5)
        omega_x  = -(mu2*(x+self.mu)*rho_1) - (self.mu*(x-mu2)*rho_2) + x
        DDx = 2*vy + omega_x
        DDz = -(mu2*z*rho_1) - (self.mu*z*rho_2)

        if abs(vy) < 1e-9:
            logger.warning(f"Denominator 'vy' is very small ({vy:.2e}). Correction step may be inaccurate.")
            vy = np.sign(vy) * 1e-9 if vy != 0 else 1e-9
            
        return np.array([[DDx],[DDz]]) @ Phi[[SynodicState.Y],:][:, (SynodicState.X,SynodicState.VY)] / vy

    def eccentricity(self) -> float:
        """Eccentricity is not a well-defined concept for halo orbits."""
        return np.nan
