"""Center manifold backend for efficient Poincare map computations.

This module provides Numba-compiled kernels for computing center manifold
trajectories in the Circular Restricted Three-Body Problem (CR3BP). The
backend implements both Runge-Kutta and symplectic integration methods
with parallel processing capabilities for high-performance computation
of Poincare maps.

The main class :class:`~hiten.algorithms.poincare.centermanifold.backend._CenterManifoldBackend` 
provides the interface for center manifold computations, while the Numba-compiled functions
handle the low-level numerical integration and section crossing detection.
"""

from typing import Callable, Literal, Optional, Tuple

import numpy as np
from numba import njit, prange

from hiten.algorithms.dynamics.hamiltonian import (_eval_dH_dP, _eval_dH_dQ,
                                                   _hamiltonian_rhs)
from hiten.algorithms.dynamics.protocols import _HamiltonianSystemProtocol
from hiten.algorithms.integrators.rk import (RK4_A, RK4_B, RK4_C, RK6_A, RK6_B,
                                             RK6_C, RK8_A, RK8_B, RK8_C)
from hiten.algorithms.integrators.symplectic import (N_SYMPLECTIC_DOF,
                                                     _integrate_symplectic)
from hiten.algorithms.poincare.centermanifold.config import _get_section_config
from hiten.algorithms.poincare.core.backend import _ReturnMapBackend
from hiten.algorithms.poincare.core.events import _SurfaceEvent
from hiten.algorithms.poincare.utils import _hermite_scalar
from hiten.algorithms.polynomial.operations import _polynomial_evaluate
from hiten.algorithms.utils.config import FASTMATH
from hiten.utils.log_config import logger


@njit(cache=False, fastmath=FASTMATH, inline="always")
def _detect_crossing(section_coord: str, state_old: np.ndarray, state_new: np.ndarray, rhs_new: np.ndarray, n_dof: int) -> Tuple[bool, float]:
    """Detect if trajectory crossed the Poincare section.

    Parameters
    ----------
    section_coord : str
        Section coordinate identifier ('q2', 'p2', 'q3', or 'p3').
    state_old : ndarray, shape (2*n_dof,)
        Previous state vector [q1, q2, q3, p1, p2, p3].
    state_new : ndarray, shape (2*n_dof,)
        Current state vector [q1, q2, q3, p1, p2, p3].
    rhs_new : ndarray, shape (2*n_dof,)
        Current right-hand side of Hamiltonian equations.
    n_dof : int
        Number of degrees of freedom (typically 3 for CR3BP).

    Returns
    -------
    bool
        True if section crossing detected.
    float
        Interpolation parameter alpha for crossing time.
    """
    if section_coord == "q3":
        f_old = state_old[2]
        f_new = state_new[2]
    elif section_coord == "p3":
        f_old = state_old[n_dof + 2]
        f_new = state_new[n_dof + 2]
    elif section_coord == "q2":
        f_old = state_old[1]
        f_new = state_new[1]
    else:  # "p2"
        f_old = state_old[n_dof + 1]
        f_new = state_new[n_dof + 1]

    # Must have sign change
    if f_old * f_new >= 0.0:
        return False, 0.0

    # Direction check
    if section_coord == "q3":
        good_dir = state_new[n_dof + 2] > 0.0
    elif section_coord == "q2":
        good_dir = state_new[n_dof + 1] > 0.0
    elif section_coord == "p3":
        good_dir = rhs_new[2] > 0.0
    else:  # "p2"
        good_dir = rhs_new[1] > 0.0

    if not good_dir:
        return False, 0.0

    alpha = f_old / (f_old - f_new)
    return True, alpha


