src/agiws_neural_quant/cache.py

"""Teacher cache — universal caching of teacher/model reference outputs.

Used by QAT, clustered quantization, and any teacher-student workflow.
All scripts import from this module — no duplication of cache logic.

Cache format (directory-based, streaming to disk during capture):
    {dataset_dir}/teacher_cache-{teacher_name}/{source_name}/
        __meta__.pt     — {"model_input": {...}, "model_output": {...}, "layer_names": [...]}
        blocks.0.norm1.pt — {"input": tensor_or_tuple, "output": tensor_or_tuple}
        blocks.0.attn.q_proj.pt
        ...

Streaming: during capture_with_hooks, each hook writes its layer's I/O to a
temporary directory IMMEDIATELY (torch.save + del tensor), so RAM/VRAM never
accumulates more than one layer's data at a time. After forward completes, the
temp directory is renamed to the final cache path.

Model-level I/O stored in __meta__.pt (small tensors — model input + output).
Per-layer I/O stored in individual {layer_name}.pt files.

Validation: cache is valid if:
  1. Directory exists with __meta__.pt
  2. cache_mtime >= source_mtime (not stale)
  3. All required layer names have corresponding .pt files

Capture granularity is configurable via CaptureConfig:
  - mode="final": only model input + model output (minimal disk)
  - mode="all": every selected module (all leaves, or recursive by depth)
  - mode="named": explicit module names
  - mode="types": modules matching isinstance(types)
  - depth: None=unlimited, 0=root, N=N levels deep
  - leaves_only: True=only terminal modules (no children)
"""

import os
import shutil
import tempfile
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import torch
import torch.nn as nn


# ---------------------------------------------------------------------------
# Path helpers
# ---------------------------------------------------------------------------

def get_cache_dir(dataset_dir: str | Path, teacher_name: str) -> Path:
    """Cache directory: {dataset_dir}/teacher_cache-{teacher_name}/"""
    return Path(dataset_dir) / f"teacher_cache-{teacher_name}"


def get_cache_path(source_path: str | Path, dataset_dir: str | Path, teacher_name: str) -> Path:
    """Cache directory path for a given source (screenshot/dataset file).

    Returns a DIRECTORY path (not a single .pt file). Inside:
        __meta__.pt + {layer_name}.pt per layer.
    """
    base = Path(source_path).stem
    return get_cache_dir(dataset_dir, teacher_name) / base


def _safe_layer_filename(name: str) -> str:
    """Convert module qualified name to a safe filename.

    Module names use dots (blocks.0.attn.q_proj) which are fine as filenames.
    """
    return name.replace("/", "__")


def _meta_path(cache_dir: Path) -> Path:
    """Path to __meta__.pt inside a cache directory."""
    return cache_dir / "__meta__.pt"


def _layer_path(cache_dir: Path, layer_name: str) -> Path:
    """Path to a single layer's .pt file inside a cache directory."""
    return cache_dir / f"{_safe_layer_filename(layer_name)}.pt"


# ---------------------------------------------------------------------------
# CaptureConfig — what to capture
# ---------------------------------------------------------------------------

@dataclass
class CaptureConfig:
    """Configuration for which modules to capture during teacher forward.

    mode:
        "final"  — only model input + model output (minimal disk)
        "all"    — all modules selected by depth/leaves_only
        "named"  — explicit module names in `names`
        "types"  — modules matching isinstance(types)

    names: list of qualified module names (for "named" mode)
    types: list of nn.Module subclasses (for "types" mode)
    depth: None=unlimited recursion, 0=root only, N=down to N levels
    leaves_only: if True, only modules with no children are captured
    """
    mode: str = "final"
    names: list[str] = field(default_factory=list)
    types: list[type] = field(default_factory=list)
    depth: int | None = None
    leaves_only: bool = False


# ---------------------------------------------------------------------------
# select_modules — build list of module names to capture
# ---------------------------------------------------------------------------

