#!/usr/bin/env python3
"""
scripts/smoke_test.py
One-Command Deterministic Smoke Test for Q-TensorFormer (Phase 15).

Verifies:
  1. Canonical research model (src.models.QTensorFormer) construction, forward pass,
     backward pass, nested TT rank slicing, 8D information state, marginal-value allocator,
     and AdaptiveKVCache.
  2. Canonical Hugging Face model (modeling_qtensorformer.QTensorFormerForCausalLM)
     construction, forward pass, autoregressive generate(), and strict=True save/load round-trip.
  3. Strict checkpoint validation: verifies that loading an incompatible state_dict with
     strict=True raises RuntimeError as required.
"""

import sys
import tempfile
from pathlib import Path

import torch

ROOT_DIR = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT_DIR))

from configuration_qtensorformer import QTensorFormerConfig
from modeling_qtensorformer import QTensorFormerForCausalLM
from src.config import ModelConfig, validate_model_config
from src.kv_cache import AdaptiveKVCache, KVPrecision
from src.models import DenseBaseline, QTensorFormer


def main():
    torch.manual_seed(42)
    print("=" * 70)
    print("Q-TensorFormer — Deterministic End-to-End Smoke Test")
    print("=" * 70)

    # 1. Canonical Research Architecture (src.models.QTensorFormer)
    cfg = ModelConfig(
        d_model=64,
        n_heads=4,
        n_layers=2,
        ff_multiplier=4,
        max_seq_len=32,
        vocab_size=256,
        tt_rank=8,
        tt_min_rank=2,
        use_quantum=True,
        n_qubits=4,
    )
    assert validate_model_config(cfg, QTensorFormer) is True
    model = QTensorFormer(cfg, preset="QTF_BALANCED")
    baseline = DenseBaseline(cfg)

    inp = torch.randint(0, 256, (2, 12))
    logits, stats = model(inp, return_stats=True)
    assert logits.shape == (2, 12, 256), f"Unexpected logits shape: {logits.shape}"
    assert len(stats) == 2

    # Verify gradient flow
    loss = logits.pow(2).mean()
    loss.backward()
    grad_ok = all(
        p.grad is not None and torch.isfinite(p.grad).all()
        for p in model.parameters()
        if p.requires_grad and p.grad is not None
    )
    assert grad_ok, "Non-finite gradients detected in QTensorFormer"
    print(
        f"[PASS] src.models.QTensorFormer forward/backward verified "
        f"(Total params={model.total_params:,}, Active params={model.active_params:,}, "
        f"Dense baseline params={baseline.total_params:,})"
    )

    # 2. AdaptiveKVCache FP16 / INT8 / INT4 verification
    cache = AdaptiveKVCache(max_capacity=16, default_precision=KVPrecision.INT4, window_size=8)
    k = torch.randn(1, 4, 20, 16)
    v = torch.randn(1, 4, 20, 16)
    k_out, v_out = cache.update(k, v)
    assert k_out.shape == (1, 4, 16, 16)
    assert cache.evicted_tokens_count == 4
    print(
        f"[PASS] AdaptiveKVCache INT4 + sink-preserving eviction verified "
        f"(Retained={cache.seq_len}, Evicted={cache.evicted_tokens_count}, Bytes={cache.current_bytes})"
    )

    # 3. Canonical Hugging Face Architecture (QTensorFormerForCausalLM)
    hf_cfg = QTensorFormerConfig(
        vocab_size=256,
        hidden_size=64,
        num_hidden_layers=2,
        num_attention_heads=4,
        num_key_value_heads=2,
        intermediate_size=128,
        max_position_embeddings=64,
        tt_ranks=[1, 4, 4, 1],
    )
    assert hf_cfg.keys_to_ignore_at_loading == []
    hf_model = QTensorFormerForCausalLM(hf_cfg)
    hf_model.eval()

    with torch.no_grad():
        hf_out = hf_model(inp)
        assert hf_out.logits.shape == (2, 12, 256)
        gen_ids = hf_model.generate(inp[:1, :4], max_new_tokens=4, do_sample=False)
        assert gen_ids.shape == (1, 8)

    with tempfile.TemporaryDirectory() as tmpdir:
        ckpt_path = Path(tmpdir) / "strict_roundtrip.pt"
        torch.save(hf_model.state_dict(), ckpt_path)
        hf_reloaded = QTensorFormerForCausalLM(hf_cfg)
        hf_reloaded.load_state_dict(
            torch.load(ckpt_path, map_location="cpu", weights_only=True), strict=True
        )
        hf_reloaded.eval()
        with torch.no_grad():
            assert torch.allclose(hf_model(inp).logits, hf_reloaded(inp).logits, atol=1e-5)

    print("[PASS] QTensorFormerForCausalLM forward, generate(), and strict=True round-trip verified")

    # 4. Verify strict=True rejects mismatched state_dict
    strict_rejected = False
    try:
        hf_model.load_state_dict({"invalid.key": torch.zeros(1)}, strict=True)
    except RuntimeError:
        strict_rejected = True
    assert strict_rejected, "strict=True failed to reject mismatched state_dict!"
    print("[PASS] strict=True checkpoint mismatch rejection verified")

    print("=" * 70)
    print("ALL SMOKE TESTS PASSED (EXIT CODE 0)")
    print("=" * 70)


if __name__ == "__main__":
    main()