def _solve_bracketed(f, a: float, b: float, xtol: float = 1e-12, max_iter: int = 200) -> float:
    """Pure-Python Brent's method solver for general callables.

    Parameters
    ----------
    f : callable
        Scalar function whose root is being searched for.
    a : float
        Left endpoint of the bracket.
    b : float
        Right endpoint of the bracket.
    xtol : float, default=1e-12
        Root finding tolerance.
    max_iter : int, default=200

    Returns
    -------
    float or None
        Root value if found, None otherwise.
    """
    # Local bind for small speedup
    f_eval = f

    # Ensure floats and evaluate endpoints
    a = float(a)
    b = float(b)
    fa = float(f_eval(a))
    fb = float(f_eval(b))

    # Handle exact hits at the endpoints
    if fa == 0.0:
        return a
    if fb == 0.0:
        return b

    # Require a valid bracket
    if fa * fb > 0.0:
        return None

    # Initialize the third point of the bracket
    c = a
    fc = fa
    d = b - a
    e = d

    eps = float(np.finfo(np.float64).eps)

    for _ in range(int(max_iter)):
        if fb == 0.0:
            return b

        # If fb and fc have the same sign, move c to a
        if fb * fc > 0.0:
            c = a
            fc = fa
            d = b - a
            e = d

        # Ensure that |fc| >= |fb|
        if abs(fc) < abs(fb):
            a, b, c = b, c, b
            fa, fb, fc = fb, fc, fb

        tol = 2.0 * eps * abs(b) + 0.5 * xtol
        m = 0.5 * (c - b)
        if abs(m) <= tol:
            return b

        if abs(e) >= tol and abs(fa) > abs(fb):
            s = fb / fa
            if a == c:
                # Secant step
                p = 2.0 * m * s
                q = 1.0 - s
            else:
                # Inverse quadratic interpolation
                q_ = fa / fc
                r = fb / fc
                p = s * (2.0 * m * q_ * (q_ - r) - (b - a) * (r - 1.0))
                q = (q_ - 1.0) * (r - 1.0) * (s - 1.0)
            if p > 0.0:
                q = -q
            else:
                p = -p
            if (2.0 * p) < min(3.0 * m * q - abs(tol * q), abs(e * q)):
                e = d
                d = p / q
            else:
                d = m
                e = m
        else:
            d = m
            e = m

        a = b
        fa = fb
        if abs(d) > tol:
            b = b + d
        else:
            b = b + (tol if m > 0.0 else -tol)
        fb = float(f_eval(b))

    # Final tolerance check
    tol = 2.0 * eps * abs(b) + 0.5 * xtol
    m = 0.5 * (c - b)
    if abs(m) <= tol or fb == 0.0:
        return b
    return None

@njit(cache=False, fastmath=FASTMATH)
def _get_rk_coefficients(order: int):
    """Return Runge-Kutta coefficients for specified order.

    Parameters
    ----------
    order : int
        Integration order (4, 6, or 8).

    Returns
    -------
    A : ndarray
        Runge-Kutta A matrix.
    B : ndarray
        Runge-Kutta B vector.
    C : ndarray
        Runge-Kutta C vector.
    """
    if order == 4:
        return RK4_A, RK4_B, RK4_C
    elif order == 6:
        return RK6_A, RK6_B, RK6_C
    else:
        return RK8_A, RK8_B, RK8_C


@njit(cache=False, fastmath=FASTMATH)
def _integrate_rk_ham(y0: np.ndarray, t_vals: np.ndarray, A: np.ndarray, B: np.ndarray, C: np.ndarray, jac_H, clmo_H):
    """Integrate Hamiltonian system using Runge-Kutta method.

    Parameters
    ----------
    y0 : ndarray, shape (2*n_dof,)
        Initial state vector [q1, q2, q3, p1, p2, p3].
    t_vals : ndarray, shape (n_steps,)
        Time values for integration.
    A : ndarray
        Runge-Kutta A matrix.
    B : ndarray
        Runge-Kutta B vector.
    C : ndarray
        Runge-Kutta C vector.
    jac_H : ndarray
        Jacobian of Hamiltonian polynomial.
    clmo_H : ndarray
        CLMO table for polynomial evaluation.

    Returns
    -------
    traj : ndarray, shape (n_steps, 2*n_dof)
        Integrated trajectory.
    """
    n_steps = t_vals.shape[0]
    dim = y0.shape[0]
    n_stages = B.shape[0]
    traj = np.empty((n_steps, dim), dtype=np.float64)
    traj[0, :] = y0.copy()

    k = np.empty((n_stages, dim), dtype=np.float64)

    n_dof = dim // 2

    for step in range(n_steps - 1):
        t_n = t_vals[step]
        h = t_vals[step + 1] - t_n

        y_n = traj[step].copy()

        for s in range(n_stages):
            y_stage = y_n.copy()
            for j in range(s):
                a_sj = A[s, j]
                if a_sj != 0.0:
                    y_stage += h * a_sj * k[j]

            Q = y_stage[0:n_dof]
            P = y_stage[n_dof: 2 * n_dof]

            dQ = _eval_dH_dP(Q, P, jac_H, clmo_H)
            dP = -_eval_dH_dQ(Q, P, jac_H, clmo_H)

            k[s, 0:n_dof] = dQ
            k[s, n_dof: 2 * n_dof] = dP

        y_np1 = y_n.copy()
        for s in range(n_stages):
            b_s = B[s]
            if b_s != 0.0:
                y_np1 += h * b_s * k[s]

        traj[step + 1] = y_np1

    return traj


