"""Prompt and FAST-lite tokenization for the minimal pi0.7 policy.

pi0.7 conditions the low-level policy on a richer prompt than earlier pi
models: task and subtask language, episode speed/quality/mistake metadata,
and control mode. This tokenizer represents those concepts
with explicit marker/value pairs so the architecture can be validated without
shipping a full language tokenizer.
"""

from __future__ import annotations

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

import torch
from torch import Tensor


ControlMode = Literal["joint", "ee"]


@dataclass(frozen=True)
class EpisodeMetadata:
    speed_steps: int | None = None
    quality: int | None = None
    mistake: bool | None = None
    control_mode: ControlMode = "joint"

    def __post_init__(self) -> None:
        if self.speed_steps is not None and self.speed_steps <= 0:
            raise ValueError("speed_steps must be positive")
        if self.quality is not None and not 1 <= self.quality <= 5:
            raise ValueError("quality must be in the inclusive range 1..5")
        if self.control_mode not in ("joint", "ee"):
            raise ValueError("control_mode must be 'joint' or 'ee'")


@dataclass(frozen=True)
class PolicySequence:
    input_ids: Tensor
    labels: Tensor
    action_prefix_mask: Tensor
    generated_token_start: int
    action_start_index: int
    prompt_length: int
    metadata_indices: tuple[int, ...]


class _LabelVocabulary:
    def __init__(self, labels: tuple[str, ...], first_token_id: int) -> None:
        self.labels = labels
        self.first_token_id = first_token_id
        self._to_token = {
            label: first_token_id + index for index, label in enumerate(labels)
        }
        self._to_label = {token: label for label, token in self._to_token.items()}

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

    def encode(self, label: str) -> int:
        try:
            return self._to_token[label]
        except KeyError as error:
            raise ValueError(f"unknown categorical label: {label}") from error

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


