#!/usr/bin/env python3
"""
PolicyDataset — loads policy label manifests for Stage 1 supervised warm-start.

Two modes:
  Image mode (default): loads raw frames; requires full VLM forward at each step.
      Used for: make_belief_cache.py, evaluate_policy.py.

  Cache mode (--belief_cache_dir):  loads pre-computed belief vectors from
      data/belief_cache/{split}.pt produced by make_belief_cache.py.
      Used for: warm_start_trainer.py  (fast, ~1000× speed-up).
      In cache mode __getitem__ returns belief/tta tensors instead of images.
"""

from __future__ import annotations

import json
import logging
from collections import Counter
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence

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

logger = logging.getLogger("Policy.dataset")

MAX_FRAMES = 8

ACTION_NAMES = {0: "SILENT", 1: "OBSERVE", 2: "ALERT"}

SAMPLING_SCHEMES = ("original", "uniform", "last_biased", "last_2s")
SOURCE_FILTERS = ("all", "nexar", "multisrc", "dada", "dad")


# ── frame loading (mirrors DPO/SFT dataset) ───────────────────────────────────

def _load_frame(src_dir: Path, frame_idx: int) -> Optional[Image.Image]:
    for fmt in ["{:03d}", "{:04d}", "{:05d}", "{:06d}", "{}"]:
        for ext in [".jpg", ".jpeg", ".png"]:
            p = src_dir / (fmt.format(frame_idx) + ext)
            if p.exists():
                try:
                    return Image.open(p).convert("RGB")
                except Exception:
                    pass
    return None