@njit(cache=False, fastmath=FASTMATH)
def _integrate_map(y0: np.ndarray, t_vals: np.ndarray, A: np.ndarray, B: np.ndarray, C: np.ndarray,
                   jac_H, clmo_H, order: int, c_omega_heuristic: float = 20.0, use_symplectic: bool = False):
    """Integrate Hamiltonian system using specified method.

    Parameters
    ----------
    y0 : ndarray, shape (2*n_dof,)
        Initial state vector [q1, q2, q3, p1, p2, p3].
    t_vals : ndarray, shape (n_steps,)
        Time values for integration.
    A : ndarray
        Runge-Kutta A matrix.
    B : ndarray
        Runge-Kutta B vector.
    C : ndarray
        Runge-Kutta C vector.
    jac_H : ndarray
        Jacobian of Hamiltonian polynomial.
    clmo_H : ndarray
        CLMO table for polynomial evaluation.
    order : int
        Integration order.
    c_omega_heuristic : float, default=20.0
        Heuristic parameter for symplectic integration.
    use_symplectic : bool, default=False
        If True, use symplectic integration; otherwise use Runge-Kutta.

    Returns
    -------
    traj : ndarray, shape (n_steps, 2*n_dof)
        Integrated trajectory.
    """

    if use_symplectic:
        traj = _integrate_symplectic(y0, t_vals, jac_H, clmo_H, order, c_omega_heuristic)
    else:
        traj = _integrate_rk_ham(y0, t_vals, A, B, C, jac_H, clmo_H)

    return traj