def select_modules(
    model: nn.Module,
    config: CaptureConfig,
) -> list[str]:
    """Build the list of qualified module names to capture based on config.

    Returns list of names like ["blocks.0.norm1", "blocks.0.attn.q_proj", ...].
    Model-level I/O (__model_input__/__model_output__) is always captured separately.
    """
    if config.mode == "final":
        return []

    if config.mode == "named":
        # Validate names exist in model
        available = dict(model.named_modules())
        result = []
        for name in config.names:
            if name in available:
                result.append(name)
            else:
                # Name not found — skip silently (or could raise)
                pass
        return result

    # "all" and "types" modes — walk module tree with depth/leaves filters
    result = []
    types_filter = config.types if config.mode == "types" else None

    for name, module in model.named_modules():
        # Skip the root module itself (empty name)
        if name == "":
            continue

        # Depth filter: count dots in name = depth level
        if config.depth is not None:
            level = name.count(".")
            if level > config.depth:
                continue

        # Leaves-only filter
        if config.leaves_only:
            has_children = len(list(module.children())) > 0
            if has_children:
                continue

        # Types filter
        if types_filter is not None:
            if not isinstance(module, tuple(types_filter)):
                continue

        result.append(name)

    return result


# ---------------------------------------------------------------------------
# capture_with_hooks — streaming to disk, single forward pass
# ---------------------------------------------------------------------------

def capture_with_hooks(
    model: nn.Module,
    inputs: dict[str, Any],
    config: CaptureConfig,
    output_extractor=None,
    tmp_dir: str | Path | None = None,
) -> dict:
    """Run a single forward pass through model, capturing I/O of selected modules.

    Streaming: each hook writes its layer's I/O to a temp directory IMMEDIATELY
    (torch.save + del tensor). After forward, the returned dict contains only
    __model_input__ and __model_output__ (small), plus a 'tmp_dir' key pointing
    to the directory with per-layer .pt files.

    Use save_capture_result() to move the temp dir to the final cache location,
    or pass the returned dict to save_layer_cache() which handles this.

    Args:
        model: the teacher model (will be set to eval mode, no_grad)
        inputs: dict of forward kwargs, e.g. {"pixel_values": pv, "grid_thw": gt}
        config: CaptureConfig specifying which modules to capture
        output_extractor: optional callable(model_output) -> dict of output parts
            If None, stores the raw output under __model_output__["raw"]
        tmp_dir: directory where per-layer .pt files are written during capture.
            If None, a temp dir is created (caller must move it to final location).

    Returns:
        dict with keys:
            "__model_input__": inputs (copy, CPU tensors)
            "__model_output__": extracted output parts (CPU tensors)
            "__layer_names__": list of captured layer names
            "__tmp_dir__": path to temp directory containing per-layer .pt files
    """
    module_names = select_modules(model, config)

    if tmp_dir is None:
        tmp_dir = Path(tempfile.mkdtemp(prefix="teacher_cache_"))
    else:
        tmp_dir = Path(tmp_dir)
        tmp_dir.mkdir(parents=True, exist_ok=True)

    captured_layer_names: list[str] = []
    hooks = []

    def _make_hook(name: str):
        def _hook(module, inp, out):
            # Detach + move to CPU WITHOUT .clone() — we save then immediately del.
            # torch.save serializes the tensor; after save the tensor is freed.
            inp_clean = tuple(t.detach().cpu() if isinstance(t, torch.Tensor) else t for t in inp)
            if isinstance(out, torch.Tensor):
                out_clean = out.detach().cpu()
            elif isinstance(out, (tuple, list)):
                out_clean = tuple(t.detach().cpu() if isinstance(t, torch.Tensor) else t for t in out)
            else:
                out_clean = out

            # Write immediately to disk — no accumulation in RAM
            layer_file = tmp_dir / f"{_safe_layer_filename(name)}.pt"
            torch.save({"input": inp_clean, "output": out_clean}, str(layer_file))

            # Free references — let GC reclaim memory
            del inp_clean, out_clean

            captured_layer_names.append(name)

        return _hook

    for name in module_names:
        module = model.get_submodule(name)
        h = module.register_forward_hook(_make_hook(name))
        hooks.append(h)

    try:
        model.eval()
        with torch.no_grad():
            output = model(**inputs)

        # Extract model output (small — kept in RAM)
        if output_extractor is not None:
            captured_output = output_extractor(output)
        else:
            if isinstance(output, torch.Tensor):
                captured_output = {"raw": output.detach().cpu()}
            else:
                captured_output = {"raw": output}

        # Model input — move to CPU (small tensors)
        captured_input = {
            k: v.detach().cpu() if isinstance(v, torch.Tensor) else v
            for k, v in inputs.items()
        }

        result = {
            "__model_input__": captured_input,
            "__model_output__": captured_output,
            "__layer_names__": captured_layer_names,
            "__tmp_dir__": str(tmp_dir),
        }
    finally:
        for h in hooks:
            h.remove()

    return result


# ---------------------------------------------------------------------------
# Save / load per-layer cache (directory-based)
# ---------------------------------------------------------------------------

