"""Dataset and collate for dual-rate training.

Each item: a random (scene, t, staleness delta) triple ->
  - Qwen input ids for the text context ending at frame index t - delta
    (header + as many recent frames as fit the token budget + suffix)
  - normalized current state at t, normalized target delta t -> t+1
  - static object features, gravity/dt conditioning vector

Frame text rendering replicates src/training/data_loader._format_frame_text
byte-for-byte (including the baseline's duplicated "Frame N:" prefix, since
frame['description'] itself starts with "Frame N:").
"""

from __future__ import annotations

from pathlib import Path

import numpy as np
import torch
from torch.utils.data import Dataset

from .normalization import wrap_angle

VELOCITY_THRESHOLD = 1.0  # matches src/data/formats.py motion description
SUFFIX = "Predict next frame:"


def render_frame_text(frame_num: int, state: np.ndarray, ids: np.ndarray) -> str:
    """state: (N, 6). Replicates the baseline text format for one frame."""
    moving = int(np.sum((np.abs(state[:, 2]) > VELOCITY_THRESHOLD)
                        | (np.abs(state[:, 3]) > VELOCITY_THRESHOLD)))
    n = state.shape[0]
    if moving == 0:
        desc = "All objects are at rest"
    elif moving == n:
        desc = "All objects are in motion"
    else:
        desc = f"{moving} of {n} objects are moving"
    lines = [f"Frame {frame_num}: Frame {frame_num}: {desc}."]
    for i in range(n):
        px, py, vx, vy, a, av = (float(x) for x in state[i])
        s = f"  obj_{int(ids[i])}: pos=({px:.4f}, {py:.4f}), vel=({vx:.4f}, {vy:.4f})"
        if abs(a) > 0.001 or abs(av) > 0.001:
            s += f", a={a:.4f}, av={av:.4f}"
        lines.append(s)
    lines.append("")
    return "\n".join(lines)


def build_context_ids(tokenizer, header_text: str, states: np.ndarray, ids: np.ndarray,
                      ctx_end: int, max_ctx_tokens: int) -> list:
    """Token ids for header + recent frames [.. ctx_end] + suffix, newest-first fill.

    Frame index i in `states` is frame number i+1 (frames are 1-indexed in data).
    """
    head = tokenizer(header_text, add_special_tokens=False)["input_ids"]
    tail = tokenizer("\n" + SUFFIX, add_special_tokens=False)["input_ids"]
    budget = max_ctx_tokens - len(head) - len(tail)
    frame_ids = []
    for i in range(ctx_end, -1, -1):
        ft = render_frame_text(i + 1, states[i], ids)
        fi = tokenizer(ft + "\n", add_special_tokens=False)["input_ids"]
        if len(fi) > budget:
            break
        frame_ids.append(fi)
        budget -= len(fi)
    frame_ids.reverse()
    out = list(head)
    for fi in frame_ids:
        out.extend(fi)
    out.extend(tail)
    return out


class DualRateDataset(Dataset):
    def __init__(self, cache_dir: str, split: str, tokenizer, cfg, normalizer,
                 max_scenes: int | None = None, seed: int = 0):
        self.files = sorted(Path(cache_dir, split).glob("*.npz"))
        if max_scenes:
            rng = np.random.default_rng(seed)
            idx = rng.permutation(len(self.files))[:max_scenes]
            self.files = [self.files[i] for i in sorted(idx)]
        if not self.files:
            raise FileNotFoundError(f"no cache files in {cache_dir}/{split}")
        self.tokenizer = tokenizer
        self.cfg = cfg
        self.norm = normalizer
        self.seed = seed
        self.deterministic = False  # tests set True: t/delta derived from index

    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):
        z = np.load(self.files[idx])
        states, static, ids = z["states"], z["static"], z["ids"]
        header_text = str(z["header_text"])
        gravity, dt = z["gravity"], float(z["dt"])
        t_max = states.shape[0] - 2  # need t+1
        min_ctx = self.cfg.min_ctx_frames

        if self.deterministic:
            rng = np.random.default_rng(self.seed + idx)
        else:
            rng = np.random.default_rng()
        t = int(rng.integers(min_ctx, t_max + 1))
        delta = int(rng.integers(0, self.cfg.tick_k))
        ctx_end = max(t - delta, min_ctx - 1)

        input_ids = build_context_ids(self.tokenizer, header_text, states, ids,
                                      ctx_end, self.cfg.max_ctx_tokens)

        cur = states[t].astype(np.float32)
        target_delta = self.norm.compute_delta(cur, states[t + 1].astype(np.float32))
        # angles accumulate unbounded in the data; wrap the *input feature* so the
        # expert sees rotation state mod 2pi (the delta target is already wrapped)
        cur_feat = cur.copy()
        cur_feat[:, 4] = wrap_angle(cur_feat[:, 4])
        gravity_dt = np.array([gravity[0] / 981.0, gravity[1] / 981.0, dt * 60.0],
                              dtype=np.float32)
        return {
            "input_ids": input_ids,
            "cur_state": torch.from_numpy(self.norm.norm_state(cur_feat)),
            "target_delta": torch.from_numpy(self.norm.norm_delta(target_delta)),
            "static": torch.from_numpy(static.astype(np.float32)),
            "gravity_dt": torch.from_numpy(gravity_dt),
        }


def collate_dualrate(batch, pad_token_id: int):
    """Left-pad text, zero-pad objects to max N in batch, build masks."""
    b = len(batch)
    max_l = max(len(x["input_ids"]) for x in batch)
    max_n = max(x["cur_state"].shape[0] for x in batch)
    s_dim = batch[0]["cur_state"].shape[-1]
    f_dim = batch[0]["static"].shape[-1]

    input_ids = torch.full((b, max_l), pad_token_id, dtype=torch.long)
    attention_mask = torch.zeros((b, max_l), dtype=torch.long)
    cur_state = torch.zeros((b, max_n, s_dim))
    target_delta = torch.zeros((b, max_n, s_dim))
    static = torch.zeros((b, max_n, f_dim))
    obj_mask = torch.zeros((b, max_n), dtype=torch.bool)
    gravity_dt = torch.stack([x["gravity_dt"] for x in batch])

    for i, x in enumerate(batch):
        l = len(x["input_ids"])
        input_ids[i, max_l - l:] = torch.tensor(x["input_ids"], dtype=torch.long)
        attention_mask[i, max_l - l:] = 1
        n = x["cur_state"].shape[0]
        cur_state[i, :n] = x["cur_state"]
        target_delta[i, :n] = x["target_delta"]
        static[i, :n] = x["static"]
        obj_mask[i, :n] = True

    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "cur_state": cur_state,
        "target_delta": target_delta,
        "static": static,
        "obj_mask": obj_mask,
        "gravity_dt": gravity_dt,
    }