#!/usr/bin/env python3
"""Analyze checkpoint weights for collapse and structure.

Quick sanity checks without running full inference:
1. Weight statistics (mean, std, range) - collapsed models have low std
2. GNN layer presence and activation
3. Output layer bias distribution - should show LLM differentiation
"""

import sys
from pathlib import Path

import numpy as np
import torch

_SCRIPT_DIR = Path(__file__).resolve().parent
_MODULE_DIR = _SCRIPT_DIR.parent


def analyze_checkpoint(name: str, model_path: Path):
    """Analyze checkpoint weights."""
    print(f"\n{'='*70}")
    print(f"CHECKPOINT: {name}")
    print(f"{'='*70}")

    if not model_path.exists():
        print(f"ERROR: File not found: {model_path}")
        return None

    checkpoint = torch.load(model_path, map_location="cpu", weights_only=False)

    # Get all weight keys
    keys = list(checkpoint.keys())
    print(f"Total layers: {len(keys)}")

    # Categorize layers
    gnn_layers = [k for k in keys if "gnn" in k.lower() or "gat" in k.lower() or "gcn" in k.lower()]
    moe_layers = [k for k in keys if "moe" in k.lower() or "expert" in k.lower() or "gate" in k.lower()]
    output_layers = [k for k in keys if "output" in k.lower() or "quality" in k.lower() or "head" in k.lower()]

    print(f"\nLayer breakdown:")
    print(f"  GNN layers: {len(gnn_layers)}")
    print(f"  MoE layers: {len(moe_layers)}")
    print(f"  Output layers: {len(output_layers)}")

    # === CHECK 1: GNN presence ===
    has_gnn = len(gnn_layers) > 0
    print(f"\n[CHECK] GNN layers present: {'YES' if has_gnn else 'NO'}")
    if gnn_layers:
        print(f"  GNN keys: {gnn_layers[:5]}...")

    # === CHECK 2: MoE presence ===
    has_moe = len(moe_layers) > 0
    print(f"[CHECK] MoE layers present: {'YES' if has_moe else 'NO'}")
    if moe_layers:
        print(f"  MoE keys: {moe_layers[:5]}...")

    # === CHECK 3: Weight statistics ===
    print(f"\n[ANALYSIS] Weight statistics:")

    all_weights = []
    for k, v in checkpoint.items():
        if isinstance(v, torch.Tensor) and v.numel() > 0:
            all_weights.extend(v.flatten().tolist())

    all_weights = np.array(all_weights)
    print(f"  Total parameters: {len(all_weights):,}")
    print(f"  Mean: {all_weights.mean():.6f}")
    print(f"  Std:  {all_weights.std():.6f}")
    print(f"  Min:  {all_weights.min():.6f}")
    print(f"  Max:  {all_weights.max():.6f}")

    # === CHECK 4: Output layer analysis ===
    print(f"\n[ANALYSIS] Output/quality layers:")
    for k in keys:
        if "quality" in k.lower() or "output" in k.lower():
            v = checkpoint[k]
            if isinstance(v, torch.Tensor):
                print(f"  {k}: shape={list(v.shape)}, mean={v.mean():.4f}, std={v.std():.4f}")

    # === CHECK 5: Check for collapsed weights (low variance) ===
    weight_std = all_weights.std()
    is_collapsed = weight_std < 0.01
    print(f"\n[CHECK] Weight collapse: {'COLLAPSED (std < 0.01)' if is_collapsed else 'HEALTHY'}")

    # === CHECK 6: Analyze specific key layers ===
    print(f"\n[ANALYSIS] Key layer inspection:")
    inspect_patterns = ["llm_transform", "query_transform", "gnn", "quality"]
    for pattern in inspect_patterns:
        matching = [k for k in keys if pattern in k.lower()][:2]
        for k in matching:
            v = checkpoint[k]
            if isinstance(v, torch.Tensor) and v.numel() > 10:
                print(f"  {k[:60]}...")
                print(f"    shape={list(v.shape)}, mean={v.mean():.4f}, std={v.std():.4f}")

    return {
        "name": name,
        "has_gnn": has_gnn,
        "has_moe": has_moe,
        "n_params": len(all_weights),
        "weight_std": weight_std,
        "is_collapsed": is_collapsed,
    }


def main():
    checkpoints = [
        ("GAT 4L + 23x BBEH", _MODULE_DIR / "model_path" / "optimal_gat_bbeh_23x_20251213.pth"),
        ("MLP + MoE4 (no GNN)", _MODULE_DIR / "model_path" / "canonical_mlp_moe4_top1_20251212.pth"),
        ("GNN Router Canonical", _MODULE_DIR / "model_path" / "gnn_router_canonical.pth"),
    ]

    results = []
    for name, path in checkpoints:
        result = analyze_checkpoint(name, path)
        if result:
            results.append(result)

    # Summary
    print(f"\n{'='*70}")
    print("SUMMARY COMPARISON")
    print(f"{'='*70}")
    print(f"{'Checkpoint':<30} {'GNN':<6} {'MoE':<6} {'Params':<12} {'Std':<10} {'Status':<10}")
    print("-" * 70)
    for r in results:
        status = "COLLAPSED" if r["is_collapsed"] else "OK"
        print(f"{r['name']:<30} {'Yes' if r['has_gnn'] else 'No':<6} "
              f"{'Yes' if r['has_moe'] else 'No':<6} {r['n_params']:>10,} "
              f"{r['weight_std']:<10.4f} {status:<10}")


if __name__ == "__main__":
    main()
