"""
TeacherOutputCache — offline knowledge distillation support.

Stores precomputed teacher outputs (logits + mean-pooled hidden states) on disk
so the teacher never needs to run during student training.

Storage layout (all under cache_dir/):
    index.pkl    — dict {abs_clip_path: row_index}  (fully in RAM, tiny)
    logits.npy   — float32 (N,)                     (fully in RAM, ~8 MB for 2M clips)
    features.npy — float16 (N, num_layers, D)       (memory-mapped, ~16 GB for 2M clips)

Usage:
    # Build (one-time, ~20-30 min on 8×A100):
    cache = TeacherOutputCache.build(
        teacher=model.teacher,
        teacher_feature_layers=[0, 7, 15, 23],
        dataset=train_dataset,
        cache_dir="/path/to/badas_clips_output/teacher_cache",
        batch_size=256,
        device="cuda",
    )

    # Load and use:
    cache = TeacherOutputCache("/path/to/badas_clips_output/teacher_cache")
    teacher_logits, teacher_features = cache.lookup_batch(paths, device)
"""

import os
import pickle
import sys
import time
from typing import List, Optional, Tuple

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


class TeacherOutputCache:
    """
    Read-only view of precomputed teacher outputs.

    Features are stored as mean-pooled per-layer vectors (global spatial+temporal
    mean over the full token sequence). This makes storage feasible (~16 GB for 2M
    clips at 4 layers × 1024d) while preserving the semantic alignment signal needed
    for feature distillation.

    Logits are stored exactly (float32) — no approximation.
    """

    INDEX_FILE    = "index.pkl"
    LOGITS_FILE   = "logits.npy"
    FEATURES_FILE = "features.npy"

    def __init__(self, cache_dir: str):
        self._dir = cache_dir
        self._load()

    def _load(self):
        index_path    = os.path.join(self._dir, self.INDEX_FILE)
        logits_path   = os.path.join(self._dir, self.LOGITS_FILE)
        features_path = os.path.join(self._dir, self.FEATURES_FILE)

        for p in (index_path, logits_path, features_path):
            if not os.path.exists(p):
                raise FileNotFoundError(f"TeacherOutputCache file missing: {p}")

        with open(index_path, "rb") as f:
            self._index: dict = pickle.load(f)

        self._logits   = np.load(logits_path)                          # (N,) float32, in RAM
        self._features = np.load(features_path, mmap_mode="r")        # (N, L, D) float16, mmap

        N, L, D = self._features.shape
        print(f"[TeacherCache] Loaded: {N:,} clips | {L} feature layers | dim={D}")
        print(f"  logits:   {self._logits.nbytes / 1e6:.1f} MB (in RAM)")
        print(f"  features: {self._features.nbytes / 1e9:.1f} GB (memory-mapped)")

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

    def __contains__(self, path: str) -> bool:
        return path in self._index

    @property
    def num_layers(self) -> int:
        return self._features.shape[1]

    @property
    def feature_dim(self) -> int:
        return self._features.shape[2]

    # ------------------------------------------------------------------
    def lookup_batch(
        self,
        paths: List[str],
        device: torch.device,
    ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
        """
        Return precomputed teacher outputs for a batch of clip paths.

        Args:
            paths: List of absolute clip paths (length B).
            device: Target device for returned tensors.

        Returns:
            logits:   (B,)              float32 teacher logits.
            features: list of L tensors, each (B, D) float32 — one per layer.
        """
        # Map paths → row indices (missing paths raise KeyError → fail loudly)
        indices = [self._index[p] for p in paths]

        # numpy fancy-index: sort for mmap efficiency, then unsort
        order    = np.argsort(indices)
        unorder  = np.argsort(order)
        sorted_i = np.array(indices)[order]

        logits_np   = self._logits[sorted_i][unorder]                      # (B,)
        features_np = self._features[sorted_i][unorder].astype(np.float32) # (B, L, D)

        logits_t = torch.from_numpy(logits_np.copy()).to(device)

        # Split into per-layer tensors: list of L × (B, D)
        features_t = [
            torch.from_numpy(features_np[:, l, :].copy()).to(device)
            for l in range(features_np.shape[1])
        ]

        return logits_t, features_t

    # ------------------------------------------------------------------
    @classmethod
    def build(
        cls,
        teacher,                          # VideoClassifierBase (frozen)
        teacher_feature_layers: List[int],
        dataset,                          # VideoDataset with return_path=True
        cache_dir: str,
        batch_size: int = 256,
        num_workers: int = 8,
        device: str = "cuda",
    ) -> "TeacherOutputCache":
        """
        Precompute teacher outputs for every clip in *dataset* and save to disk.

        Args:
            teacher:               Frozen teacher model.
            teacher_feature_layers: Layer indices for hidden state extraction.
            dataset:               Dataset whose __getitem__ returns (video, label, path).
            cache_dir:             Directory to write index.pkl / logits.npy / features.npy.
            batch_size:            Videos per GPU batch (no gradient → can be large).
            num_workers:           DataLoader worker count.
            device:                CUDA device string.

        Returns:
            Loaded TeacherOutputCache instance.
        """
        os.makedirs(cache_dir, exist_ok=True)
        dev = torch.device(device)

        loader = DataLoader(
            dataset,
            batch_size=batch_size,
            shuffle=False,
            num_workers=num_workers,
            pin_memory=True,
            drop_last=False,
        )

        N          = len(dataset)
        num_layers = len(teacher_feature_layers)
        teacher_dim = teacher.feature_dim

        # Pre-allocate in RAM; we'll write features to mmap after the loop
        all_paths    = []
        all_logits   = np.zeros(N, dtype=np.float32)
        all_features = np.zeros((N, num_layers, teacher_dim), dtype=np.float16)

        teacher = teacher.to(dev)
        teacher.eval()

        idx  = 0
        t0   = time.time()

        print(f"\n[TeacherCache.build] {N:,} clips | batch={batch_size} | workers={num_workers}")
        print(f"  device={device}  layers={teacher_feature_layers}  dim={teacher_dim}")
        print(f"  output dir: {cache_dir}")
        sys.stdout.flush()

        with torch.no_grad():
            for step, batch in enumerate(loader):
                videos, labels, paths = batch
                B = videos.shape[0]
                videos = videos.to(dev, non_blocking=True)

                # ── teacher logits ──────────────────────────────────────────
                logits = teacher(videos)           # (B,) or (B, 1)
                logits = logits.view(-1)

                # ── hidden states at selected layers ────────────────────────
                backbone_out = teacher.backbone(
                    pixel_values_videos=videos,
                    skip_predictor=True,
                    output_hidden_states=True,
                    output_attentions=False,
                )

                # Extract + mean-pool each selected layer: (B, seq_len, D) → (B, D)
                feat_buf = np.zeros((B, num_layers, teacher_dim), dtype=np.float16)
                for li, layer_idx in enumerate(teacher_feature_layers):
                    hs = backbone_out.hidden_states[layer_idx]   # (B, seq_len, D)
                    pooled = hs.mean(dim=1).float().cpu().numpy()  # (B, D)
                    feat_buf[:, li, :] = pooled.astype(np.float16)

                all_paths.extend(list(paths))
                all_logits[idx:idx + B]      = logits.float().cpu().numpy()
                all_features[idx:idx + B]    = feat_buf

                idx += B

                if step % 20 == 0 or idx == N:
                    elapsed = time.time() - t0
                    rate    = idx / max(elapsed, 1e-3)
                    eta     = (N - idx) / max(rate, 1e-3)
                    ts      = time.strftime("%H:%M:%S")
                    print(
                        f"  [{ts}] step {step:4d} | {idx:7,}/{N:,} clips"
                        f" | {rate:.0f} clips/s | ETA {eta/60:.1f} min"
                    )
                    sys.stdout.flush()

        # ── save ─────────────────────────────────────────────────────────
        print(f"\n[TeacherCache.build] Saving to {cache_dir} ...")
        sys.stdout.flush()

        index = {path: i for i, path in enumerate(all_paths)}
        with open(os.path.join(cache_dir, cls.INDEX_FILE), "wb") as f:
            pickle.dump(index, f, protocol=4)

        np.save(os.path.join(cache_dir, cls.LOGITS_FILE),   all_logits)
        np.save(os.path.join(cache_dir, cls.FEATURES_FILE), all_features)

        total = time.time() - t0
        feat_gb = all_features.nbytes / 1e9
        print(f"  Done in {total/60:.1f} min | features: {feat_gb:.1f} GB")
        sys.stdout.flush()

        return cls(cache_dir)
