"""Memory-mapped reader that turns native-rate packs into model-ready batches.

Shards hold decoded source media: faces and labels at the native frame rate,
unnormalized 16 kHz audio. Everything MuVAP+ specific happens here - sampling
onto the 25 Hz model timeline, loudness normalization, chunking, augmentation -
so the same shards can serve a model with different conventions.
"""

import json
import math
import random
from dataclasses import dataclass
from pathlib import Path
from typing import Sequence

import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, Sampler

from data.alignment import model_frame_count, source_indices
from data.augment import AugmentConfig, augment_visual, mix_audio
from data.media import WaveformStats, apply_normalization
from preprocess.asd.schema import FORMAT_NAME, MODEL_FPS, SAMPLE_RATE, SUPPORTED_VERSIONS


def _clip_id(sample_id: str) -> str:
    """Strip the entity suffix so same-recording samples share a key.

    AVA and WASD both name tracks `<clip>_<start>_<end>:<entity>`, where every
    entity of a clip may share one audio recording.
    """
    return sample_id.rsplit(":", 1)[0]


@dataclass(frozen=True)
class _Entry:
    root: Path
    shard: str
    metadata: dict


class PackedASDDataset(Dataset):
    """Random-access dataset backed by OS-page-cached NumPy memmaps."""

    def __init__(
        self,
        roots: str | Path | Sequence[str | Path],
        projection_window=None,
        window: int | None = None,
        context: int = 0,
        jitter: bool = False,
        augment: AugmentConfig | None = None,
        model_fps: float = MODEL_FPS,
        normalize: bool = True,
    ):
        """Read packed shards, optionally as fixed-length training chunks.

        `window` splits every track longer than it into near-equal chunks, so
        one epoch still covers every frame while batches stay uniform enough to
        pack tightly.

        `context` frames are read *before* each chunk and marked invalid in the
        returned mask. The model is causal, so a frame near the start of a chunk
        would otherwise be asked to predict from history the encoder was never
        shown. Feeding that history as unscored context costs a little compute
        and no storage, where overlapping the chunks would duplicate the data
        and still train on the truncated copies. A track's real first frames get
        no prefix: there the missing history is genuine, and the model has to
        cope with it at inference too.

        `jitter` slides the chunk grid within its track on every access, which
        varies the boundaries across epochs at no extra cost. All spans are
        counted in model frames, not source frames.
        """
        if isinstance(roots, (str, Path)):
            roots = [roots]
        self.projection_window = projection_window
        self.window = window
        self.context = max(0, context)
        self.jitter = jitter
        self.augment = augment
        self.model_fps = model_fps
        self.normalize = normalize
        self.samples_per_frame = round(SAMPLE_RATE / model_fps)
        self.entries: list[_Entry] = []
        self._arrays: dict[
            tuple[Path, str], tuple[np.ndarray, np.ndarray, np.ndarray]
        ] = {}

        for raw_root in roots:
            root = Path(raw_root)
            with (root / "dataset.json").open() as handle:
                manifest = json.load(handle)
            if (
                manifest.get("format") != FORMAT_NAME
                or manifest.get("version") not in SUPPORTED_VERSIONS
            ):
                raise ValueError(f"unsupported packed dataset format: {root}")
            if manifest.get("timeline") != "native":
                raise ValueError(f"pack is not on the native timeline: {root}")
            if manifest.get("audio_sample_rate") != SAMPLE_RATE:
                raise ValueError(f"unexpected audio sample rate: {root}")
            if manifest.get("audio_normalization") is not None:
                raise ValueError(f"pack has baked-in audio normalization: {root}")

            for shard_dir in sorted(root.glob("shard-*")):
                with (shard_dir / "index.json").open() as handle:
                    for metadata in json.load(handle):
                        self.entries.append(_Entry(root, shard_dir.name, metadata))

        self.model_frames = np.array(
            [
                model_frame_count(
                    entry.metadata["source_frames"],
                    entry.metadata["source_fps"],
                    model_fps,
                )
                for entry in self.entries
            ],
            dtype=np.int64,
        )
        self.chunks = self._build_chunks()
        self._clips = [_clip_id(entry.metadata["sample_id"]) for entry in self.entries]

    @property
    def predict_window(self) -> int | None:
        """Frames scored per chunk; the context prefix is read on top of this."""
        if self.window is None:
            return None
        return max(1, self.window - self.context)

    def _build_chunks(self) -> list[tuple[int, int, int]]:
        """Split tracks into near-equal *scored* spans.

        A chunk is identified by the span it predicts. Whatever context the
        loader prepends is read on top and never scored, so the spans still tile
        the track exactly and every frame is predicted once per epoch.
        """
        chunks: list[tuple[int, int, int]] = []
        span = self.predict_window
        for index in range(len(self.entries)):
            frames = int(self.model_frames[index])
            if frames == 0:
                continue
            if span is None or frames <= span:
                chunks.append((index, 0, frames))
                continue
            edges = np.linspace(
                0, frames, math.ceil(frames / span) + 1
            ).round().astype(int)
            for start, end in zip(edges[:-1], edges[1:]):
                chunks.append((index, int(start), int(end - start)))
        return chunks

    def _fed_length(self, start: int, length: int) -> int:
        return length + min(self.context, start)

    def chunk_lengths(self) -> np.ndarray:
        """Frames actually fed per chunk - what bounds memory and batch size."""
        return np.fromiter(
            (self._fed_length(start, length) for _, start, length in self.chunks),
            dtype=np.int64,
            count=len(self.chunks),
        )

    def __len__(self) -> int:
        return len(self.chunks)

    def _open(self, entry: _Entry):
        key = (entry.root, entry.shard)
        arrays = self._arrays.get(key)
        if arrays is None:
            shard_dir = entry.root / entry.shard
            faces = np.load(shard_dir / "faces.npy", mmap_mode="r", allow_pickle=False)
            audio = np.load(shard_dir / "audio.npy", mmap_mode="r", allow_pickle=False)
            labels = np.load(shard_dir / "labels.npy", mmap_mode="r", allow_pickle=False)
            arrays = faces, audio, labels
            self._arrays[key] = arrays
        return arrays

    def _stats(self, entry: _Entry) -> WaveformStats:
        return WaveformStats(
            mean=entry.metadata["audio_mean"],
            rms=entry.metadata["audio_rms"],
            peak=entry.metadata["audio_peak"],
        )

    def _read_audio(self, entry: _Entry, start: int, length: int) -> torch.Tensor:
        """Read one model-frame span and restore the recording's own gain."""
        _, audio, _ = self._open(entry)
        offset = entry.metadata["audio_offset"] + start * self.samples_per_frame
        available = entry.metadata["audio_samples"] - start * self.samples_per_frame
        wanted = length * self.samples_per_frame
        block = np.array(audio[offset : offset + min(wanted, max(available, 0))], copy=True)
        waveform = torch.from_numpy(block).float().div_(32768.0).unsqueeze(0)
        if waveform.shape[-1] < wanted:
            waveform = F.pad(waveform, (0, wanted - waveform.shape[-1]))
        if self.normalize:
            waveform = apply_normalization(waveform, self._stats(entry))
        return waveform

    def _donor_audio(self, clip: str, frames: int, rng) -> np.ndarray:
        """Read audio from a different clip than `clip`.

        Drawing from the whole dataset rather than the mini-batch matters here:
        every face track in a WASD clip shares one clip-level recording, and the
        batch sampler groups by length, which puts those tracks together. A
        batch-local donor would often be the very same waveform.
        """
        for _ in range(8):
            index = rng.randrange(len(self.entries))
            if self._clips[index] == clip:
                continue
            entry = self.entries[index]
            span = int(self.model_frames[index])
            start = rng.randrange(max(span - frames, 0) + 1) if span > frames else 0
            return self._read_audio(entry, start, min(frames, max(span, 1)))[0].numpy()
        return np.zeros(frames * self.samples_per_frame, dtype=np.float32)

    def __getitem__(self, index: int):
        entry_index, predict_start, predict_length = self.chunks[index]
        entry = self.entries[entry_index]
        faces, _, labels = self._open(entry)
        frames = int(self.model_frames[entry_index])
        if self.jitter and predict_length < frames:
            predict_start = random.randint(0, frames - predict_length)

        # Read a history prefix where one exists; at a track's real start there
        # is none, and that truncation is genuine rather than an artefact.
        prefix = min(self.context, predict_start)
        start = predict_start - prefix
        length = prefix + predict_length
        scored = torch.zeros(length, dtype=torch.bool)
        scored[prefix:] = True

        # One shared mapping puts faces and labels on the model timeline, which
        # is what keeps a label attached to the face it describes.
        rows = source_indices(
            entry.metadata["source_frames"], entry.metadata["source_fps"], self.model_fps
        )
        span = rows[start : start + length]
        face_block = np.asarray(faces[entry.metadata["face_offset"] + span])
        # Labels are one byte a frame, so the whole track is cheap to read. The
        # projection window needs 50 frames of history and 15 of future; taking
        # them from the track rather than the chunk keeps a chunk boundary from
        # fabricating silence in the targets.
        track_labels = np.asarray(
            labels[entry.metadata["label_offset"] + rows], dtype=np.uint8
        ).astype(np.float32)
        waveform = self._read_audio(entry, start, length)

        if self.augment is not None:
            # The module RNG is what Lightning seeds per worker, so augmentation
            # is reproducible from the run seed.
            rng = random
            clip = self._clips[entry_index]
            audio_block = waveform[0].numpy()
            if rng.random() < self.augment.substitute_audio_prob:
                audio_block = self._donor_audio(clip, length, rng)
                track_labels = np.zeros_like(track_labels)
            elif rng.random() < self.augment.audio_mix_prob:
                audio_block = mix_audio(
                    audio_block,
                    self._donor_audio(clip, length, rng),
                    rng.uniform(*self.augment.audio_snr_db),
                )
            waveform = torch.from_numpy(np.ascontiguousarray(audio_block)).unsqueeze(0)
            if self.augment.visual:
                face_block = augment_visual(face_block, self.augment, rng)

        visual = torch.from_numpy(np.ascontiguousarray(face_block)).float()
        track_vad = torch.from_numpy(track_labels)
        vad = track_vad[start : start + length]
        if self.projection_window is None:
            return waveform, visual, vad, scored, entry.metadata
        with torch.no_grad():
            vap = self.projection_window.get_labels(track_vad.view(1, 1, frames))
        return (
            waveform,
            visual,
            vap[0, start : start + length],
            vad,
            scored,
            entry.metadata,
        )