def _resample_indices(
    base: Sequence[int],
    n_frames: int,
    scheme: str = "original",
) -> List[int]:
    """Resample frame indices within the window [base[0], base[-1]].

    `base` is the manifest's baked frame_indices (typically length 8, event-window
    biased). We treat its min/max as the sampling window and redraw `n_frames`
    indices inside that window, rounded to int.

    Schemes:
        original      — return `base[:n_frames]` (classic behavior)
        uniform       — evenly spaced indices across [min, max]
        last_biased   — 25% of frames from first half, 75% from second half
        last_2s       — all frames crammed into the last 2 s (~60 frames @ 30 fps)
                        assumed at the tail of the window
    """
    if not base:
        return []
    if scheme == "original" or n_frames == len(base):
        return list(base[:n_frames])

    lo, hi = int(base[0]), int(base[-1])
    if hi <= lo:
        return [lo] * n_frames

    if scheme == "uniform":
        idx = np.linspace(lo, hi, n_frames)
    elif scheme == "last_biased":
        n_head = max(1, n_frames // 4)
        n_tail = n_frames - n_head
        mid = (lo + hi) // 2
        head = np.linspace(lo, mid, n_head, endpoint=False)
        tail = np.linspace(mid, hi, n_tail)
        idx = np.concatenate([head, tail])
    elif scheme == "last_2s":
        # assume 30 fps → last 2 s = last 60 frames, clamped to available window
        two_s = min(hi - lo, 60)
        start = hi - two_s
        idx = np.linspace(start, hi, n_frames)
    else:
        raise ValueError(f"unknown sampling scheme: {scheme}")

    return [int(round(x)) for x in idx]


def _load_frames(
    source_dir: str,
    frame_indices: List[int],
    n_frames: int = MAX_FRAMES,
) -> List[Image.Image]:
    src = Path(source_dir)
    imgs = []
    for idx in frame_indices[:n_frames]:
        img = _load_frame(src, idx)
        if img is not None:
            imgs.append(img)
    if not imgs:
        imgs = [Image.new("RGB", (384, 384), (64, 64, 64))]
    return imgs


# ── dataset ───────────────────────────────────────────────────────────────────

class PolicyDataset(Dataset):
    """
    Args:
        manifests          : list of paths to JSON files from make_policy_labels.py
        split              : "train" or "val" (for logging)
        belief_cache_path  : optional path to .pt file from make_belief_cache.py;
                             when supplied, __getitem__ returns cached tensors
                             instead of PIL images (fast training mode)
        debug              : if True, truncate to first debug_samples
        debug_samples      : cap on samples in debug mode
    """

    def __init__(
        self,
        manifests:          List[Any],
        split:              str           = "train",
        belief_cache_path:  Optional[Any] = None,
        debug:              bool          = False,
        debug_samples:      int           = 64,
        n_frames:           int           = MAX_FRAMES,
        sampling:           str           = "original",
        source_filter:      str           = "all",
    ):
        self.split          = split
        self.n_frames       = int(n_frames)
        self.sampling       = sampling
        self.source_filter  = source_filter
        assert sampling in SAMPLING_SCHEMES, (
            f"sampling must be one of {SAMPLING_SCHEMES}, got {sampling}")
        assert source_filter in SOURCE_FILTERS, (
            f"source_filter must be one of {SOURCE_FILTERS}, got {source_filter}")
        self.samples: List[dict] = []

        for m in manifests:
            m = Path(m)
            if not m.exists():
                logger.warning(f"Policy label manifest not found: {m}")
                continue
            with open(m) as f:
                data = json.load(f)
            chunk = data.get("samples", data if isinstance(data, list) else [])
            self.samples.extend(chunk)
            logger.info(
                f"Loaded {len(chunk)} samples from {m.name}  "
                f"labels={data.get('label_counts', {}) if isinstance(data, dict) else 'n/a'}  "
                f"excluded={data.get('excluded', {}) if isinstance(data, dict) else 'n/a'}"
            )

        # Source filter (applied after manifest load so filter obeys the
        # naming convention in the Stage K/L plan).
        if source_filter != "all":
            keep = {
                "nexar":    {"nexar"},
                "multisrc": {"nexar", "dada"},  # balanced63k already controls ratio
                "dada":     {"dada"},
                "dad":      {"dad"},
            }[source_filter]
            before = len(self.samples)
            self.samples = [s for s in self.samples if s.get("source") in keep]
            logger.info(
                f"source_filter={source_filter}: {before} → {len(self.samples)} samples"
            )

        if debug:
            self.samples = self.samples[:debug_samples]

        # ── optional belief cache ─────────────────────────────────────────────
        self._cache: Optional[dict] = None
        if belief_cache_path is not None:
            p = Path(belief_cache_path)
            if not p.exists():
                raise FileNotFoundError(f"Belief cache not found: {p}")
            cache = torch.load(p, map_location="cpu", weights_only=True)
            # Trim cache to match current sample count (debug mode may shrink samples)
            n = len(self.samples)
            # Clip caches ship key "beliefs" [N, D]; per-frame caches from
            # make_belief_cache_v2 ship "beliefs_frame" [N, T, D] + "valid_frames" [N, T].
            # Accept either: mean-pool across frames if given a per-frame cache so
            # downstream v3/v5/v6/v7 heads (which read _cache["beliefs"]) all work.
            if "beliefs" in cache:
                beliefs = cache["beliefs"][:n]
                if beliefs.dim() == 3:
                    dtype = beliefs.dtype
                    beliefs = beliefs.float().mean(dim=1).to(dtype)
            elif "beliefs_frame" in cache:
                raw = cache["beliefs_frame"][:n]  # [N, T, D]
                dtype = raw.dtype
                vf = cache.get("valid_frames")
                if vf is not None:
                    vmask = vf[:n].float().unsqueeze(-1)           # [N, T, 1]
                    denom = vmask.sum(dim=1).clamp(min=1.0)        # [N, 1]
                    beliefs = (raw.float() * vmask).sum(dim=1) / denom
                else:
                    beliefs = raw.float().mean(dim=1)
                beliefs = beliefs.to(dtype)
            else:
                raise KeyError(
                    f"Belief cache {p.name} has neither 'beliefs' nor 'beliefs_frame' key"
                )
            self._cache = {
                "beliefs":   beliefs,
                "tta_means": cache["tta_means"][:n],
                "tta_vars":  cache["tta_vars"][:n],
            }
            logger.info(
                f"Loaded belief cache from {p.name}  ({n} entries, "
                f"belief_dim={self._cache['beliefs'].shape[-1]})"
            )

        label_dist = Counter(ACTION_NAMES[s["action_label"]] for s in self.samples)
        cat_dist   = Counter(s["category"] for s in self.samples)
        mode       = "cached" if self._cache is not None else "image"
        logger.info(
            f"PolicyDataset [{split}, {mode}]: {len(self.samples)} samples.  "
            f"Labels: {dict(label_dist)}.  Categories: {dict(cat_dist)}"
        )

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

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        s = self.samples[idx]
        base = {
            "action_label": int(s["action_label"]),
            "ce_weight":    float(s["ce_weight"]),
            "category":     s["category"],
            "tta_raw":      float(s["tta_raw"]),   # -1.0 for non_ego / safe_neg
            "video_id":     s["video_id"],
        }

        if self._cache is not None:
            # Fast path: return pre-computed belief tensors
            base["belief"]   = self._cache["beliefs"][idx]    # [hidden_dim]
            base["tta_mean"] = self._cache["tta_means"][idx]  # scalar
            base["tta_var"]  = self._cache["tta_vars"][idx]   # scalar
        else:
            # Slow path: return raw images for VLM encoding.
            # Resample frame_indices if a non-default sampling scheme is requested.
            base_idx = s["frame_indices"]
            if self.sampling != "original" or self.n_frames != MAX_FRAMES:
                frame_idx = _resample_indices(base_idx, self.n_frames, self.sampling)
            else:
                frame_idx = base_idx
            base["images"]   = _load_frames(s["source_dir"], frame_idx,
                                            n_frames=self.n_frames)
            base["metadata"] = s.get("metadata", {})
            base["frame_indices_used"] = frame_idx

        return base


def policy_collate_fn(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
    """Works for both image-mode and cache-mode batches."""
    out: Dict[str, Any] = {
        "action_labels": torch.tensor([b["action_label"] for b in batch], dtype=torch.long),
        "ce_weights":    torch.tensor([b["ce_weight"]    for b in batch], dtype=torch.float32),
        "categories":    [b["category"]  for b in batch],
        "tta_raws":      torch.tensor([b["tta_raw"]      for b in batch], dtype=torch.float32),
        "video_ids":     [b["video_id"]  for b in batch],
    }
    if "belief" in batch[0]:
        # Cache mode
        out["beliefs"]   = torch.stack([b["belief"]   for b in batch])   # [B, H]
        out["tta_means"] = torch.stack([b["tta_mean"] for b in batch])   # [B]
        out["tta_vars"]  = torch.stack([b["tta_var"]  for b in batch])   # [B]
    else:
        # Image mode
        out["images"]    = [b["images"]   for b in batch]
        out["metadata"]  = [b["metadata"] for b in batch]
    return out