"""Provide result classes for connection discovery data in CR3BP.

This module provides data structures for storing and presenting the results
of connection discovery between manifolds in the Circular Restricted Three-Body
Problem (CR3BP). It includes both individual connection result records and
collection classes with convenient access and formatting methods.

The result classes provide a clean interface for accessing connection data
including transfer types, Delta-V requirements, intersection points, and
full 6D state information at connection points.

All coordinates and velocities are in nondimensional CR3BP rotating-frame units.

See Also
--------
:mod:`~hiten.algorithms.connections.base`
    Main ConnectionPipeline class that produces these results.
:mod:`~hiten.algorithms.connections.engine`
    ConnectionPipeline engine that generates result objects.
:mod:`~hiten.algorithms.connections.backends`
    Backend algorithms that compute connection data.
"""

from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Iterator, Literal, Mapping, Sequence, Tuple

import numpy as np

from hiten.algorithms.types.core import _DomainPayload

if TYPE_CHECKING:
    from hiten.system.manifold import Manifold


@dataclass
class ConnectionsBackendRequest:
    """Structured request for the connections backend."""

    points_u: np.ndarray
    points_s: np.ndarray
    states_u: np.ndarray
    states_s: np.ndarray
    traj_indices_u: np.ndarray | None
    traj_indices_s: np.ndarray | None
    eps: float
    dv_tol: float
    bal_tol: float
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class ConnectionsBackendResponse:
    """Structured response returned by the connections backend."""

    results: list["_ConnectionResult"]
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class _ConnectionResult:
    """Store an individual connection result between two manifolds.

    This dataclass stores all the information about a single discovered
    connection between source and target manifolds, including the transfer
    type, velocity change requirement, geometric location, and full state
    information at the connection point.

    Parameters
    ----------
    kind : {"impulsive", "ballistic"}
        Type of transfer. "ballistic" for very low Delta-V transfers
        (typically ||Delta-V|| <= ballistic_tol), "impulsive" for transfers
        requiring finite velocity changes.
    delta_v : float
        Magnitude of velocity change required for the transfer, in
        nondimensional CR3BP velocity units. ||v_source - v_target||.
    point2d : tuple of float
        2D coordinates (x, y) of the connection point on the synodic section
        plane, in nondimensional CR3BP distance units.
    state_u : np.ndarray, shape (6,)
        6D phase space state [x, y, z, vx, vy, vz] at the connection point
        on the source (typically unstable) manifold, in nondimensional
        CR3BP units.
    state_s : np.ndarray, shape (6,)
        6D phase space state [x, y, z, vx, vy, vz] at the connection point
        on the target (typically stable) manifold, in nondimensional
        CR3BP units.
    index_u : int
        Index of the connection point within the source manifold's
        section intersection data.
    index_s : int
        Index of the connection point within the target manifold's
        section intersection data.
    trajectory_index_u : int
        Index of the trajectory within the source manifold that
        produced the connection point.
    trajectory_index_s : int
        Index of the trajectory within the target manifold that
        produced the connection point.

    Notes
    -----
    ConnectionPipeline results are typically generated by the backend algorithms
    and sorted by increasing Delta-V requirement. The classification
    into "ballistic" vs "impulsive" is based on the ballistic_tol
    parameter in the search configuration.

    The 2D point coordinates correspond to the intersection location
    on the synodic section plane, while the 6D states provide the
    full phase space information needed for trajectory propagation.

    Examples
    --------
    >>> # Accessing connection data
    >>> print(f"Transfer type: {result.kind}")
    >>> print(f"Delta-V: {result.delta_v:.6f}")
    >>> print(f"Section point: {result.point2d}")
    >>> print(f"Source state: {result.state_u}")
    >>> print(f"Target state: {result.state_s}")

    See Also
    --------
    :class:`~hiten.algorithms.connections.types.Connections`
        Collection class for multiple connection results.
    :class:`~hiten.algorithms.connections.config._SearchConfig`
        Configuration that determines ballistic vs impulsive classification.
    """
    kind: Literal["impulsive", "ballistic"]
    delta_v: float
    point2d: Tuple[float, float]
    state_u: np.ndarray
    state_s: np.ndarray
    index_u: int
    index_s: int
    trajectory_index_u: int
    trajectory_index_s: int


