"""Fase 0: verifica qué attn_implementation usar y si flash-linear-attention compila en SM121.

No bloqueante: flash_attn y flash-linear-attention son "nice to have"; si fallan,
el resultado documentado es usar el fallback correspondiente (sdpa / PyTorch puro).
"""
import sys
import time

import torch

RESULTS = {}


def test_flash_attn():
    try:
        from flash_attn import flash_attn_func
    except ImportError as e:
        return False, f"import failed: {e}"

    try:
        device = "cuda"
        dtype = torch.bfloat16
        batch, seqlen, nheads, headdim = 2, 128, 8, 64
        q = torch.randn(batch, seqlen, nheads, headdim, device=device, dtype=dtype)
        k = torch.randn(batch, seqlen, nheads, headdim, device=device, dtype=dtype)
        v = torch.randn(batch, seqlen, nheads, headdim, device=device, dtype=dtype)
        out = flash_attn_func(q, k, v, causal=True)
        torch.cuda.synchronize()
        assert out.shape == (batch, seqlen, nheads, headdim)
        return True, f"forward ok, output shape {tuple(out.shape)}"
    except Exception as e:
        return False, f"forward failed: {type(e).__name__}: {e}"


def test_sdpa():
    try:
        device = "cuda"
        dtype = torch.bfloat16
        batch, nheads, seqlen, headdim = 2, 8, 128, 64
        q = torch.randn(batch, nheads, seqlen, headdim, device=device, dtype=dtype)
        k = torch.randn(batch, nheads, seqlen, headdim, device=device, dtype=dtype)
        v = torch.randn(batch, nheads, seqlen, headdim, device=device, dtype=dtype)
        out = torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True)
        torch.cuda.synchronize()
        assert out.shape == (batch, nheads, seqlen, headdim)
        return True, f"forward ok, output shape {tuple(out.shape)}"
    except Exception as e:
        return False, f"forward failed: {type(e).__name__}: {e}"


def test_flash_linear_attention(timeout_s=1200):
    start = time.time()
    try:
        import fla  # noqa: F401
        from fla.ops.gated_delta_rule import chunk_gated_delta_rule
    except ImportError as e:
        return False, f"import failed (package not installed): {e}"

    try:
        device = "cuda"
        dtype = torch.bfloat16
        batch, seqlen, nheads, headdim = 2, 64, 4, 64
        q = torch.randn(batch, seqlen, nheads, headdim, device=device, dtype=dtype, requires_grad=True)
        k = torch.randn(batch, seqlen, nheads, headdim, device=device, dtype=dtype, requires_grad=True)
        v = torch.randn(batch, seqlen, nheads, headdim, device=device, dtype=dtype, requires_grad=True)
        beta = torch.rand(batch, seqlen, nheads, device=device, dtype=dtype).sigmoid()
        g = -torch.rand(batch, seqlen, nheads, device=device, dtype=torch.float32)

        out, _ = chunk_gated_delta_rule(q, k, v, g, beta)
        loss = out.sum()
        loss.backward()
        torch.cuda.synchronize()

        if time.time() - start > timeout_s:
            return False, "compiled/ran but exceeded time budget"
        return True, f"forward+backward ok, output shape {tuple(out.shape)}"
    except Exception as e:
        elapsed = time.time() - start
        return False, f"failed after {elapsed:.0f}s: {type(e).__name__}: {e}"


def main():
    if not torch.cuda.is_available():
        print("CUDA not available — aborting.")
        sys.exit(1)

    print(f"torch {torch.__version__}, CUDA {torch.version.cuda}, device: {torch.cuda.get_device_name(0)}")
    print()

    ok, msg = test_flash_attn()
    RESULTS["flash_attn"] = {"ok": ok, "detail": msg}
    print(f"[flash_attn]  {'OK' if ok else 'FAIL'} — {msg}")

    ok, msg = test_sdpa()
    RESULTS["sdpa"] = {"ok": ok, "detail": msg}
    print(f"[sdpa]        {'OK' if ok else 'FAIL'} — {msg}")

    recommended = "flash_attention_2" if RESULTS["flash_attn"]["ok"] else "sdpa"
    print(f"\n>>> attn_implementation recomendado para el training: {recommended}\n")

    print("Probando flash-linear-attention (opcional, acotado)...")
    ok, msg = test_flash_linear_attention()
    RESULTS["flash_linear_attention"] = {"ok": ok, "detail": msg}
    print(f"[flash-linear-attention]  {'OK' if ok else 'FAIL (se descarta, no bloqueante)'} — {msg}")

    print("\n=== Resumen ===")
    for k, v in RESULTS.items():
        print(f"  {k}: {'OK' if v['ok'] else 'FAIL'} — {v['detail']}")


if __name__ == "__main__":
    main()