# ─────────────────────────────────────────────────────────
# EGen Labs | ErebusTN
# AXIOM-1 — EGen-Core/AXIOM-1
# File: axiom1/evals/loop_stats.py
# ─────────────────────────────────────────────────────────

"""
Loop differentiation and ADC depth statistics for AXIOM-1.

Monitors whether each loop step meaningfully changes the hidden state.
If cos_sim ≈ 1.0 everywhere → loop collapse. Stop training and diagnose.
"""

import torch
import torch.nn.functional as F

from axiom1.model.model import AXIOM1


@torch.no_grad()
def check_loop_differentiation(
    model: AXIOM1,
    input_ids: torch.Tensor,
    collapse_threshold: float = 0.999
) -> dict:
    """
    Verifies that each loop step meaningfully changes the hidden state.

    If cos_sim ≈ 1.0 everywhere → loop collapse. Stop training and diagnose.

    Args:
        model: AXIOM1 model instance
        input_ids: Token IDs (batch, seq_len)
        collapse_threshold: cos_sim above this value = collapse risk

    Returns:
        Dict with 'steps' (list of per-step metrics), 'collapsed' (bool),
        'collapse_ratio' (float: fraction of steps showing collapse risk).
    """
    model.eval()

    h = model.token_embed(input_ids) * (model.cfg.d_model ** 0.5)
    h = model.pre_loop_norm(h)
    prev = h.clone()

    step_results: list[dict[str, float]] = []
    collapse_count = 0

    print("Loop differentiation check:")
    for step in range(model.cfg.num_loops):
        lse = model.loop_step_embed(torch.tensor(step, device=h.device))
        h = h + lse.unsqueeze(0).unsqueeze(0)
        h = model.block(h)

        delta = (h - prev).norm().item()
        cos_sim = F.cosine_similarity(h.flatten(), prev.flatten(), dim=0).item()
        status = "OK" if delta > 0.1 and cos_sim < collapse_threshold else "⚠ COLLAPSE RISK"
        if cos_sim >= collapse_threshold:
            collapse_count += 1
        print(f"  Step {step:2d}: Δ={delta:.4f}, cos_sim={cos_sim:.5f}  [{status}]")

        step_results.append({
            'step': step,
            'delta': delta,
            'cos_sim': cos_sim,
            'status': status
        })
        prev = h.clone()

    collapse_ratio = collapse_count / max(model.cfg.num_loops, 1)
    collapsed = collapse_ratio > 0.5

    if collapsed:
        print(f"  ╔═══════════════════════════════════════════════╗")
        print(f"  ║  CRITICAL: LOOP COLLAPSE DETECTED             ║")
        print(f"  ║  {collapse_count}/{model.cfg.num_loops} steps show cos_sim ≥ {collapse_threshold}  ║")
        print(f"  ║  Stop training and diagnose immediately.       ║")
        print(f"  ╚═══════════════════════════════════════════════╝")
    else:
        print(f"  Loop health: {model.cfg.num_loops - collapse_count}/{model.cfg.num_loops} steps OK")

    return {
        'steps': step_results,
        'collapsed': collapsed,
        'collapse_ratio': collapse_ratio,
    }


@torch.no_grad()
def check_adc_depth_stats(
    model: AXIOM1,
    input_ids: torch.Tensor,
    max_loops: int = 16
) -> dict[str, float]:
    """
    Compute ADC (Adaptive Depth Control) statistics.

    Only meaningful during Phase 3 when adaptive depth is enabled.
    Reports average loop depth and distribution.

    Args:
        model: AXIOM1 model instance with ADC enabled
        input_ids: Token IDs (batch, seq_len)
        max_loops: Maximum number of loops for ADC

    Returns:
        Dict with 'avg_depth', 'min_depth', 'max_depth', 'depth_ratio'.
    """
    model.eval()

    if not getattr(model.cfg, 'use_adaptive_depth', False):
        print("[AXIOM-1] ADC not enabled — skipping depth stats")
        return {'avg_depth': float(model.cfg.num_loops),
                'min_depth': float(model.cfg.num_loops),
                'max_depth': float(model.cfg.num_loops),
                'depth_ratio': 1.0}

    # For ADC, we would need the adaptive_forward method
    # This is a placeholder that returns fixed-depth stats when ADC is not yet
    # implemented (Phase 1 & 2)
    avg_depth = float(model.cfg.num_loops)
    print(f"[AXIOM-1] ADC depth stats: avg={avg_depth:.1f}, "
          f"ratio={avg_depth / max_loops:.2f}")

    return {
        'avg_depth': avg_depth,
        'min_depth': avg_depth,
        'max_depth': avg_depth,
        'depth_ratio': avg_depth / max_loops
    }
