"""Provide Lie series transformations for polynomial normal form computations.

This module provides the core implementation of Lie series-based normalization
for polynomial Hamiltonians in the Circular Restricted Three-Body Problem (CR3BP).
The routines implement the fundamental operations needed for constructing normal
forms around equilibrium points using canonical transformations.

The Lie series method transforms polynomial Hamiltonians through canonical
transformations generated by polynomial generating functions. This enables
the construction of normal forms that simplify the dynamics around equilibrium
points by eliminating non-resonant terms.

See Also
--------
:func:`~hiten.algorithms.hamiltonian.lie._apply_poly_transform`
    Apply the computed generating function via Lie series.
:mod:`~hiten.algorithms.hamiltonian.center`
    Center manifold normal forms using these Lie transformations.
:mod:`~hiten.algorithms.hamiltonian.normal`
    Complete normal form computations around equilibrium points.
:mod:`~hiten.algorithms.polynomial.operations`
    Polynomial operations including Poisson brackets.

References
----------
Murdoch, N., Wiggins, S. (1993). The normal form for Hamiltonian systems with
small perturbations. Journal of Differential Equations, 106(1), 123-156.

Jorba, A., Masdemont, J. (1999). Dynamics in the center manifold of the
collinear points of the restricted three body problem. Physica D, 132(1-2), 189-213.
"""


import numpy as np
from numba import njit
from numba.typed import List

from hiten.algorithms.polynomial.base import (_decode_multiindex, _factorial,
                                              _make_poly)
from hiten.algorithms.polynomial.operations import (
    _polynomial_clean, _polynomial_poisson_bracket, _polynomial_zero_list)
from hiten.algorithms.utils.config import FASTMATH


@njit(fastmath=FASTMATH, cache=False)
def _solve_homological_equation(
p_elim: np.ndarray, 
n: int, 
eta: np.ndarray, 
clmo: np.ndarray) -> np.ndarray:
    """Solve homological equation for polynomial generating function coefficients.

    Compute the generating function coefficients that eliminate specified terms
    from a polynomial Hamiltonian through canonical transformation. The homological
    equation relates the generating function to the terms being eliminated.

    Parameters
    ----------
    p_elim : ndarray
        Coefficient array of polynomial terms to be eliminated, in nondimensional
        energy units. Each element corresponds to a monomial term.
    n : int
        Degree of the homogeneous polynomial terms being processed.
    eta : ndarray, shape (3,)
        Complex eigenvalue array [lambda, i*omega1, i*omega2] where lambda is
        the real eigenvalue and omega1, omega2 are imaginary frequencies, all
        in nondimensional frequency units.
    clmo : ndarray
        Packed multi-index table from :func:`~hiten.algorithms.polynomial.base._init_index_tables`.

    Returns
    -------
    ndarray
        Coefficient array for the generating function of degree n, in
        nondimensional action units (energy * time).

    Notes
    -----
    The homological equation for canonical transformations has the form:
    {H_0, G} + H_1 = N

    where H_0 is the quadratic normal form, G is the generating function,
    H_1 contains terms to eliminate, and N is the new normal form.

    For each monomial with multi-index k = [k0, k1, k2, k3, k4, k5], the
    generating function coefficient is:
    g_k = -h_k / ((k3-k0)*lambda + (k4-k1)*i*omega1 + (k5-k2)*i*omega2)

    Small divisor terms (denominator < 1e-14) are skipped to avoid numerical
    instability, corresponding to near-resonant terms that remain in the
    normal form.

    See Also
    --------
    :func:`~hiten.algorithms.hamiltonian.lie._apply_poly_transform`
        Apply the computed generating function via Lie series.
    :func:`~hiten.algorithms.polynomial.base._decode_multiindex`
        Decode packed multi-indices used in coefficient computation.

    References
    ----------
    Meyer, K.R., Hall, G.R. (1992). Introduction to Hamiltonian Dynamical
    Systems and the N-Body Problem. Springer-Verlag, Chapter 4.
    """
    p_G = np.zeros_like(p_elim)
    for i in range(p_elim.shape[0]):
        c = p_elim[i]
        if c != 0.0:
            k = _decode_multiindex(i, n, clmo)
            denom = ((k[3]-k[0]) * eta[0] +
                     (k[4]-k[1]) * eta[1] +
                     (k[5]-k[2]) * eta[2])

            if abs(denom) < 1e-14:
                continue
            p_G[i] = -c / denom
    return p_G