@dataclass
class ConnectionResults:
    """Carrier for engine-level connection results.

    Encapsulates the backend-produced connection records so engines can
    return a standardized result object. Provides a helper to obtain the
    user-facing collection wrapper.

    Parameters
    ----------
    connections : list of :class:`~hiten.algorithms.connections.types._ConnectionResult`
        The connection results.

    Returns
    -------
    :class:`~hiten.algorithms.connections.types.Connections`
        The connection results.
    """

    connections: list[_ConnectionResult]

    def to_results(self) -> "Connections":
        """Convert to a :class:`~hiten.algorithms.connections.types.Connections` object.
        
        Returns
        -------
        :class:`~hiten.algorithms.connections.types.Connections`
            The connection results.
        """
        return Connections(self.connections)


class Connections:
    """Provide a collection of connection results with convenient access and formatting.

    This class provides a read-only sequence-like interface over a collection
    of :class:`~hiten.algorithms.connections.results._ConnectionResult` objects, with enhanced formatting capabilities
    for analysis and presentation. It behaves like a standard Python sequence
    while providing specialized string representations optimized for connection
    data.

    Parameters
    ----------
    results : sequence of :class:`~hiten.algorithms.connections.results._ConnectionResult` or None
        Collection of connection results to wrap. If None or empty,
        creates an empty results collection.

    Notes
    -----
    This class implements the sequence protocol, supporting:
    
    - Length queries: ``len(results)``
    - Iteration: ``for result in results``
    - Indexing: ``results[i]``
    - Boolean evaluation: ``bool(results)``
    
    The string representation provides a formatted table showing key
    connection information in a compact, readable format suitable for
    analysis and reporting.

    Examples
    --------
    >>> # Create results collection
    >>> results = Connections(connection_list)
    >>> print(f"Found {len(results)} connections")
    >>> 
    >>> # Access individual results
    >>> best_connection = results[0]  # Lowest Delta-V
    >>> print(f"Best Delta-V: {best_connection.delta_v:.6f}")
    >>> 
    >>> # Iterate over results
    >>> for result in results:
    ...     if result.kind == "ballistic":
    ...         print(f"Ballistic connection: {result.delta_v:.3e}")
    >>> 
    >>> # Formatted display
    >>> print(results)  # Shows formatted table

    See Also
    --------
    :class:`~hiten.algorithms.connections.results._ConnectionResult`
        Individual connection result data structure.
    :class:`~hiten.algorithms.connections.base.ConnectionPipeline`
        Main class that produces these result collections.
    """

    def __init__(self, results: Sequence[_ConnectionResult] | None):
        self._results: list[_ConnectionResult] = list(results) if results else []

    # Sequence-like methods
    def __len__(self) -> int:
        return len(self._results)

    def __iter__(self) -> Iterator[_ConnectionResult]:
        return iter(self._results)

    def __getitem__(self, idx: int) -> _ConnectionResult:
        return self._results[idx]

    def __bool__(self) -> bool:
        return bool(self._results)

    def __repr__(self) -> str:
        n_total = len(self._results)
        n_ballistic = sum(1 for r in self._results if r.kind == "ballistic")
        n_impulsive = n_total - n_ballistic
        return (
            f"Connections(n={n_total}, ballistic={n_ballistic}, impulsive={n_impulsive})"
        )

    def __str__(self) -> str:
        """Return a formatted table representation of all connection results.

        Returns
        -------
        str
            Multi-line formatted table showing connection details including
            index, transfer type, Delta-V, manifold indices, and state vectors.
            Returns "<no connection results>" if the collection is empty.

        Notes
        -----
        The table format includes:
        
        - Index: Sequential numbering of results
        - Kind: Transfer type ("ballistic" or "impulsive")
        - Delta-V: Velocity change magnitude in scientific notation
        - idx_u/idx_s: Manifold intersection indices
        - state_u/state_s: 6D state vectors with limited precision
        
        State vectors are formatted with 6 decimal places for readability
        while maintaining sufficient precision for analysis.
        """
        if not self._results:
            return "<no connection results>"

        # Header
        headers = ("#", "kind", "Delta-V", "idx_u", "idx_s", "state_u", "state_s")
        rows: list[tuple[str, ...]] = [headers]

        # Build rows with limited precision for readability
        for i, r in enumerate(self._results):
            su = np.asarray(r.state_u).ravel()
            ss = np.asarray(r.state_s).ravel()
            su_str = "[" + ", ".join(f"{v:.6f}" for v in su) + "]"
            ss_str = "[" + ", ".join(f"{v:.6f}" for v in ss) + "]"
            rows.append(
                (
                    str(i),
                    r.kind,
                    f"{r.delta_v:.3e}",
                    str(r.index_u),
                    str(r.index_s),
                    su_str,
                    ss_str,
                )
            )

        # Compute column widths
        col_widths = [max(len(row[c]) for row in rows) for c in range(len(headers))]

        def fmt_row(row: tuple[str, ...]) -> str:
            return "  ".join(cell.rjust(col_widths[i]) for i, cell in enumerate(row))

        # Assemble table
        lines = [fmt_row(rows[0])]
        lines.append("  ".join("-" * w for w in col_widths))
        for row in rows[1:]:
            lines.append(fmt_row(row))

        return "\n".join(lines)