class ASDCollator:
    """Bring a batch to one length, by cropping or by padding.

    Training crops to the shortest member. That keeps the model free of a
    padding mask, and it is affordable because chunking bounds the spread of
    lengths and the batch sampler groups similar lengths together; measured on
    AVA plus WASD the pair loses 0.01% of frames per epoch, against 64% for
    random batches.

    Evaluation cannot crop: the official export needs every frame of every
    track, which is why it would otherwise run one track at a time. `pad=True`
    right-pads to the longest member instead and marks the filler unscored.
    Every stage of the model is causal - the CPC and Mimi frontends, the visual
    TCN, and both transformer stacks - so appended frames cannot change any
    earlier output, and the visual frontend's batch norm is on running
    statistics in eval mode. A test pins both properties.
    """

    def __init__(
        self, samples_per_frame: int = SAMPLE_RATE // int(MODEL_FPS), pad: bool = False
    ):
        self.samples_per_frame = samples_per_frame
        self.pad = pad

    @staticmethod
    def _to_length(tensor: torch.Tensor, length: int, axis: int) -> torch.Tensor:
        """Crop or right-pad one stream along `axis` (counted from the end)."""
        current = tensor.shape[-axis]
        if current == length:
            return tensor
        if current > length:
            index = [slice(None)] * tensor.ndim
            index[-axis] = slice(0, length)
            return tensor[tuple(index)]
        padding = [0, 0] * (axis - 1) + [0, length - current]
        return F.pad(tensor, padding)

    def __call__(self, batch):
        lengths = [item[1].shape[0] for item in batch]
        frames = max(lengths) if self.pad else min(lengths)
        samples = frames * self.samples_per_frame

        audio = torch.stack([self._to_length(item[0], samples, 1) for item in batch])
        visual = torch.stack([self._to_length(item[1], frames, 3) for item in batch])
        vap = torch.stack([self._to_length(item[2], frames, 2) for item in batch])
        vad = torch.stack([self._to_length(item[3], frames, 1) for item in batch])
        scored = torch.stack([self._to_length(item[4], frames, 1) for item in batch])
        return audio, visual, vap, vad, scored, [item[5] for item in batch]