@njit(cache=False, fastmath=FASTMATH)
def _poincare_step(q2: float, p2: float, q3: float, p3: float, dt: float,
                   jac_H, clmo, order: int, max_steps: int, use_symplectic: bool,
                   n_dof: int, section_coord: str, c_omega_heuristic: float = 20.0):
    """Perform one Poincare map step for center manifold integration.

    Parameters
    ----------
    q2 : float
        q2 coordinate (nondimensional units).
    p2 : float
        p2 coordinate (nondimensional units).
    q3 : float
        q3 coordinate (nondimensional units).
    p3 : float
        p3 coordinate (nondimensional units).
    dt : float
        Integration time step (nondimensional units).
    jac_H : ndarray
        Jacobian of Hamiltonian polynomial.
    clmo : ndarray
        CLMO table for polynomial evaluation.
    order : int
        Integration order.
    max_steps : int
        Maximum number of integration steps.
    use_symplectic : bool
        If True, use symplectic integration.
    n_dof : int
        Number of degrees of freedom.
    section_coord : str
        Section coordinate identifier.
    c_omega_heuristic : float, default=20.0
        Heuristic parameter for symplectic integration.

    Returns
    -------
    flag : int
        1 if section crossing found, 0 otherwise.
    q2p : float
        q2 coordinate at crossing.
    p2p : float
        p2 coordinate at crossing.
    q3p : float
        q3 coordinate at crossing.
    p3p : float
        p3 coordinate at crossing.
    t_cross : float
        Time of section crossing.
    """

    state_old = np.zeros(2 * n_dof, dtype=np.float64)
    state_old[1] = q2
    state_old[2] = q3
    state_old[n_dof + 1] = p2
    state_old[n_dof + 2] = p3

    elapsed = 0.0
    for _ in range(max_steps):
        c_A, c_B, c_C = _get_rk_coefficients(order)
        traj = _integrate_map(y0=state_old, t_vals=np.array([0.0, dt]), A=c_A, B=c_B, C=c_C,
                              jac_H=jac_H, clmo_H=clmo, order=order,
                              c_omega_heuristic=c_omega_heuristic, use_symplectic=use_symplectic)
        state_new = traj[1]

        rhs_new = _hamiltonian_rhs(state_new, jac_H, clmo, n_dof)
        crossed, alpha = _detect_crossing(section_coord, state_old, state_new, rhs_new, n_dof)

        if crossed:
            rhs_old = _hamiltonian_rhs(state_old, jac_H, clmo, n_dof)

            q2p = _hermite_scalar(alpha, state_old[1],       state_new[1],       rhs_old[1],       rhs_new[1],       dt)
            p2p = _hermite_scalar(alpha, state_old[n_dof+1], state_new[n_dof+1], rhs_old[n_dof+1], rhs_new[n_dof+1], dt)
            q3p = _hermite_scalar(alpha, state_old[2],       state_new[2],       rhs_old[2],       rhs_new[2],       dt)
            p3p = _hermite_scalar(alpha, state_old[n_dof+2], state_new[n_dof+2], rhs_old[n_dof+2], rhs_new[n_dof+2], dt)

            t_cross = elapsed + alpha * dt
            return 1, q2p, p2p, q3p, p3p, t_cross

        state_old = state_new
        elapsed += dt

    return 0, 0.0, 0.0, 0.0, 0.0, 0.0


@njit(parallel=True, cache=False)
def _poincare_map(seeds: np.ndarray, dt: float, jac_H, clmo, order: int, max_steps: int,
                  use_symplectic: bool, n_dof: int, section_coord: str, c_omega_heuristic: float):
    """Compute Poincare map for multiple center manifold seeds in parallel.

    Parameters
    ----------
    seeds : ndarray, shape (n_seeds, 4)
        Array of initial seeds [q2, p2, q3, p3] (nondimensional units).
    dt : float
        Integration time step (nondimensional units).
    jac_H : ndarray
        Jacobian of Hamiltonian polynomial.
    clmo : ndarray
        CLMO table for polynomial evaluation.
    order : int
        Integration order.
    max_steps : int
        Maximum number of integration steps.
    use_symplectic : bool
        If True, use symplectic integration.
    n_dof : int
        Number of degrees of freedom.
    section_coord : str
        Section coordinate identifier.
    c_omega_heuristic : float
        Heuristic parameter for symplectic integration.

    Returns
    -------
    success : ndarray, shape (n_seeds,)
        Success flags (1 if crossing found, 0 otherwise).
    q2p_out : ndarray, shape (n_seeds,)
        q2 coordinates at crossings.
    p2p_out : ndarray, shape (n_seeds,)
        p2 coordinates at crossings.
    q3p_out : ndarray, shape (n_seeds,)
        q3 coordinates at crossings.
    p3p_out : ndarray, shape (n_seeds,)
        p3 coordinates at crossings.
    t_out : ndarray, shape (n_seeds,)
        Times of section crossings.
    """

    n_seeds = seeds.shape[0]
    success = np.zeros(n_seeds, dtype=np.int64)
    q2p_out = np.empty(n_seeds, dtype=np.float64)
    p2p_out = np.empty(n_seeds, dtype=np.float64)
    q3p_out = np.empty(n_seeds, dtype=np.float64)
    p3p_out = np.empty(n_seeds, dtype=np.float64)
    t_out = np.empty(n_seeds, dtype=np.float64)

    for i in prange(n_seeds):
        q2, p2, q3, p3 = seeds[i, 0], seeds[i, 1], seeds[i, 2], seeds[i, 3]

        flag, q2_new, p2_new, q3_new, p3_new, t_cross = _poincare_step(
            q2, p2, q3, p3, dt, jac_H, clmo, order, max_steps, use_symplectic,
            n_dof, section_coord, c_omega_heuristic
        )

        if flag == 1:
            success[i] = 1
            q2p_out[i] = q2_new
            p2p_out[i] = p2_new
            q3p_out[i] = q3_new
            p3p_out[i] = p3_new
            t_out[i] = t_cross

    return success, q2p_out, p2p_out, q3p_out, p3p_out, t_out