@dataclass(frozen=True)
class ConnectionDomainPayload(_DomainPayload):
    """Domain payload carrying raw connection results and metadata."""

    @classmethod
    def _from_mapping(cls, data: Mapping[str, object]) -> "ConnectionDomainPayload":
        return cls(data=data)

    @property
    def connections(self) -> Tuple[_ConnectionResult, ...]:
        return tuple(self.require("connections"))

    @property
    def source(self):
        return self.get("source")

    @property
    def target(self):
        return self.get("target")


@dataclass(frozen=True)
class _ConnectionProblem:
    """Define a problem specification for connection discovery between two manifolds.

    This dataclass encapsulates all the parameters needed to define a connection
    discovery problem, including the source and target manifolds, compile-time
    configuration (section structure), and runtime options (search tolerances).

    Parameters
    ----------
    source : :class:`~hiten.system.manifold.Manifold` | :class:`~hiten.system.orbits.base.PeriodicOrbit`
        Source manifold (typically unstable manifold).
    target : :class:`~hiten.system.manifold.Manifold` | :class:`~hiten.system.orbits.base.PeriodicOrbit`
        Target manifold (typically stable manifold).
    section_axis : str
        Axis for the synodic section (e.g., "x", "y", "z"). From config.
    section_offset : float
        Offset value for the synodic section. From config.
    plane_coords : tuple of str
        Coordinate labels for the section plane projection. From config.
    direction : int or None
        Direction for section crossings (1, -1, or None for both). From config.
    delta_v_tol : float
        Maximum Delta-V tolerance for accepting a connection. From options.
    ballistic_tol : float
        Threshold for classifying connections as ballistic vs impulsive. From options.
    eps2d : float
        Radius for initial 2D pairing of points on the synodic section. From options.
    n_workers : int or None
        Number of parallel workers for computation. From options.

    Notes
    -----
    This class serves as a data container that packages all the necessary
    information for the connection engine to process. It combines compile-time
    structure (config) and runtime tuning (options) into a single problem
    specification.

    The problem specification is typically created by the high-level
    :class:`~hiten.algorithms.connections.base.ConnectionPipeline` class and passed
    to the engine for processing.

    Examples
    --------
    >>> problem = _ConnectionProblem(
    ...     source=unstable_manifold,
    ...     target=stable_manifold,
    ...     section_axis="x",
    ...     section_offset=0.8,
    ...     plane_coords=("y", "z"),
    ...     direction=1,
    ...     delta_v_tol=1e-3,
    ...     ballistic_tol=1e-8,
    ...     eps2d=1e-4,
    ...     n_workers=1
    ... )

    See Also
    --------
    :class:`~hiten.algorithms.connections.engine._ConnectionEngine`
        Engine class that processes this problem specification.
    :class:`~hiten.algorithms.connections.base.ConnectionPipeline`
        High-level class that creates these problem specifications.
    """
    source: "Manifold"
    target: "Manifold"
    section_axis: str
    section_offset: float
    plane_coords: tuple[str, str]
    direction: int | None
    delta_v_tol: float
    ballistic_tol: float
    eps2d: float
    n_workers: int | None
