src/agiws_neural_quant/analysis.py

"""LayerAnalyzer — per-layer cosine analysis with error accumulation tracking.

Analyzes how quantization error accumulates across layers of a model.
For each layer: runs quantized layer on cached teacher input, compares
output to cached teacher output, computes cosine similarity.

Detects:
  - Explosion points: layers where cosine drops sharply (> threshold)
  - Cascade zones: consecutive layers with monotonic cosine decline
  - Per-layer sensitivity: which layers lose most accuracy

Writes reports in JSON (programmatic) and Markdown (human-readable).
"""

import json
import math
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import torch
import torch.nn as nn


@dataclass
class LayerResult:
    """Per-layer analysis result."""
    layer_name: str
    cosine: float
    input_shape: str
    output_shape: str
    layer_type: str
    error: Optional[str] = None  # if layer execution failed


@dataclass
class AnomalyReport:
    """Detected anomalies in error accumulation."""
    explosion_points: List[Dict[str, Any]] = field(default_factory=list)
    cascade_zones: List[Dict[str, Any]] = field(default_factory=list)
    worst_layer: Optional[Dict[str, Any]] = None
    best_layer: Optional[Dict[str, Any]] = None


@dataclass
class SplitReport:
    """Full report for one quantization split."""
    split_label: str
    value_bits: int
    cluster_id_bits: int
    B: int
    K: int
    full_model_cosine: float
    per_layer: List[LayerResult] = field(default_factory=list)
    anomalies: Optional[AnomalyReport] = None
    quant_time: float = 0.0
    analysis_time: float = 0.0


