r"""Tiny discrete state, subtask, advantage, and FAST-lite tokenization.

The layout mirrors the ordering used by :math:`\pi^*_{0.6}`: the model first
predicts a subtask, then receives an optional advantage indicator, and only
then predicts the discrete action representation.  ``omitted`` is represented
by the physical absence of an indicator token, as required for unconditional
classifier-free guidance training.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Iterable, Literal

import torch
from torch import Tensor


Advantage = Literal["positive", "negative", "omitted"]


@dataclass(frozen=True)
class CoordinateRange:
    minimum: float
    maximum: float

    def __post_init__(self) -> None:
        if self.maximum <= self.minimum:
            raise ValueError("coordinate maximum must be greater than minimum")


class CoordinateTokenizer:
    """Uniformly discretize XY state values into literal byte tokens 0..255."""

    bins = 256

    def __init__(
        self,
        x_range: tuple[float, float] = (-10.0, 10.0),
        y_range: tuple[float, float] = (-5.0, 5.0),
    ) -> None:
        self.ranges = (CoordinateRange(*x_range), CoordinateRange(*y_range))

    def encode(self, coordinates: Tensor | Iterable[Iterable[float]]) -> Tensor:
        values = torch.as_tensor(coordinates, dtype=torch.float32)
        if values.shape[-1] != 2:
            raise ValueError("coordinates must end in an XY dimension of size 2")
        low = values.new_tensor([axis.minimum for axis in self.ranges])
        high = values.new_tensor([axis.maximum for axis in self.ranges])
        unit = ((values - low) / (high - low)).clamp(0.0, 1.0)
        return torch.round(unit * (self.bins - 1)).long()

    def decode(self, tokens: Tensor | Iterable[Iterable[int]]) -> Tensor:
        encoded = torch.as_tensor(tokens, dtype=torch.long)
        if encoded.shape[-1] != 2:
            raise ValueError("coordinate tokens must end in an XY dimension of size 2")
        if encoded.numel() and (encoded.min() < 0 or encoded.max() >= self.bins):
            raise ValueError("coordinate tokens must be in the inclusive range 0..255")
        low = torch.tensor([axis.minimum for axis in self.ranges])
        high = torch.tensor([axis.maximum for axis in self.ranges])
        return low + encoded.float() / (self.bins - 1) * (high - low)


class SubtaskVocabulary:
    labels = (
        "keep_lane",
        "steer_left",
        "steer_right",
        "slow_down",
        "stop",
        "accelerate",
        "unknown",
    )

    def __init__(self, first_token_id: int = 260) -> None:
        self.first_token_id = int(first_token_id)
        self._label_to_token = {
            label: self.first_token_id + index for index, label in enumerate(self.labels)
        }
        self._token_to_label = {token: label for label, token in self._label_to_token.items()}

    @property
    def token_ids(self) -> tuple[int, ...]:
        return tuple(self._label_to_token[label] for label in self.labels)

    def classify(self, reasoning: str) -> str:
        text = reasoning.casefold()
        if any(term in text for term in ("stop", "halt", "stationary")):
            return "stop"
        if any(term in text for term in ("slow", "deceler", "reduce speed", "brake")):
            return "slow_down"
        if any(term in text for term in ("accelerat", "speed up")):
            return "accelerate"
        if any(term in text for term in ("left", "counterclockwise")):
            return "steer_left"
        if any(term in text for term in ("right", "clockwise")):
            return "steer_right"
        if any(term in text for term in ("lane", "straight", "follow", "continue", "proceed")):
            return "keep_lane"
        return "unknown"

    def encode(self, label_or_reasoning: str) -> int:
        label = (
            label_or_reasoning
            if label_or_reasoning in self._label_to_token
            else self.classify(label_or_reasoning)
        )
        return self._label_to_token[label]

    def decode(self, token_id: int | Tensor) -> str:
        try:
            return self._token_to_label[int(token_id)]
        except KeyError as error:
            raise ValueError(f"unknown subtask token: {int(token_id)}") from error


@dataclass(frozen=True)
class PolicySequence:
    """One policy token sequence plus masks needed by ``Pi06Mini``.

    ``labels`` is already shifted for next-token prediction.  The advantage
    token itself is never a prediction target because it is externally
    supplied conditioning.  ``action_prefix_mask`` allows the continuous
    expert to read observation, predicted subtask, and advantage (if present),
    while insulating it from teacher-forced FAST-lite tokens.
    """

    input_ids: Tensor
    labels: Tensor
    generated_token_start: int
    action_prefix_mask: Tensor
    subtask_index: int
    advantage_index: int | None
    action_start_index: int


class MiniTokenizer:
    r"""Compose the minimal :math:`\pi^*_{0.6}` policy vocabulary."""

    BOS = 256
    IMAGE = 257
    STATE = 258
    SUBTASK_QUERY = 259
    ADVANTAGE_POSITIVE = SUBTASK_QUERY + 1 + len(SubtaskVocabulary.labels)
    ADVANTAGE_NEGATIVE = ADVANTAGE_POSITIVE + 1
    ACTION_START = ADVANTAGE_NEGATIVE + 1
    ACTION_END = ACTION_START + 1
    ACTION_TOKEN_START = ACTION_END + 1
    ACTION_BINS = 64
    ACTION_TOKEN_END = ACTION_TOKEN_START + ACTION_BINS - 1
    METADATA_DEMO = ACTION_TOKEN_END + 1
    METADATA_UNDERSTEER = METADATA_DEMO + 1
    METADATA_STOP = METADATA_DEMO + 2
    METADATA_INTERVENTION = METADATA_DEMO + 3
    ACTION_MINIMUM = -4.0
    ACTION_MAXIMUM = 4.0

    def __init__(
        self,
        coordinate_tokenizer: CoordinateTokenizer | None = None,
        subtasks: SubtaskVocabulary | None = None,
    ) -> None:
        self.coordinates = coordinate_tokenizer or CoordinateTokenizer()
        self.subtasks = subtasks or SubtaskVocabulary(first_token_id=self.SUBTASK_QUERY + 1)

    @property
    def vocab_size(self) -> int:
        return self.METADATA_INTERVENTION + 1

    def encode_metadata(self, source_type: str | None) -> Tensor | None:
        if source_type is None:
            return None
        if source_type == "demo_success": return torch.tensor(self.METADATA_DEMO)
        if "understeer" in source_type: return torch.tensor(self.METADATA_UNDERSTEER)
        if "stop" in source_type: return torch.tensor(self.METADATA_STOP)
        if "intervention" in source_type: return torch.tensor(self.METADATA_INTERVENTION)
        raise ValueError(f"unknown metadata source_type: {source_type}")

    @property
    def action_token_ids(self) -> range:
        return range(self.ACTION_TOKEN_START, self.ACTION_TOKEN_END + 1)

    def build_observation(self, history_xy: Tensor, metadata: str | None = None) -> Tensor:
        coordinate_tokens = self.coordinates.encode(history_xy).reshape(-1).cpu()
        specials = torch.tensor([self.BOS, self.IMAGE, self.STATE], dtype=torch.long)
        query = torch.tensor([self.SUBTASK_QUERY], dtype=torch.long)
        parts = [specials, coordinate_tokens]
        metadata_token = self.encode_metadata(metadata)
        if metadata_token is not None: parts.append(metadata_token.reshape(1))
        parts.append(query)
        return torch.cat(parts)

    def build_value_input(self, history_xy: Tensor) -> Tensor:
        """Return observation tokens for the critic without outcome metadata.

        RECAP's value function estimates ``V(o, task)``.  Experience provenance
        and success/failure labels are training targets, not observable inputs;
        including them would let the critic solve the toy buffer by leakage.
        The single-task minimal experiment keeps the task command implicit.
        """

        coordinate_tokens = self.coordinates.encode(history_xy).reshape(-1).cpu()
        specials = torch.tensor([self.BOS, self.IMAGE, self.STATE], dtype=torch.long)
        return torch.cat((specials, coordinate_tokens))

    # Backward-compatible name used by the pi0.5 data utilities.
    build_prefix = build_observation

    def encode_subtask(self, reasoning: str) -> Tensor:
        return torch.tensor(self.subtasks.encode(reasoning), dtype=torch.long)

    def encode_advantage(self, advantage: Advantage | bool | None) -> Tensor | None:
        if advantage is True or advantage == "positive":
            return torch.tensor(self.ADVANTAGE_POSITIVE, dtype=torch.long)
        if advantage is False or advantage == "negative":
            return torch.tensor(self.ADVANTAGE_NEGATIVE, dtype=torch.long)
        if advantage is None or advantage == "omitted":
            return None
        raise ValueError("advantage must be positive, negative, omitted, bool, or None")

    def build_subtask_query(self, history_xy: Tensor) -> Tensor:
        """Observation prefix whose final token predicts the next subtask."""

        return self.build_observation(history_xy, metadata=None)

    def build_action_condition(
        self,
        history_xy: Tensor,
        subtask_token: int | Tensor,
        advantage: Advantage | bool | None,
    ) -> Tensor:
        """Build the inference prefix ``observation -> subtask -> advantage``."""

        pieces = [
            self.build_observation(history_xy, metadata=None),
            torch.as_tensor(subtask_token, dtype=torch.long).reshape(1).cpu(),
        ]
        advantage_token = self.encode_advantage(advantage)
        if advantage_token is not None:
            pieces.append(advantage_token.reshape(1))
        return torch.cat(pieces)

    @staticmethod
    def _orthonormal_dct_matrix(
        horizon: int, *, device: torch.device, dtype: torch.dtype
    ) -> Tensor:
        if horizon <= 0:
            raise ValueError("DCT horizon must be positive")
        time = torch.arange(horizon, device=device, dtype=dtype)
        frequency = torch.arange(horizon, device=device, dtype=dtype).unsqueeze(1)
        matrix = torch.cos(math.pi / horizon * (time + 0.5) * frequency)
        matrix[0] *= math.sqrt(1.0 / horizon)
        if horizon > 1:
            matrix[1:] *= math.sqrt(2.0 / horizon)
        return matrix

    def dct_actions(self, actions: Tensor | Iterable[Iterable[float]]) -> Tensor:
        values = torch.as_tensor(actions, dtype=torch.float32)
        if values.ndim != 2 or values.shape[-1] != 2 or values.shape[0] == 0:
            raise ValueError("actions must have non-empty shape [horizon, 2]")
        transform = self._orthonormal_dct_matrix(
            values.shape[0], device=values.device, dtype=values.dtype
        )
        return transform @ values

    def idct_actions(self, coefficients: Tensor | Iterable[Iterable[float]]) -> Tensor:
        values = torch.as_tensor(coefficients, dtype=torch.float32)
        if values.ndim != 2 or values.shape[-1] != 2 or values.shape[0] == 0:
            raise ValueError("DCT coefficients must have non-empty shape [horizon, 2]")
        transform = self._orthonormal_dct_matrix(
            values.shape[0], device=values.device, dtype=values.dtype
        )
        return transform.transpose(0, 1) @ values

    def encode_actions(self, actions: Tensor | Iterable[Iterable[float]]) -> Tensor:
        values = torch.as_tensor(actions, dtype=torch.float32)
        if values.ndim != 2 or values.shape[-1] != 2 or values.shape[0] == 0:
            raise ValueError("actions must have non-empty shape [horizon, 2]")
        values = values.clamp(self.ACTION_MINIMUM, self.ACTION_MAXIMUM)
        coefficients = self.dct_actions(values) / math.sqrt(values.shape[0])
        coefficients = torch.where(
            coefficients.abs() < 1e-6, torch.zeros_like(coefficients), coefficients
        )
        unit = (
            coefficients.clamp(self.ACTION_MINIMUM, self.ACTION_MAXIMUM)
            - self.ACTION_MINIMUM
        ) / (self.ACTION_MAXIMUM - self.ACTION_MINIMUM)
        payload = torch.round(unit * (self.ACTION_BINS - 1)).long().reshape(-1)
        payload = payload + self.ACTION_TOKEN_START
        boundaries = torch.tensor(
            [self.ACTION_START, self.ACTION_END], dtype=torch.long, device=payload.device
        )
        return torch.cat((boundaries[:1], payload, boundaries[1:]))

    def decode_actions(self, tokens: Tensor | Iterable[int]) -> Tensor:
        encoded = torch.as_tensor(tokens, dtype=torch.long)
        if encoded.ndim != 1:
            raise ValueError("action tokens must be a one-dimensional sequence")
        if encoded.numel() < 4 or encoded[0] != self.ACTION_START or encoded[-1] != self.ACTION_END:
            raise ValueError("action tokens must be enclosed by ACTION_START and ACTION_END")
        payload = encoded[1:-1]
        if payload.numel() % 2:
            raise ValueError("action token payload must contain complete XY pairs")
        if payload.min() < self.ACTION_TOKEN_START or payload.max() > self.ACTION_TOKEN_END:
            raise ValueError("action token payload is outside the reserved action token interval")
        quantized = payload - self.ACTION_TOKEN_START
        coefficients = (
            quantized.float() / (self.ACTION_BINS - 1)
            * (self.ACTION_MAXIMUM - self.ACTION_MINIMUM)
            + self.ACTION_MINIMUM
        )
        horizon = payload.numel() // 2
        return self.idct_actions(coefficients.reshape(horizon, 2) * math.sqrt(horizon))

    def build_policy_sequence(
        self,
        history_xy: Tensor,
        subtask: str,
        actions: Tensor,
        advantage: Advantage | bool | None,
        metadata: str | None = None,
    ) -> PolicySequence:
        observation = self.build_observation(history_xy, metadata)
        subtask_token = self.encode_subtask(subtask).reshape(1)
        advantage_token = self.encode_advantage(advantage)
        action_tokens = self.encode_actions(actions).cpu()
        pieces = [observation, subtask_token]
        if advantage_token is not None:
            pieces.append(advantage_token.reshape(1))
        pieces.append(action_tokens)
        input_ids = torch.cat(pieces)

        subtask_index = observation.numel()
        advantage_index = subtask_index + 1 if advantage_token is not None else None
        action_start_index = subtask_index + 1 + int(advantage_token is not None)

        labels = torch.full_like(input_ids, -100)
        # SUBTASK_QUERY predicts the subtask.  The external advantage indicator
        # is deliberately skipped as a target.  Advantage/subtask then predicts
        # ACTION_START and the remaining FAST-lite stream autoregressively.
        labels[subtask_index - 1] = input_ids[subtask_index]
        labels[action_start_index - 1 : -1] = input_ids[action_start_index:]

        action_prefix_mask = torch.zeros_like(input_ids, dtype=torch.bool)
        action_prefix_mask[:action_start_index] = True
        return PolicySequence(
            input_ids=input_ids,
            labels=labels,
            generated_token_start=subtask_index,
            action_prefix_mask=action_prefix_mask,
            subtask_index=subtask_index,
            advantage_index=advantage_index,
            action_start_index=action_start_index,
        )