@njit(fastmath=FASTMATH, cache=False)
def _apply_poly_transform(
poly_H: List[np.ndarray], 
p_G_n: np.ndarray, 
deg_G: int, 
N_max: int, 
psi: np.ndarray, 
clmo: np.ndarray, 
encode_dict_list: List[dict], 
tol: float) -> List[np.ndarray]:
    """Apply Lie series transformation to polynomial Hamiltonian.

    Transform a polynomial Hamiltonian using a canonical transformation generated
    by a polynomial generating function via the Lie series method. This implements
    the exponential Lie operator through iterated Poisson brackets.

    Parameters
    ----------
    poly_H : List[ndarray]
        Original Hamiltonian polynomial coefficients organized by degree,
        in nondimensional energy units.
    p_G_n : ndarray
        Coefficient array for the generating function of degree deg_G,
        in nondimensional action units (energy * time).
    deg_G : int
        Degree of the generating function polynomial.
    N_max : int
        Maximum polynomial degree to retain in the transformed Hamiltonian.
    psi : ndarray
        Combinatorial index table from :func:`~hiten.algorithms.polynomial.base._init_index_tables`.
    clmo : ndarray
        Packed multi-index table from :func:`~hiten.algorithms.polynomial.base._init_index_tables`.
    encode_dict_list : List[dict]
        Lookup dictionaries mapping multi-indices to coefficient positions.
    tol : float
        Numerical tolerance for cleaning small coefficients to prevent
        accumulation of round-off errors.

    Returns
    -------
    List[ndarray]
        Transformed Hamiltonian polynomial coefficients organized by degree,
        in nondimensional energy units.

    Notes
    -----
    The Lie series transformation implements the canonical transformation:
    H' = exp(L_G) H = H + {H,G} + (1/2!)*{{H,G},G} + (1/3!)*{{{H,G},G},G} + ...

    where:
    - L_G is the Lie operator associated with generating function G
    - {A,B} denotes the Poisson bracket of polynomials A and B
    - The series is truncated when terms exceed the maximum degree N_max

    The transformation preserves the Hamiltonian structure and implements
    a canonical coordinate change that eliminates specific polynomial terms
    while maintaining the symplectic nature of the flow.

    The number of terms in the Lie series is determined by the degree
    constraints: for deg_G > 2, at most (N_max - deg_G)/(deg_G - 2) + 1
    terms can contribute to the final result within the degree bound.

    See Also
    --------
    :func:`~hiten.algorithms.hamiltonian.lie._solve_homological_equation`
        Compute generating function coefficients for term elimination.
    :func:`~hiten.algorithms.polynomial.operations._polynomial_poisson_bracket`
        Poisson bracket computation used in the Lie series.

    References
    ----------
    Deprit, A. (1969). Canonical transformations depending on a small parameter.
    Celestial Mechanics, 1(1), 12-30.

    Hori, G. (1966). Theory of general perturbations with unspecified canonical
    variables. Publications of the Astronomical Society of Japan, 18, 287-296.
    """
    # Initialize result by copying input polynomial
    poly_result = List()
    for i in range(N_max + 1):
        if i < len(poly_H):
            poly_result.append(poly_H[i].copy())
        else:
            poly_result.append(_make_poly(i, psi))
    
    # Build complete generator polynomial from single degree
    poly_G = _polynomial_zero_list(N_max, psi)
    if deg_G < len(poly_G):
        poly_G[deg_G] = p_G_n.copy()
    
    # Determine number of terms in Lie series
    if deg_G > 2:
        K = max(N_max, (N_max - deg_G) // (deg_G - 2) + 1)
    else:
        K = 1
    
    # Precompute factorials
    factorials = [_factorial(k) for k in range(K + 1)]
    
    # Initialize with H for Poisson bracket iteration
    poly_bracket = List()
    for i in range(len(poly_H)):
        poly_bracket.append(poly_H[i].copy())
    
    # Apply Lie series: H + {H,G} + (1/2!){{H,G},G} + ...
    for k in range(1, K + 1):
        # Compute next Poisson bracket
        poly_bracket = _polynomial_poisson_bracket(
            poly_bracket,
            poly_G,
            N_max,
            psi,
            clmo,
            encode_dict_list
        )
        poly_bracket = _polynomial_clean(poly_bracket, tol)
        
        # Add to result with factorial coefficient
        coeff = 1.0 / factorials[k]
        for d in range(min(len(poly_bracket), len(poly_result))):
            poly_result[d] += coeff * poly_bracket[d]
    
    return _polynomial_clean(poly_result, tol)