class MiniTokenizer:
    """Explicit rich-prompt tokens plus a DCT-based FAST-lite action stream."""

    BOS = 256
    TASK_MARKER = 257
    SUBTASK_MARKER = 258
    SPEED_MARKER = 259
    QUALITY_MARKER = 260
    MISTAKE_MARKER = 261
    CONTROL_MARKER = 262
    NEXT_SUBTASK_MARKER = 263

    TASK_LABELS = (
        "drive",
        "pick_place",
        "fold",
        "open_close",
        "clean",
        "cook",
        "unknown",
    )
    TASK_TOKEN_START = 264
    TASK_TOKEN_END = TASK_TOKEN_START + len(TASK_LABELS) - 1

    SUBTASK_LABELS = (
        "keep_lane",
        "steer_left",
        "steer_right",
        "slow_down",
        "stop",
        "accelerate",
        "pick_place",
        "fold",
        "open_close",
        "unknown",
    )
    SUBTASK_TOKEN_START = TASK_TOKEN_END + 1
    SUBTASK_TOKEN_END = SUBTASK_TOKEN_START + len(SUBTASK_LABELS) - 1

    SPEED_BINS = (500, 1000, 2000, 4000, 8000, 16000)
    SPEED_TOKEN_START = SUBTASK_TOKEN_END + 1
    SPEED_TOKEN_END = SPEED_TOKEN_START + len(SPEED_BINS) - 1
    QUALITY_TOKEN_START = SPEED_TOKEN_END + 1
    QUALITY_TOKEN_END = QUALITY_TOKEN_START + 4
    MISTAKE_FALSE = QUALITY_TOKEN_END + 1
    MISTAKE_TRUE = MISTAKE_FALSE + 1
    CONTROL_JOINT = MISTAKE_TRUE + 1
    CONTROL_EE = CONTROL_JOINT + 1

    ACTION_START = CONTROL_EE + 1
    ACTION_END = ACTION_START + 1
    ACTION_TOKEN_START = ACTION_END + 1
    ACTION_BINS = 64
    ACTION_TOKEN_END = ACTION_TOKEN_START + ACTION_BINS - 1
    ACTION_MINIMUM = -4.0
    ACTION_MAXIMUM = 4.0

    def __init__(self) -> None:
        self.tasks = _LabelVocabulary(self.TASK_LABELS, self.TASK_TOKEN_START)
        self.subtasks = _LabelVocabulary(self.SUBTASK_LABELS, self.SUBTASK_TOKEN_START)

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

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

    @staticmethod
    def _classify_task(description: str) -> str:
        text = description.casefold()
        if any(word in text for word in ("drive", "lane", "road", "steer", "vehicle")):
            return "drive"
        if any(word in text for word in ("fold", "shirt", "towel", "laundry")):
            return "fold"
        if any(word in text for word in ("open", "close", "door", "drawer", "fridge")):
            return "open_close"
        if any(word in text for word in ("clean", "wipe", "wash", "tidy")):
            return "clean"
        if any(word in text for word in ("cook", "food", "recipe", "microwave")):
            return "cook"
        if any(word in text for word in ("pick", "place", "move", "put")):
            return "pick_place"
        return description if description in MiniTokenizer.TASK_LABELS else "unknown"

    @staticmethod
    def _classify_subtask(description: str) -> str:
        if description in MiniTokenizer.SUBTASK_LABELS:
            return description
        text = description.casefold()
        if any(word in text for word in ("stop", "halt", "stationary")):
            return "stop"
        if any(word in text for word in ("slow", "brake", "deceler")):
            return "slow_down"
        if any(word in text for word in ("accelerat", "speed up")):
            return "accelerate"
        if "left" in text:
            return "steer_left"
        if "right" in text:
            return "steer_right"
        if any(word in text for word in ("lane", "straight", "continue", "follow")):
            return "keep_lane"
        if "fold" in text:
            return "fold"
        if any(word in text for word in ("open", "close")):
            return "open_close"
        if any(word in text for word in ("pick", "place", "put")):
            return "pick_place"
        return "unknown"

    def encode_prompt(
        self,
        task: str,
        subtask: str | None,
        metadata: EpisodeMetadata,
        *,
        include_episode_metadata: bool = True,
        include_subtask: bool = True,
    ) -> tuple[Tensor, tuple[int, ...]]:
        """Encode prompt and return positions controlled by metadata CFG.

        Control mode is deliberately retained even in the unconditional branch,
        matching pi0.7 training where control-mode dropout is not applied.
        """

        ids = [self.BOS, self.TASK_MARKER, self.tasks.encode(self._classify_task(task))]
        if include_subtask and subtask is not None:
            ids.extend(
                (self.SUBTASK_MARKER, self.subtasks.encode(self._classify_subtask(subtask)))
            )

        metadata_indices: list[int] = []

        def add_metadata(marker: int, value: int) -> None:
            metadata_indices.extend((len(ids), len(ids) + 1))
            ids.extend((marker, value))

        if include_episode_metadata:
            if metadata.speed_steps is not None:
                bin_index = min(
                    range(len(self.SPEED_BINS)),
                    key=lambda index: abs(self.SPEED_BINS[index] - metadata.speed_steps),
                )
                add_metadata(self.SPEED_MARKER, self.SPEED_TOKEN_START + bin_index)
            if metadata.quality is not None:
                add_metadata(
                    self.QUALITY_MARKER, self.QUALITY_TOKEN_START + metadata.quality - 1
                )
            if metadata.mistake is not None:
                add_metadata(
                    self.MISTAKE_MARKER,
                    self.MISTAKE_TRUE if metadata.mistake else self.MISTAKE_FALSE,
                )
        control = self.CONTROL_JOINT if metadata.control_mode == "joint" else self.CONTROL_EE
        ids.extend((self.CONTROL_MARKER, control))
        return torch.tensor(ids, dtype=torch.long), tuple(metadata_indices)

    def encode_high_level_prompt(
        self,
        task: str,
        past_subtasks: Iterable[str],
        metadata: EpisodeMetadata,
    ) -> Tensor:
        """Encode task, past semantic decisions, metadata, and a next-subtask query."""

        prompt, _ = self.encode_prompt(
            task,
            subtask=None,
            metadata=metadata,
            include_subtask=False,
        )
        control_pair = prompt[-2:]
        body = prompt[:-2]
        history: list[int] = []
        for subtask in past_subtasks:
            history.extend(
                (self.SUBTASK_MARKER, self.subtasks.encode(self._classify_subtask(subtask)))
            )
        return torch.cat(
            (
                body,
                torch.tensor(history, dtype=torch.long),
                control_pair,
                torch.tensor([self.NEXT_SUBTASK_MARKER], dtype=torch.long),
            )
        )

    @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 FAST-lite 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,
        task: str,
        subtask: str | None,
        actions: Tensor,
        metadata: EpisodeMetadata,
        *,
        include_episode_metadata: bool = True,
        include_subtask: bool = True,
    ) -> PolicySequence:
        prompt, metadata_indices = self.encode_prompt(
            task,
            subtask,
            metadata,
            include_episode_metadata=include_episode_metadata,
            include_subtask=include_subtask,
        )
        action_tokens = self.encode_actions(actions).cpu()
        input_ids = torch.cat((prompt, action_tokens))
        action_start_index = prompt.numel()
        labels = torch.full_like(input_ids, -100)
        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,
            action_prefix_mask=action_prefix_mask,
            generated_token_start=0,
            action_start_index=action_start_index,
            prompt_length=prompt.numel(),
            metadata_indices=metadata_indices,
        )

    def build_metadata_cfg_pair(
        self,
        task: str,
        subtask: str | None,
        actions: Tensor,
        metadata: EpisodeMetadata,
    ) -> tuple[PolicySequence, PolicySequence]:
        """Return conditioned and episode-metadata-dropped policy sequences."""

        conditioned = self.build_policy_sequence(task, subtask, actions, metadata)
        unconditional = self.build_policy_sequence(
            task,
            subtask,
            actions,
            metadata,
            include_episode_metadata=False,
        )
        return conditioned, unconditional