class LayerAnalyzer:
    """Analyzes per-layer quantization error accumulation.

    Two modes:
      1. Dict mode (legacy): pass teacher_cache dict (all layers in RAM).
         Use for tests / small models.
      2. Lazy mode: pass TeacherCache object + source_path.
         Loads one layer at a time from disk — low RAM footprint.
         Use for large models on Colab.

    Usage:
        # Lazy (recommended for Colab):
        analyzer = LayerAnalyzer(cache=teacher_cache_obj, source_path=img_path)
        # Dict (legacy):
        analyzer = LayerAnalyzer(teacher_cache=cache_dict)
    """

    def __init__(
        self,
        teacher_cache: Optional[Dict[str, Any]] = None,
        compute_dtype: str = "fp32",
        explosion_threshold: float = 0.1,
        cascade_min_length: int = 3,
        cache=None,
        source_path: Optional[str] = None,
    ):
        """
        Args:
            teacher_cache: loaded cache dict (dict mode). Can be None if
                           using lazy mode (cache + source_path).
            compute_dtype: "fp32" or "fp16"
            explosion_threshold: cosine drop > this = explosion point
            cascade_min_length: min consecutive declining layers for cascade zone
            cache: TeacherCache object (lazy mode)
            source_path: source file path (lazy mode, passed to cache.load_layer_io)
        """
        self.teacher_cache = teacher_cache
        self.compute_dtype = compute_dtype
        self.explosion_threshold = explosion_threshold
        self.cascade_min_length = cascade_min_length
        # Lazy mode
        self._lazy_cache = cache
        self._lazy_source_path = source_path
        self._lazy_layer_names: Optional[List[str]] = None
        self._lazy_model_input: Optional[Dict] = None
        self._lazy_model_output: Optional[Dict] = None

        if cache is not None and source_path is not None:
            # Lazy mode: preload only meta (model I/O + layer name list)
            self._lazy_init()

    def _lazy_init(self):
        """In lazy mode, load only __meta__.pt (small) to get layer names + model I/O."""
        from agiws_neural_quant.cache import _meta_path, _layer_path
        cache_dir = self._lazy_cache.get_path(self._lazy_source_path)
        meta = torch.load(str(_meta_path(cache_dir)), weights_only=False)
        self._lazy_layer_names = [
            n for n in meta.get("__layer_names__", [])
            if not n.startswith("__")
        ]
        # If no layer names in meta, discover from directory
        if not self._lazy_layer_names:
            self._lazy_layer_names = [
                f.stem for f in cache_dir.glob("*.pt")
                if f.name != "__meta__.pt"
            ]
        self._lazy_model_input = meta.get("__model_input__", {})
        self._lazy_model_output = meta.get("__model_output__", {})

    def _lazy_get_layer(self, layer_name: str) -> Optional[Dict]:
        """In lazy mode, load one layer from disk. Returns {'input':..., 'output':...}."""
        if self._lazy_cache is None:
            return None
        try:
            inp, out = self._lazy_cache.load_layer_io(self._lazy_source_path, layer_name)
            return {"input": inp, "output": out}
        except (KeyError, FileNotFoundError):
            return None

    def _get_layer_entry(self, layer_name: str) -> Optional[Dict]:
        """Get layer entry from cache — lazy or dict mode."""
        if self._lazy_cache is not None:
            return self._lazy_get_layer(layer_name)
        if self.teacher_cache is not None:
            return self.teacher_cache.get(layer_name)
        return None

    def _get_model_input(self) -> Optional[Dict]:
        """Get model input — lazy or dict mode."""
        if self._lazy_model_input is not None:
            return self._lazy_model_input
        if self.teacher_cache is not None:
            return self.teacher_cache.get("__model_input__")
        return None

    def _get_model_output(self) -> Optional[Dict]:
        """Get model output — lazy or dict mode."""
        if self._lazy_model_output is not None:
            return self._lazy_model_output
        if self.teacher_cache is not None:
            return self.teacher_cache.get("__model_output__")
        return None

    def _get_layer_names(self) -> List[str]:
        """Get list of layer names — lazy or dict mode."""
        if self._lazy_layer_names is not None:
            return self._lazy_layer_names
        if self.teacher_cache is not None:
            return [k for k in self.teacher_cache.keys() if not k.startswith("__")]
        return []

    def analyze_split(
        self,
        quantized_model: nn.Module,
        split_label: str,
        value_bits: int,
        cluster_id_bits: int,
        quant_time: float = 0.0,
        max_layers: Optional[int] = None,
    ) -> SplitReport:
        """Analyze one quantization split: per-layer cosine + anomaly detection.

        Args:
            quantized_model: model already quantized with this split
            split_label: human-readable label (e.g. "4v+0c")
            value_bits, cluster_id_bits: split parameters
            quant_time: time spent on quantization (for report)
            max_layers: limit number of layers to analyze (None = all)

        Returns: SplitReport with per-layer cosines and anomalies
        """
        B = value_bits + cluster_id_bits
        K = 1 << cluster_id_bits
        report = SplitReport(
            split_label=split_label,
            value_bits=value_bits,
            cluster_id_bits=cluster_id_bits,
            B=B,
            K=K,
            full_model_cosine=0.0,
            quant_time=quant_time,
        )

        t0 = time.time()

        # Full-model cosine (using cached model output)
        full_cos = self._compute_full_model_cosine(quantized_model)
        report.full_model_cosine = full_cos

        # Per-layer cosine
        layer_names = self._get_layer_names()
        if max_layers is not None:
            layer_names = layer_names[:max_layers]

        for layer_name in layer_names:
            lr = self._analyze_single_layer(quantized_model, layer_name)
            report.per_layer.append(lr)

        # Anomaly detection
        report.anomalies = self._detect_anomalies(report.per_layer)
        report.analysis_time = time.time() - t0

        return report

    def _compute_full_model_cosine(self, model: nn.Module) -> float:
        """Compute full-model cosine vs cached teacher output."""
        mi = self._get_model_input()
        mo = self._get_model_output()

        if not mi or not mo or "pooler_output" not in mo:
            return 0.0

        pv = mi.get("pixel_values") or mi.get("hidden_states")
        gt = mi.get("grid_thw")
        if pv is None or gt is None:
            return 0.0

        # Determine model device
        try:
            dev = next(model.parameters()).device
        except StopIteration:
            dev = torch.device("cpu")
        pv = pv.to(dev)
        gt = gt.to(dev)

        model.eval()
        # Force all buffers to model device (.to() may miss lazy buffers)
        try:
            mdev = next(model.parameters()).device
        except StopIteration:
            mdev = torch.device("cpu")
        for b in model.buffers():
            b.data = b.data.to(mdev)

        with torch.no_grad():
            out = model(pv, grid_thw=gt) if "pixel_values" in mi else model(hidden_states=pv, grid_thw=gt)

        if not hasattr(out, "pooler_output"):
            return 0.0

        ref = mo["pooler_output"].float().flatten()
        test = out.pooler_output.float().flatten()
        # Ensure both on same device (ref from cache=CPU, test from model=GPU)
        ref = ref.to(test.device)
        return torch.nn.functional.cosine_similarity(
            ref.unsqueeze(0), test.unsqueeze(0)
        ).item()

    def _analyze_single_layer(
        self,
        model: nn.Module,
        layer_name: str,
    ) -> LayerResult:
        """Analyze one layer: run quantized layer on cached input, compare output."""
        entry = self._get_layer_entry(layer_name)
        if entry is None:
            return LayerResult(
                layer_name=layer_name,
                cosine=0.0,
                input_shape="N/A",
                output_shape="N/A",
                layer_type="unknown",
                error="not in cache",
            )

        cached_inp = entry.get("input")
        cached_out = entry.get("output")
        if cached_inp is None or cached_out is None:
            return LayerResult(
                layer_name=layer_name,
                cosine=0.0,
                input_shape="N/A",
                output_shape="N/A",
                layer_type="unknown",
                error="cache entry missing input/output",
            )

        # Get quantized module
        try:
            q_module = model.get_submodule(layer_name)
        except Exception as e:
            return LayerResult(
                layer_name=layer_name,
                cosine=0.0,
                input_shape="N/A",
                output_shape="N/A",
                layer_type="missing",
                error=f"get_submodule failed: {e}",
            )

        layer_type = type(q_module).__name__

        # Determine device of the quantized module
        try:
            dev = next(q_module.parameters()).device
        except StopIteration:
            dev = torch.device("cpu")

        # Move cached input to module device
        def _to_dev(x):
            if isinstance(x, torch.Tensor):
                return x.to(dev)
            return x

        if isinstance(cached_inp, (tuple, list)):
            cached_inp_dev = tuple(_to_dev(t) for t in cached_inp)
        else:
            cached_inp_dev = _to_dev(cached_inp)

        # Run quantized layer on cached input
        try:
            with torch.no_grad():
                if isinstance(cached_inp_dev, (tuple, list)) and len(cached_inp_dev) > 0:
                    q_out = q_module(*cached_inp_dev)
                else:
                    q_out = q_module(cached_inp_dev)
        except Exception as e:
            in_shape = "N/A"
            if isinstance(cached_inp, (tuple, list)) and len(cached_inp) > 0:
                in_shape = str(getattr(cached_inp[0], "shape", "N/A"))
            return LayerResult(
                layer_name=layer_name,
                cosine=0.0,
                input_shape=in_shape,
                output_shape="N/A",
                layer_type=layer_type,
                error=f"forward failed: {type(e).__name__}: {e}",
            )

        # Cosine comparison (both on same device)
        cos = 0.0
        if isinstance(cached_out, torch.Tensor) and isinstance(q_out, torch.Tensor):
            ref = cached_out.float().to(dev).flatten()
            test = q_out.float().flatten()
            if ref.numel() > 0 and test.numel() > 0:
                c = torch.nn.functional.cosine_similarity(
                    ref.unsqueeze(0), test.unsqueeze(0)
                ).item()
                # Guard against NaN/Inf (zero vectors → 0/0 = NaN)
                if not (math.isnan(c) or math.isinf(c)):
                    cos = c

        in_shape = "N/A"
        if isinstance(cached_inp, (tuple, list)) and len(cached_inp) > 0:
            in_shape = str(getattr(cached_inp[0], "shape", "N/A"))
        out_shape = str(getattr(q_out, "shape", "N/A"))

        return LayerResult(
            layer_name=layer_name,
            cosine=cos,
            input_shape=in_shape,
            output_shape=out_shape,
            layer_type=layer_type,
        )

    def _detect_anomalies(self, layer_results: List[LayerResult]) -> AnomalyReport:
        """Detect explosion points and cascade zones in error accumulation.

        Explosion points and cascade zones are computed between REAL adjacent
        layers (by index in layer_results), skipping error-layers. An error-layer
        does NOT create a false explosion between its neighbours — it breaks
        adjacency (neighbours across an error are not compared).
        """
        report = AnomalyReport()

        # Build list of (index, result) for valid layers — preserve original positions
        indexed_valid = [
            (i, lr) for i, lr in enumerate(layer_results) if lr.error is None
        ]
        if not indexed_valid:
            return report

        valid = [lr for _, lr in indexed_valid]
        positions = [idx for idx, _ in indexed_valid]

        # Worst and best
        worst = min(valid, key=lambda x: x.cosine)
        best = max(valid, key=lambda x: x.cosine)
        report.worst_layer = {"name": worst.layer_name, "cosine": worst.cosine}
        report.best_layer = {"name": best.layer_name, "cosine": best.cosine}

        # Explosion points: sharp cosine drop between REAL adjacent layers
        # (positions must be consecutive: positions[i] == positions[i-1]+1)
        for i in range(1, len(valid)):
            if positions[i] != positions[i - 1] + 1:
                continue  # not real neighbours (error-layer between them)
            drop = valid[i - 1].cosine - valid[i].cosine
            if drop > self.explosion_threshold:
                report.explosion_points.append({
                    "layer": valid[i].layer_name,
                    "prev_cosine": valid[i - 1].cosine,
                    "cosine": valid[i].cosine,
                    "drop": drop,
                })

        # Cascade zones: consecutive real-adjacent layers with monotonic decline
        zone_start = None
        for i in range(1, len(valid)):
            is_real_adjacent = positions[i] == positions[i - 1] + 1
            if is_real_adjacent and valid[i].cosine < valid[i - 1].cosine:
                if zone_start is None:
                    zone_start = i - 1
            else:
                if zone_start is not None and (i - zone_start) >= self.cascade_min_length:
                    report.cascade_zones.append({
                        "start": valid[zone_start].layer_name,
                        "end": valid[i - 1].layer_name,
                        "length": i - zone_start,
                        "start_cosine": valid[zone_start].cosine,
                        "end_cosine": valid[i - 1].cosine,
                        "total_drop": valid[zone_start].cosine - valid[i - 1].cosine,
                    })
                zone_start = None
        # Check trailing zone
        if zone_start is not None and (len(valid) - zone_start) >= self.cascade_min_length:
            report.cascade_zones.append({
                "start": valid[zone_start].layer_name,
                "end": valid[-1].layer_name,
                "length": len(valid) - zone_start,
                "start_cosine": valid[zone_start].cosine,
                "end_cosine": valid[-1].cosine,
                "total_drop": valid[zone_start].cosine - valid[-1].cosine,
            })

        return report

    # ---- Report writing ----

    @staticmethod
    def write_report(
        reports: List[SplitReport],
        output_path: str | Path,
        fmt: str = "markdown",
    ):
        """Write analysis report to file.

        Args:
            reports: list of SplitReport (one per quantization split)
            output_path: file path
            fmt: "markdown" or "json"
        """
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        if fmt == "json":
            LayerAnalyzer._write_json(reports, output_path)
        elif fmt == "markdown":
            LayerAnalyzer._write_markdown(reports, output_path)
        else:
            raise ValueError(f"Unknown format: {fmt}")

    @staticmethod
    def _write_json(reports: List[SplitReport], path: Path):
        """Write JSON report (programmatic analysis)."""
        data = {
            "report_type": "layer_analysis",
            "timestamp": time.time(),
            "splits": [],
        }
        for r in reports:
            split_data = {
                "split_label": r.split_label,
                "value_bits": r.value_bits,
                "cluster_id_bits": r.cluster_id_bits,
                "B": r.B,
                "K": r.K,
                "full_model_cosine": r.full_model_cosine,
                "quant_time": r.quant_time,
                "analysis_time": r.analysis_time,
                "per_layer": [
                    {
                        "layer_name": lr.layer_name,
                        "cosine": lr.cosine,
                        "layer_type": lr.layer_type,
                        "error": lr.error,
                    }
                    for lr in r.per_layer
                ],
                "anomalies": {
                    "explosion_points": r.anomalies.explosion_points if r.anomalies else [],
                    "cascade_zones": r.anomalies.cascade_zones if r.anomalies else [],
                    "worst_layer": r.anomalies.worst_layer if r.anomalies else None,
                    "best_layer": r.anomalies.best_layer if r.anomalies else None,
                },
            }
            data["splits"].append(split_data)

        with open(path, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2, ensure_ascii=False)

    @staticmethod
    def _write_markdown(reports: List[SplitReport], path: Path):
        """Write Markdown report (human-readable)."""
        lines = []
        lines.append("# Per-Layer Quantization Analysis Report")
        lines.append("")
        lines.append(f"Generated: {time.strftime('%Y-%m-%d %H:%M:%S')}")
        lines.append("")

        # Summary table
        lines.append("## Summary")
        lines.append("")
        lines.append("| Split | B | K | Full-model cosine | Layers analyzed | Worst layer | Best layer |")
        lines.append("|-------|---|---|------------------|-----------------|-------------|------------|")
        for r in reports:
            worst = r.anomalies.worst_layer if r.anomalies and r.anomalies.worst_layer else {"name": "N/A", "cosine": 0}
            best = r.anomalies.best_layer if r.anomalies and r.anomalies.best_layer else {"name": "N/A", "cosine": 0}
            n_valid = len([lr for lr in r.per_layer if lr.error is None])
            lines.append(
                f"| {r.split_label} | {r.B} | {r.K} | {r.full_model_cosine:.6f} | "
                f"{n_valid} | {worst['name']} ({worst['cosine']:.4f}) | "
                f"{best['name']} ({best['cosine']:.4f}) |"
            )
        lines.append("")

        # Per-split details
        for r in reports:
            lines.append(f"## {r.split_label} (B={r.B}, K={r.K})")
            lines.append("")
            lines.append(f"Full-model cosine: {r.full_model_cosine:.6f}")
            lines.append(f"Quant time: {r.quant_time:.1f}s, Analysis time: {r.analysis_time:.1f}s")
            lines.append("")

            # Anomalies
            if r.anomalies:
                if r.anomalies.explosion_points:
                    lines.append("### Explosion Points (sharp cosine drops)")
                    lines.append("")
                    for ep in r.anomalies.explosion_points:
                        lines.append(
                            f"- **{ep['layer']}**: {ep['prev_cosine']:.4f} -> {ep['cosine']:.4f} "
                            f"(drop {ep['drop']:.4f})"
                        )
                    lines.append("")

                if r.anomalies.cascade_zones:
                    lines.append("### Cascade Zones (monotonic decline)")
                    lines.append("")
                    for cz in r.anomalies.cascade_zones:
                        lines.append(
                            f"- **{cz['start']} -> {cz['end']}** ({cz['length']} layers): "
                            f"{cz['start_cosine']:.4f} -> {cz['end_cosine']:.4f} "
                            f"(total drop {cz['total_drop']:.4f})"
                        )
                    lines.append("")

            # Per-layer cosine table
            lines.append("### Per-Layer Cosine")
            lines.append("")
            lines.append("| Layer | Type | Cosine | Error |")
            lines.append("|-------|------|--------|-------|")
            for lr in r.per_layer:
                err = lr.error or ""
                lines.append(
                    f"| {lr.layer_name} | {lr.layer_type} | {lr.cosine:.6f} | {err} |"
                )
            lines.append("")

        with open(path, "w", encoding="utf-8") as f:
            f.write("\n".join(lines))