"""01_smoke_test.py -- verify the kernel loads and runs.

Loads NeuronMamba3Mixer from the HF Hub kernel (or a local clone), instantiates
a small SISO mixer, and runs a single forward on device="neuron".

Requirements:
  - trn2.3xlarge (or larger) AWS Neuron instance
  - Neuron SDK Beta 3+ with PyTorch Native (torch-neuronx 2.11.3+)
  - kernels library >= 0.15.2

Run:
  python 01_smoke_test.py
"""

import time
import torch
from _loader import load_kernel


def main():
    print("=" * 60)
    print("mamba3-neuron-kernels smoke test")
    print("=" * 60)

    mamba3 = load_kernel()
    print(f"kernel version: {mamba3.__version__}")

    torch.manual_seed(0)
    # Small SISO mixer at reference config
    mixer = mamba3.NeuronMamba3Mixer(
        d_model=1024,
        d_state=128,
        headdim=64,
        chunk_size=64,
        mimo_rank=1,
        use_nki_ssd=True,
    ).to("neuron")
    print(f"instantiated SISO mixer: d_model=1024, nheads={mixer.nheads}, NKI kernels enabled")

    u = torch.randn(1, 128, 1024, device="neuron")

    print("\nrunning first forward (includes kernel compile)...")
    t0 = time.perf_counter()
    y, cache = mixer(u)
    try:
        torch.neuron.synchronize()
    except Exception:
        pass
    first_ms = (time.perf_counter() - t0) * 1000
    print(f"  first call: {first_ms:.0f} ms")

    print("running warm forward...")
    t0 = time.perf_counter()
    y, cache = mixer(u)
    try:
        torch.neuron.synchronize()
    except Exception:
        pass
    warm_ms = (time.perf_counter() - t0) * 1000
    print(f"  warm call:  {warm_ms:.2f} ms")

    assert y.shape == (1, 128, 1024), f"y shape unexpected: {y.shape}"
    assert cache.ssm_state.shape == (1, 32, 64, 128), f"ssm_state shape unexpected: {cache.ssm_state.shape}"

    print(f"\ny.shape = {y.shape}")
    print(f"cache.ssm_state.shape = {cache.ssm_state.shape}")
    print(f"cache.prev_Bx.shape = {cache.prev_Bx.shape}")
    print(f"cache.cum_angle.shape = {cache.cum_angle.shape}")

    print("\n=== SMOKE TEST PASS ===")


if __name__ == "__main__":
    main()