class _CenterManifoldBackend(_ReturnMapBackend):
    """Backend for center manifold computations in the CR3BP.

    This backend provides efficient computation of center manifold trajectories
    using Numba-compiled kernels for Hamiltonian integration and Poincare map
    evaluation. It supports both Runge-Kutta and symplectic integration methods.

    Parameters
    ----------
    dynsys : :class:`~hiten.algorithms.dynamics.protocols._HamiltonianSystemProtocol`
        Hamiltonian system providing polynomial representation and CLMO tables.
    surface : :class:`~hiten.algorithms.poincare.core.events._SurfaceEvent`
        Poincare section surface definition.
    section_coord : str
        Section coordinate identifier ('q2', 'p2', 'q3', or 'p3').
    h0 : float
        Energy level for center manifold (nondimensional units).
    forward : int, default=1
        Integration direction (1 for forward, -1 for backward).
    max_steps : int, default=2000
        Maximum integration steps per trajectory.
    method : {'fixed', 'symplectic', 'adaptive'}, default='adaptive'
        Integration method.
    order : int, default=8
        Integration order for Runge-Kutta methods.
    pre_steps : int, default=1000
        Pre-integration steps for trajectory stabilization.
    refine_steps : int, default=3000
        Refinement steps for root finding.
    bracket_dx : float, default=1e-10
        Initial bracket size for root finding.
    max_expand : int, default=500
        Maximum bracket expansion iterations.
    c_omega_heuristic : float, default=20.0
        Heuristic parameter for symplectic integration.

    Notes
    -----
    State vectors are ordered as [q1, q2, q3, p1, p2, p3].
    """

    def __init__(
        self,
        *,
        dynsys: "_HamiltonianSystemProtocol",
        surface: "_SurfaceEvent",
        section_coord: str,
        h0: float,
        forward: int = 1,
        max_steps: int = 2000,
        method: Literal["fixed", "adaptive", "symplectic"] = "adaptive",
        order: int = 8,
        pre_steps: int = 1000,
        refine_steps: int = 3000,
        bracket_dx: float = 1e-10,
        max_expand: int = 500,
        c_omega_heuristic: float = 20.0,
    ) -> None:
        super().__init__(
            dynsys=dynsys,
            surface=surface,
            forward=forward,
            method=method,
            order=order,
            pre_steps=pre_steps,
            refine_steps=refine_steps,
            bracket_dx=bracket_dx,
            max_expand=max_expand,
        )

        self._section_cfg = _get_section_config(section_coord)
        self._h0 = h0
        self._H_blocks = dynsys.poly_H()
        self._clmo_table = dynsys.clmo_table
        self._jac_H = dynsys.jac_H
        self._order = order
        self._max_steps = max_steps
        self._use_symplectic = method == "symplectic"
        self._n_dof = N_SYMPLECTIC_DOF
        self._c_omega_heuristic = c_omega_heuristic

    def _lift_plane_point(self, plane: tuple[float, float]) -> Optional[tuple[float, float, float, float]]:
        """Convert a 2-D plane point into a 4-tuple center manifold state.

        Parameters
        ----------
        plane : tuple[float, float]
            Plane coordinates (q2, p2) or (q3, p3) depending on section.

        Returns
        -------
        tuple[float, float, float, float] or None
            Center manifold state (q2, p2, q3, p3) if valid, None if outside
            Hill box or root finding failed.
        """

        cfg = self._section_cfg

        constraints = cfg.build_constraint_dict(**{
            cfg.plane_coords[0]: plane[0],
            cfg.plane_coords[1]: plane[1],
        })

        # _solve_missing_coord is a bound method; it already knows energy and
        # Hamiltonian data via self.  Passing them again shifts positional
        # parameters and breaks the signature.  Only provide the variable
        # name and the dict of fixed values.
        missing_val = self._solve_missing_coord(cfg.missing_coord, constraints)

        if missing_val is None:
            return None

        other_vals = [0.0, 0.0]
        idx = 0 if cfg.missing_coord == cfg.other_coords[0] else 1
        other_vals[idx] = missing_val

        return cfg.build_state(plane, tuple(other_vals))

    def _bracket_root(
        self,
        f: Callable[[float], float],
        initial: float = 1e-3,
        factor: float = 2.0,
        max_expand: int = 40,
        *,
        symmetric: bool = False,
        xtol: float = 1e-12,
    ) -> Optional[float]:
        """Find a positive root of function f using bracketing and root finding.

        Parameters
        ----------
        f : callable
            Function to find root of.
        initial : float, default=1e-3
            Initial bracket size.
        factor : float, default=2.0
            Bracket expansion factor.
        max_expand : int, default=40
            Maximum expansion iterations.
        symmetric : bool, default=False
            If True, use symmetric bracketing.
        xtol : float, default=1e-12
            Root finding tolerance.

        Returns
        -------
        float or None
            Root value if found, None otherwise.
        """

        # We require f(0) <= 0 so that a root can lie in (0, x).
        if f(0.0) > 0.0:
            return None

        a, b = self._expand_bracket(
            f,
            0.0,
            dx0=initial,
            grow=factor,
            max_expand=max_expand,
            crossing_test=lambda prev, curr: prev <= 0.0 < curr,
            symmetric=symmetric,
        )

        return float(_solve_bracketed(f, float(a), float(b), float(xtol), 200))


    def _solve_missing_coord(
        self,
        varname: str,
        fixed_vals: dict[str, float],
        initial_guess: float = 1e-3,
        expand_factor: float = 2.0,
        max_expand: int = 40,
        *,
        symmetric: bool = False,
        xtol: float = 1e-12,
    ) -> Optional[float]:
        """Solve for the turning-point value of one center manifold coordinate.

        Parameters
        ----------
        varname : str
            Variable name to solve for ('q1', 'q2', 'q3', 'p1', 'p2', 'p3').
        fixed_vals : dict[str, float]
            Dictionary of fixed coordinate values.
        initial_guess : float, default=1e-3
            Initial guess for root finding.
        expand_factor : float, default=2.0
            Bracket expansion factor.
        max_expand : int, default=40
            Maximum expansion iterations.
        symmetric : bool, default=False
            If True, use symmetric bracketing.
        xtol : float, default=1e-12
            Root finding tolerance.

        Returns
        -------
        float or None
            Turning-point coordinate value if found, None otherwise.

        Notes
        -----
        The turning point is found by solving H(q,p) = h0 for the specified
        coordinate, where H is the Hamiltonian and h0 is the energy level.
        """

        var_indices = {
            "q1": 0, "q2": 1, "q3": 2,
            "p1": 3, "p2": 4, "p3": 5,
        }

        if varname not in var_indices:
            raise ValueError(f"Unknown variable: {varname}")

        solve_idx = var_indices[varname]

        def _residual(x: float) -> float:
            state = np.zeros(6, dtype=np.complex128)

            # Apply fixed values.
            for name, val in fixed_vals.items():
                if name in var_indices:
                    state[var_indices[name]] = val

            state[solve_idx] = x

            return _polynomial_evaluate(self._H_blocks, state, self._clmo_table).real - self._h0

        root = self._bracket_root(
            _residual,
            initial=initial_guess,
            factor=expand_factor,
            max_expand=max_expand,
            symmetric=symmetric,
            xtol=xtol,
        )

        if root is None:
            logger.warning("Failed to locate %s turning point within search limits", varname)
            return None

        return root

    def _find_turning(
        self,
        q_or_p: str,
        initial_guess: float = 1e-3,
        expand_factor: float = 2.0,
        max_expand: int = 40,
        *,
        symmetric: bool = False,
        xtol: float = 1e-12,
    ) -> float:
        """Find the absolute turning-point value for specified coordinate.

        Parameters
        ----------
        q_or_p : str
            Coordinate name ('q2', 'p2', 'q3', or 'p3').
        initial_guess : float, default=1e-3
            Initial guess for root finding.
        expand_factor : float, default=2.0
            Bracket expansion factor.
        max_expand : int, default=40
            Maximum expansion iterations.
        symmetric : bool, default=False
            If True, use symmetric bracketing.
        xtol : float, default=1e-12
            Root finding tolerance.

        Returns
        -------
        float
            Absolute turning-point coordinate value.

        Raises
        ------
        RuntimeError
            If root finding fails to converge.

        Notes
        -----
        This method finds the turning point for the Hill boundary of the
        center manifold, where all other coordinates are set to zero.
        """

        fixed_vals = {coord: 0.0 for coord in ("q2", "p2", "q3", "p3") if coord != q_or_p}

        root = self._solve_missing_coord(
            q_or_p,
            fixed_vals,
            initial_guess,
            expand_factor,
            max_expand,
            symmetric=symmetric,
            xtol=xtol,
        )

        if root is None:
            logger.warning("Failed to locate %s turning point within search limits", q_or_p)
            raise RuntimeError("Root finding for Hill boundary did not converge.")

        return root

    def step_to_section(
        self,
        seeds: np.ndarray,
        *,
        dt: float = 1e-2,
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """Propagate center manifold seeds until the next Poincare section crossing.

        Parameters
        ----------
        seeds : ndarray, shape (n_seeds, 4)
            Array of initial seeds [q2, p2, q3, p3] (nondimensional units).
        dt : float, default=1e-2
            Integration time step (nondimensional units).

        Returns
        -------
        pts : ndarray, shape (n_crossings, 2)
            Plane coordinates of the section crossings.
        cm_states : ndarray, shape (n_crossings, 4)
            Center manifold coordinates of the crossings.
        times : ndarray, shape (n_crossings,)
            Times of section crossings.
        flags : ndarray, shape (n_seeds,)
            Success flags (1 if crossing found, 0 otherwise).

        Notes
        -----
        This method uses parallel Numba compilation for efficient computation
        of multiple trajectories. The integration continues until either a
        section crossing is detected or the maximum number of steps is reached.
        """

        if seeds.size == 0:
            return (
                np.empty((0, 2)),
                np.empty((0, 4)),
                np.empty((0,)),
                np.empty((0,), dtype=np.int64),
            )

        max_steps = int(np.ceil(20.0 / dt))

        flags, q2p_arr, p2p_arr, q3p_arr, p3p_arr, t_arr = _poincare_map(
            np.ascontiguousarray(seeds, dtype=np.float64),
            dt,
            self._jac_H,
            self._clmo_table,
            self._order,
            max_steps,
            self._use_symplectic,
            N_SYMPLECTIC_DOF,
            self._section_cfg.section_coord,
            self._c_omega_heuristic,
        )

        cfg = self._section_cfg
        pts_list: list[tuple[float, float]] = []
        states_list: list[tuple[float, float, float, float]] = []
        times_list: list[float] = []

        for i in range(flags.shape[0]):
            if flags[i]:
                state = (q2p_arr[i], p2p_arr[i], q3p_arr[i], p3p_arr[i])
                # Ensure the returned seed lies EXACTLY on the section plane so that
                # the following iteration does not detect the same crossing again.
                # We zero the coordinate that defines the section (q2, p2, q3 or p3)
                # to avoid re-registering the same hit in subsequent steps.
                if cfg.section_coord == "q3":
                    state = (state[0], state[1], 0.0, state[3])
                elif cfg.section_coord == "p3":
                    state = (state[0], state[1], state[2], 0.0)
                elif cfg.section_coord == "q2":
                    state = (0.0, state[1], state[2], state[3])
                else:  # "p2"
                    state = (state[0], 0.0, state[2], state[3])
                states_list.append(state)
                if cfg.plane_coords == ("q2", "p2"):
                    pts_list.append((state[0], state[1]))
                else:
                    pts_list.append((state[2], state[3]))
                times_list.append(float(t_arr[i]))

        return (
            np.asarray(pts_list, dtype=np.float64),
            np.asarray(states_list, dtype=np.float64),
            np.asarray(times_list, dtype=np.float64),
            np.asarray(flags, dtype=np.int64),
        )