def save_layer_cache(
    source_path: str | Path,
    dataset_dir: str | Path,
    teacher_name: str,
    data: dict,
) -> Path:
    """Save capture result to a cache DIRECTORY on disk.

    If data contains '__tmp_dir__', per-layer .pt files are moved from the temp
    directory into the final cache directory. __model_input__ / __model_output__
    / __layer_names__ are saved to __meta__.pt.

    If data does NOT contain '__tmp_dir__' (legacy/old format with inline tensors),
    each layer entry is saved as a separate .pt file, and model I/O to __meta__.pt.

    Returns the cache directory path.
    """
    cache_dir = get_cache_path(source_path, dataset_dir, teacher_name)

    # Remove old cache if exists
    if cache_dir.exists():
        shutil.rmtree(cache_dir)
    cache_dir.mkdir(parents=True, exist_ok=True)

    # Save model-level I/O + layer name list to __meta__.pt
    meta = {
        "__model_input__": data.get("__model_input__", {}),
        "__model_output__": data.get("__model_output__", {}),
        "__layer_names__": data.get("__layer_names__", []),
    }
    torch.save(meta, str(_meta_path(cache_dir)))

    # Handle per-layer files
    if "__tmp_dir__" in data:
        # Streaming capture — move .pt files from temp dir to cache dir
        tmp_dir = Path(data["__tmp_dir__"])
        if tmp_dir.exists():
            for pt_file in tmp_dir.glob("*.pt"):
                shutil.copy2(str(pt_file), str(cache_dir / pt_file.name))
            # Clean up temp dir
            shutil.rmtree(tmp_dir, ignore_errors=True)
    else:
        # Legacy format — layer data is inline in dict
        for key in data:
            if key.startswith("__"):
                continue
            layer_data = data[key]
            torch.save(layer_data, str(_layer_path(cache_dir, key)))

    return cache_dir


def load_layer_cache(
    source_path: str | Path,
    dataset_dir: str | Path,
    teacher_name: str,
) -> dict:
    """Load full per-layer cache from disk into a single dict.

    Reads __meta__.pt (model I/O) and all per-layer .pt files, merging them
    into a flat dict:
        {
            "__model_input__": {...},
            "__model_output__": {...},
            "blocks.0.norm1": {"input": tensor, "output": tensor},
            ...
        }

    Note: this loads ALL layers into RAM. For large caches, prefer load_layer_io()
    to read a single layer at a time.
    """
    cache_dir = get_cache_path(source_path, dataset_dir, teacher_name)

    # Load meta
    meta = torch.load(str(_meta_path(cache_dir)), weights_only=False)

    result = {
        "__model_input__": meta.get("__model_input__", {}),
        "__model_output__": meta.get("__model_output__", {}),
    }

    # Load all per-layer .pt files (excluding __meta__.pt)
    layer_names = meta.get("__layer_names__", [])
    if not layer_names:
        # Fallback: discover layer files by listing directory
        layer_names = [
            f.stem for f in cache_dir.glob("*.pt")
            if f.name != "__meta__.pt"
        ]

    for name in layer_names:
        layer_file = _layer_path(cache_dir, name)
        if layer_file.exists():
            result[name] = torch.load(str(layer_file), weights_only=False)

    return result


def load_layer_io(
    source_path: str | Path,
    dataset_dir: str | Path,
    teacher_name: str,
    layer_name: str,
) -> tuple:
    """Load input and output of a single layer from cache.

    Reads ONLY the requested layer's .pt file — no full cache load.

    Returns: (input, output) where each may be a tensor, tuple, or other.
    """
    cache_dir = get_cache_path(source_path, dataset_dir, teacher_name)
    layer_file = _layer_path(cache_dir, layer_name)
    if not layer_file.exists():
        # Try listing available layers from meta
        meta = torch.load(str(_meta_path(cache_dir)), weights_only=False)
        available = meta.get("__layer_names__", [])
        raise KeyError(f"Layer '{layer_name}' not in cache. Available: {available}")
    entry = torch.load(str(layer_file), weights_only=False)
    return entry["input"], entry["output"]


def list_cached_names(data: dict) -> list[str]:
    """List all layer names in a loaded cache dict (excluding __model_*__).

    Works with both loaded dicts (from load_layer_cache) and meta dicts.
    """
    keys = [k for k in data.keys() if not k.startswith("__")]
    # If data is a meta dict (has __layer_names__), use that
    if "__layer_names__" in data:
        keys = data["__layer_names__"]
    return keys


