"""Variable-width trajectory contract shared by statistics and fitting."""

from __future__ import annotations

from collections.abc import Sequence
from dataclasses import dataclass

import numpy as np

from ._validation import active_dimensions
from .configuration import ActionCodec2Config

__all__ = ["ActiveTrajectory", "TrajectoryLike", "prepare_trajectory"]


@dataclass(frozen=True)
class ActiveTrajectory:
    """One complete trajectory over an explicit subset of codebook dimensions."""

    values: np.ndarray
    active_dims: tuple[int, ...]
    previous_first_order: np.ndarray | None = None

    def __init__(
        self,
        values: np.ndarray,
        active_dims: Sequence[int],
        *,
        previous_first_order: np.ndarray | None = None,
    ) -> None:
        object.__setattr__(self, "values", np.asarray(values))
        object.__setattr__(self, "active_dims", tuple(active_dims))
        previous = (
            None
            if previous_first_order is None
            else np.asarray(previous_first_order, dtype=np.float32)
        )
        if previous is not None and (
            previous.shape != (len(self.active_dims),) or not np.all(np.isfinite(previous))
        ):
            raise ValueError("previous_first_order must be finite with shape [D_active]")
        object.__setattr__(self, "previous_first_order", previous)

    @classmethod
    def from_padded(
        cls,
        values: np.ndarray,
        *,
        active_dims: Sequence[int],
        previous_first_order: np.ndarray | None = None,
    ) -> ActiveTrajectory:
        """Drop engineering-only dense padding before the semantic boundary."""
        array = np.asarray(values)
        dims = tuple(int(value) for value in active_dims)
        if array.ndim != 2 or not dims or min(dims) < 0 or max(dims) >= array.shape[1]:
            raise ValueError("padded values and active_dims do not define a valid trajectory")
        previous = None if previous_first_order is None else np.take(previous_first_order, dims)
        return cls(np.take(array, dims, axis=1), dims, previous_first_order=previous)


def trajectory_boundary(
    trajectory: TrajectoryLike, config: ActionCodec2Config, active_dims: tuple[int, ...]
) -> np.ndarray | None:
    """Expand explicit compact state to ``[D]`` for fitting and validation."""
    previous = getattr(trajectory, "previous_first_order", None)
    active = set(active_dims)
    required = any(
        c.primitive_order == 2 and active.intersection(c.dims)
        for c in config.physical.components.values()
    )
    if previous is None and required:
        raise ValueError("second-order trajectories require explicit previous_first_order")
    if previous is None and not any(
        c.primitive_order == 2 for c in config.physical.components.values()
    ):
        return None
    dense = np.zeros(config.physical.input_dim, dtype=np.float32)
    if previous is not None:
        dense[list(active_dims)] = previous
    return dense


TrajectoryLike = np.ndarray | ActiveTrajectory


def prepare_trajectory(
    trajectory: TrajectoryLike,
    *,
    dimension: int,
    index: int,
    dtype: np.dtype | type = np.float32,
) -> tuple[np.ndarray, tuple[int, ...]]:
    """Validate a trajectory and return compact values plus global dimension ids."""
    if isinstance(trajectory, ActiveTrajectory):
        raw = trajectory.values
        active_dims = trajectory.active_dims
    else:
        raw = trajectory
        array = np.asarray(raw)
        active_dims = tuple(range(array.shape[1])) if array.ndim == 2 else ()

    values = np.asarray(raw, dtype=dtype)
    if values.ndim != 2 or len(values) < 1 or values.shape[1] < 1:
        raise ValueError(
            f"trajectory {index} must have non-empty shape [T, D_active], got {values.shape}"
        )
    if len(active_dims) != values.shape[1]:
        raise ValueError(
            f"trajectory {index} active_dims has {len(active_dims)} entries for "
            f"width {values.shape[1]}"
        )
    if any(isinstance(item, bool) or not isinstance(item, int) for item in active_dims):
        raise ValueError(f"trajectory {index} active_dims must contain integers")
    active_dims = active_dimensions(
        active_dims, dimension, subject=f"trajectory {index} active_dims"
    )
    if not np.all(np.isfinite(values)):
        raise ValueError(f"trajectory {index} contains a non-finite value")
    return np.ascontiguousarray(values), active_dims