#!/usr/bin/env python3
"""
NexarDataset — loads pre-computed belief caches + collision labels.

Two data sources:
  train set  : nexar-collision-prediction/train.csv  (has time_of_event, target)
  test set   : nexar-collision-prediction/test.csv   (no labels, only IDs)

Cache format (from nexar_extractor.py):
  {
    "video_ids": [str, ...],
    "features": {
      vid_id: {
        "beliefs":   FloatTensor [n_windows, H]
        "tta_means": FloatTensor [n_windows]
        "tta_vars":  FloatTensor [n_windows]
        "p_alert":   FloatTensor [n_windows]
        "p_obs":     FloatTensor [n_windows]
        "p_silent":  FloatTensor [n_windows]
      }
    }
  }
"""
from __future__ import annotations

import logging
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import pandas as pd
import torch
from torch.utils.data import Dataset

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


class NexarTrainDataset(Dataset):
    """
    Dataset for Nexar TRAIN set domain adaptation.

    Each sample = one clip × its binary collision label.
    Positive class (target=1): a window ending at TTE before the collision.
    Negative class (target=0): a window from a non-collision video.
    """

    def __init__(
        self,
        cache_pos_file: str,    # .pt from nexar_extractor for positive videos
        cache_neg_file: str,    # .pt from nexar_extractor for negative videos
        n_windows: int = 3,
    ):
        self.n_windows = n_windows
        self.samples: List[dict] = []

        pos_cache = torch.load(cache_pos_file, map_location="cpu", weights_only=False)
        neg_cache = torch.load(cache_neg_file, map_location="cpu", weights_only=False)

        for vid_id, feat in pos_cache["features"].items():
            self.samples.append({
                "video_id": vid_id,
                "label": 1,
                "features": feat,
            })

        for vid_id, feat in neg_cache["features"].items():
            self.samples.append({
                "video_id": vid_id,
                "label": 0,
                "features": feat,
            })

        n_pos = sum(1 for s in self.samples if s["label"] == 1)
        n_neg = sum(1 for s in self.samples if s["label"] == 0)
        logger.info(f"NexarTrainDataset: {len(self.samples)} clips  pos={n_pos}  neg={n_neg}")

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

    def __getitem__(self, idx: int) -> dict:
        s = self.samples[idx]
        feat = s["features"]
        # Pad / truncate to n_windows
        n = feat["beliefs"].shape[0]
        if n >= self.n_windows:
            beliefs   = feat["beliefs"][-self.n_windows:]
            tta_means = feat["tta_means"][-self.n_windows:]
            tta_vars  = feat["tta_vars"][-self.n_windows:]
            p_alerts  = feat["p_alert"][-self.n_windows:]
        else:
            pad = self.n_windows - n
            beliefs   = torch.cat([feat["beliefs"][0:1].expand(pad, -1), feat["beliefs"]])
            tta_means = torch.cat([feat["tta_means"][0:1].expand(pad), feat["tta_means"]])
            tta_vars  = torch.cat([feat["tta_vars"][0:1].expand(pad),  feat["tta_vars"]])
            p_alerts  = torch.cat([feat["p_alert"][0:1].expand(pad),   feat["p_alert"]])

        return {
            "video_id": s["video_id"],
            "label":    torch.tensor(s["label"], dtype=torch.float32),
            "beliefs":  beliefs.float(),    # [T, H]
            "tta_means": tta_means.float(), # [T]
            "tta_vars":  tta_vars.float(),  # [T]
            "p_alerts":  p_alerts.float(),  # [T]
        }


class NexarTestDataset(Dataset):
    """Dataset for generating submission scores on the test set."""

    def __init__(self, cache_file: str, n_windows: int = 3):
        self.n_windows = n_windows
        cache = torch.load(cache_file, map_location="cpu", weights_only=False)
        self.video_ids = cache["video_ids"]
        self.features  = cache["features"]
        logger.info(f"NexarTestDataset: {len(self.video_ids)} test clips")

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

    def __getitem__(self, idx: int) -> dict:
        vid_id = self.video_ids[idx]
        feat = self.features[vid_id]
        n = feat["beliefs"].shape[0]
        if n >= self.n_windows:
            beliefs   = feat["beliefs"][-self.n_windows:]
            tta_means = feat["tta_means"][-self.n_windows:]
            tta_vars  = feat["tta_vars"][-self.n_windows:]
            p_alerts  = feat["p_alert"][-self.n_windows:]
        else:
            pad = self.n_windows - n
            beliefs   = torch.cat([feat["beliefs"][0:1].expand(pad, -1), feat["beliefs"]])
            tta_means = torch.cat([feat["tta_means"][0:1].expand(pad), feat["tta_means"]])
            tta_vars  = torch.cat([feat["tta_vars"][0:1].expand(pad),  feat["tta_vars"]])
            p_alerts  = torch.cat([feat["p_alert"][0:1].expand(pad),   feat["p_alert"]])

        return {
            "video_id": vid_id,
            "beliefs":  beliefs.float(),
            "tta_means": tta_means.float(),
            "tta_vars":  tta_vars.float(),
            "p_alerts":  p_alerts.float(),
        }


def nexar_collate_train(batch: List[dict]) -> dict:
    return {
        "video_ids": [b["video_id"] for b in batch],
        "labels":    torch.stack([b["label"]    for b in batch]),  # [B]
        "beliefs":   torch.stack([b["beliefs"]  for b in batch]),  # [B, T, H]
        "tta_means": torch.stack([b["tta_means"] for b in batch]), # [B, T]
        "tta_vars":  torch.stack([b["tta_vars"]  for b in batch]), # [B, T]
        "p_alerts":  torch.stack([b["p_alerts"]  for b in batch]), # [B, T]
    }


def nexar_collate_test(batch: List[dict]) -> dict:
    return {
        "video_ids": [b["video_id"] for b in batch],
        "beliefs":   torch.stack([b["beliefs"]  for b in batch]),
        "tta_means": torch.stack([b["tta_means"] for b in batch]),
        "tta_vars":  torch.stack([b["tta_vars"]  for b in batch]),
        "p_alerts":  torch.stack([b["p_alerts"]  for b in batch]),
    }