# ---------------------------------------------------------------------------
# Validation — mtime + contents check (directory-based)
# ---------------------------------------------------------------------------

def validate_cache_contents(
    source_path: str | Path,
    dataset_dir: str | Path,
    teacher_name: str,
    required_names: list[str],
) -> list[str]:
    """Check if cache directory has all required layer .pt files.

    Returns list of MISSING names (empty = all present = valid).
    Does NOT check mtime — use is_cache_valid for full validation.
    """
    cache_dir = get_cache_path(source_path, dataset_dir, teacher_name)
    if not _meta_path(cache_dir).exists():
        return list(required_names)

    # Check each required name has a .pt file
    missing = []
    for name in required_names:
        if not _layer_path(cache_dir, name).exists():
            missing.append(name)
    return missing


def is_cache_valid(
    source_path: str | Path,
    dataset_dir: str | Path,
    teacher_name: str,
    required_names: list[str] | None = None,
) -> bool:
    """Full validation: directory exists with __meta__.pt, mtime fresh, contents present.

    Args:
        required_names: list of layer names that must have .pt files.
            If None, only mtime+existence is checked (backward compat).
    """
    cache_dir = get_cache_path(source_path, dataset_dir, teacher_name)
    if not _meta_path(cache_dir).exists():
        return False
    # mtime of cache dir (or __meta__.pt) must be >= source mtime
    cache_mtime = cache_dir.stat().st_mtime
    if cache_mtime < Path(source_path).stat().st_mtime:
        return False
    if required_names is not None:
        missing = validate_cache_contents(source_path, dataset_dir, teacher_name, required_names)
        if missing:
            return False
    return True


def get_sources_needing_cache(
    source_paths: list[str],
    dataset_dir: str | Path,
    teacher_name: str,
    required_names: list[str] | None = None,
) -> list[str]:
    """Filter source paths that need (re)caching: missing, stale, or incomplete."""
    return [
        s for s in source_paths
        if not is_cache_valid(s, dataset_dir, teacher_name, required_names)
    ]


# ---------------------------------------------------------------------------
# Cleanup — remove old-format cache files
# ---------------------------------------------------------------------------

def clean_old_cache(dataset_dir: str | Path, teacher_name: str) -> int:
    """Remove old-format cache files (single .pt files instead of directories).

    Old format: {source_name}.pt (single file with all layers inline)
    New format: {source_name}/ directory with per-layer .pt files

    Also removes corrupted cache directories (missing __meta__.pt).

    Returns count of removed items (files + directories).
    """
    cache_dir = get_cache_dir(dataset_dir, teacher_name)
    if not cache_dir.exists():
        return 0

    removed = 0
    for item in cache_dir.iterdir():
        if item.is_file() and item.suffix == ".pt":
            # Old format — single .pt file (new format uses directories)
            try:
                data = torch.load(str(item), weights_only=False)
                if isinstance(data, dict) and "__model_input__" not in data:
                    # Old format — remove
                    item.unlink()
                    removed += 1
            except (RuntimeError, EOFError, Exception):
                # Corrupted — remove
                item.unlink()
                removed += 1
        elif item.is_dir():
            # New format directory — check if it has __meta__.pt
            if not _meta_path(item).exists():
                # Corrupted/incomplete — remove
                shutil.rmtree(item, ignore_errors=True)
                removed += 1
    return removed


# ---------------------------------------------------------------------------
# TeacherCache — OOP wrapper, eliminates dataset_dir/teacher_name duplication
# ---------------------------------------------------------------------------

