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

"""
INT8 quantization for AXIOM-1 FFN layers.

Only FFN layers (gate, up, down) are quantized — attention projections
and embeddings remain in bfloat16 to preserve quality.
INT4 is NEVER allowed (error compounds per loop iteration).

Usage:
    model = AXIOM1(cfg)
    model.load_state_dict(checkpoint['model_state_dict'])
    model = quantize_ffn_int8(model)
    # Verify quality:
    ok = verify_ppl_after_quantization(model, val_data, threshold=1.5)
"""

import torch
import torch.nn as nn

from axiom1.model.model import AXIOM1
from axiom1.model.config import AXIOM1Config


def quantize_ffn_int8(model: AXIOM1) -> AXIOM1:
    """
    Quantize FFN layers (gate, up, down) to INT8 using dynamic quantization.

    Only targets the SwiGLU FFN inside the shared LoopBlock.
    Attention projections and embeddings are left untouched.

    This is safe for AXIOM-1 because:
    - INT8 error per-layer is small enough that 12 loop iterations
      don't accumulate significant drift
    - Only FFN weights are quantized, not activations

    Args:
        model: AXIOM1 model instance (will be modified in-place)

    Returns:
        Quantized model (same instance, modified)
    """
    # Target only the FFN's Linear layers inside the shared block
    ffn = model.block.ffn

    # Dynamic quantization on the FFN linear layers
    ffn.gate = torch.ao.quantization.quantize_dynamic(
        nn.Sequential(ffn.gate), {nn.Linear}, dtype=torch.qint8
    )[0]
    ffn.up = torch.ao.quantization.quantize_dynamic(
        nn.Sequential(ffn.up), {nn.Linear}, dtype=torch.qint8
    )[0]
    ffn.down = torch.ao.quantization.quantize_dynamic(
        nn.Sequential(ffn.down), {nn.Linear}, dtype=torch.qint8
    )[0]

    print("[AXIOM-1] FFN layers quantized to INT8")
    print(f"  Quantized: block.ffn.gate, block.ffn.up, block.ffn.down")

    return model


@torch.no_grad()
def verify_ppl_after_quantization(
    model: AXIOM1,
    input_ids: torch.Tensor,
    targets: torch.Tensor,
    ppl_threshold: float = 1.5
) -> bool:
    """
    Verify that quantization did not degrade perplexity beyond threshold.

    Computes PPL on the provided data and checks it stays within
    ppl_threshold × baseline (where baseline is typical post-training PPL).

    Args:
        model: Quantized AXIOM1 model
        input_ids: Validation input IDs (batch, seq_len)
        targets: Validation target IDs (batch, seq_len)
        ppl_threshold: Maximum PPL multiplier (1.5 = allow 50% degradation)

    Returns:
        True if PPL is acceptable, False if quantization degraded too much.
    """
    import math
    import torch.nn.functional as F

    model.eval()
    logits = model(input_ids)

    # Compute cross-entropy loss
    loss = F.cross_entropy(
        logits[:, :-1].reshape(-1, logits.size(-1)),
        targets[:, 1:].reshape(-1),
        ignore_index=0  # Ignore padding
    )

    ppl = math.exp(min(loss.item(), 20.0))  # Cap to prevent overflow

    print(f"[AXIOM-1] Post-quantization PPL: {ppl:.2f}")
    print(f"  Threshold: {ppl_threshold:.1f}x baseline")

    if ppl > 100.0 * ppl_threshold:
        print(f"  ⚠ WARNING: PPL too high ({ppl:.2f}). Quantization may be harmful.")
        return False

    print(f"  ✓ PPL within acceptable range")
    return True