class FrameBudgetBatchSampler(Sampler):
    """Group similar-length chunks into batches of roughly constant frame count.

    Holding `batch * frames` near a budget keeps three things stable that a
    fixed batch size does not: activation memory, the `B*T` population the
    visual frontend's batch norm sees, and the per-frame gradient noise of the
    mean-reduced losses. The batch size is clamped because the contrastive term
    is a `B`-way classification whose difficulty would otherwise swing with the
    chunk length.
    """

    def __init__(
        self,
        lengths,
        frame_budget: int,
        min_batch: int = 4,
        max_batch: int = 16,
        shuffle: bool = True,
        drop_last: bool = False,
        seed: int = 0,
        num_replicas: int = 1,
        rank: int = 0,
    ):
        if min_batch < 1 or max_batch < min_batch:
            raise ValueError("min_batch must be positive and not exceed max_batch")
        self.lengths = np.asarray(lengths, dtype=np.int64)
        self.frame_budget = frame_budget
        self.min_batch = min_batch
        self.max_batch = max_batch
        self.shuffle = shuffle
        self.drop_last = drop_last
        self.seed = seed
        self.num_replicas = num_replicas
        self.rank = rank
        self.epoch = 0
        self._batches = self._build()

    def set_epoch(self, epoch: int) -> None:
        self.epoch = epoch
        self._batches = self._build()

    def _build(self) -> list[list[int]]:
        order = np.argsort(-self.lengths, kind="stable")
        batches = []
        cursor = 0
        while cursor < len(order):
            longest = int(self.lengths[order[cursor]])
            size = min(
                max(self.frame_budget // max(longest, 1), self.min_batch), self.max_batch
            )
            batch = order[cursor : cursor + size]
            cursor += size
            if self.drop_last and len(batch) < self.min_batch:
                continue
            batches.append([int(index) for index in batch])

        if self.shuffle:
            rng = random.Random(self.seed + self.epoch)
            rng.shuffle(batches)
        if self.num_replicas > 1:
            # Every rank must run the same number of steps or DDP deadlocks.
            remainder = len(batches) % self.num_replicas
            if remainder:
                batches += batches[: self.num_replicas - remainder]
            batches = batches[self.rank :: self.num_replicas]
        return batches

    def __iter__(self):
        return iter(self._batches)

    def __len__(self) -> int:
        return len(self._batches)