class TeacherCache:
    """Bound cache for a specific teacher model + dataset directory.

    Encapsulates dataset_dir and teacher_name so scripts create ONE instance
    and call methods without passing these every time. All read/write/validation
    goes through this class — no duplication in scripts.

    Usage:
        cache = TeacherCache(dataset_dir, "qwen3_vl_vision_fp32")
        cache.clean_old()
        need = cache.get_sources_needing_cache(screenshots, required_names)
        cache.save(source_path, data)
        data = cache.load(source_path)
        inp, out = cache.load_layer_io(source_path, "blocks.0.norm")
    """

    def __init__(self, dataset_dir: str | Path, teacher_name: str):
        self.dataset_dir = Path(dataset_dir)
        self.teacher_name = teacher_name
        self.cache_dir = get_cache_dir(dataset_dir, teacher_name)

    def get_path(self, source_path: str | Path) -> Path:
        """Cache directory path for a source file (directory, not .pt file)."""
        return get_cache_path(source_path, self.dataset_dir, self.teacher_name)

    def is_valid(
        self,
        source_path: str | Path,
        required_names: list[str] | None = None,
    ) -> bool:
        """Full validation: exists, mtime fresh, contents present."""
        return is_cache_valid(
            source_path, self.dataset_dir, self.teacher_name, required_names
        )

    def validate_contents(
        self,
        source_path: str | Path,
        required_names: list[str],
    ) -> list[str]:
        """Returns list of MISSING layer names (empty = all present)."""
        return validate_cache_contents(
            source_path, self.dataset_dir, self.teacher_name, required_names
        )

    def get_sources_needing_cache(
        self,
        source_paths: list[str],
        required_names: list[str] | None = None,
    ) -> list[str]:
        """Filter sources that need (re)caching."""
        return get_sources_needing_cache(
            source_paths, self.dataset_dir, self.teacher_name, required_names
        )

    def save(self, source_path: str | Path, data: dict) -> Path:
        """Save capture result to cache directory. Returns cache directory path.

        If data has '__tmp_dir__', per-layer .pt files are moved from temp dir.
        Otherwise, layer entries in data are saved as individual .pt files.
        """
        return save_layer_cache(
            source_path, self.dataset_dir, self.teacher_name, data
        )

    def load(self, source_path: str | Path) -> dict:
        """Load full per-layer cache dict from disk (ALL layers into RAM).

        For large caches, prefer load_layer_io() to read one layer at a time.
        """
        return load_layer_cache(
            source_path, self.dataset_dir, self.teacher_name
        )

    def load_layer_io(
        self,
        source_path: str | Path,
        layer_name: str,
    ) -> tuple:
        """Load (input, output) of a single layer from cache (lazy, one file)."""
        return load_layer_io(
            source_path, self.dataset_dir, self.teacher_name, layer_name
        )

    def list_names(self, data: dict) -> list[str]:
        """List layer names in a loaded cache dict or meta dict (excluding __model_*__)."""
        return list_cached_names(data)

    def clean_old(self) -> int:
        """Remove old-format cache files/directories. Returns count removed."""
        return clean_old_cache(self.dataset_dir, self.teacher_name)

    def capture(
        self,
        model: nn.Module,
        inputs: dict[str, Any],
        config: CaptureConfig,
        output_extractor=None,
        tmp_dir: str | Path | None = None,
    ) -> dict:
        """Run forward pass with hooks, capture I/O to temp dir. Does NOT save.

        Returns dict with __model_input__, __model_output__, __layer_names__,
        __tmp_dir__. Use save() to persist to final cache location.
        """
        return capture_with_hooks(model, inputs, config, output_extractor, tmp_dir=tmp_dir)

    def capture_and_save(
        self,
        source_path: str | Path,
        model: nn.Module,
        inputs: dict[str, Any],
        config: CaptureConfig,
        output_extractor=None,
        tmp_dir: str | Path | None = None,
    ) -> Path:
        """Capture I/O via hooks (streaming to temp dir) AND save to final cache.

        Hooks write each layer to disk immediately (no RAM accumulation).
        After forward, temp dir is moved to final cache location.

        Returns cache directory path.
        """
        data = capture_with_hooks(model, inputs, config, output_extractor, tmp_dir=tmp_dir)
        return self.save(source_path, data)


# ---------------------------------------------------------------------------
# Backward-compat wrappers (old API, delegates to new directory-based cache)
# ---------------------------------------------------------------------------

def save_cache(
    source_path: str | Path,
    dataset_dir: str | Path,
    teacher_name: str,
    pixel_values,
    grid_thw,
    teacher_parts: list[torch.Tensor],
) -> Path:
    """Backward-compat: save model-level I/O only (final mode).

    Constructs a minimal cache dict with __model_input__ and __model_output__,
    then saves as directory-based cache (meta only, no per-layer files).
    """
    data = {
        "__model_input__": {"pixel_values": pixel_values, "grid_thw": grid_thw},
        "__model_output__": {"teacher_parts": teacher_parts},
        "__layer_names__": [],
    }
    return save_layer_cache(source_path, dataset_dir, teacher_name, data)


def load_cache(
    source_path: str | Path,
    dataset_dir: str | Path,
    teacher_name: str,
) -> tuple:
    """Backward-compat: load model-level I/O (final mode).

    Returns: (pixel_values, grid_thw, teacher_parts)
    """
    data = load_layer_cache(source_path, dataset_dir, teacher_name)
    mi = data["__model_input__"]
    mo = data["__model_output__"]
    return mi["pixel_values"], mi["grid_thw"], mo["teacher_parts"]