"""
ArcisVLM AutoKernel Benchmark Suite

Runs all Triton kernel benchmarks at multiple sizes and reports a
comparison table showing PyTorch vs Triton performance for each kernel.

Usage:
    python -m autokernel.benchmark
    python autokernel/benchmark.py
"""

import sys
import os

# Ensure project root is on path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from autokernel.kernels.fused_attention import benchmark_attention, validate_attention
from autokernel.kernels.fused_moe import benchmark_moe, validate_moe
from autokernel.kernels.fused_layernorm import benchmark_layernorm, validate_layernorm
from autokernel.kernels.fused_mlp import benchmark_mlp, validate_mlp


def print_header():
    print("=" * 80)
    print("  ArcisVLM AutoKernel Benchmark Suite")
    print("=" * 80)
    print()

    try:
        import triton
        print(f"  Triton version: {triton.__version__}")
    except ImportError:
        print("  Triton: NOT INSTALLED (using PyTorch fallback for all kernels)")

    import torch
    print(f"  PyTorch version: {torch.__version__}")
    print(f"  CUDA available: {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"  GPU: {torch.cuda.get_device_name(0)}")
    print()


def run_correctness_tests():
    """Run correctness validation for all kernels."""
    print("-" * 80)
    print("  Correctness Validation")
    print("-" * 80)

    tests = [
        ("fused_attention (bidir)", lambda: validate_attention(is_causal=False)),
        ("fused_attention (causal)", lambda: validate_attention(is_causal=True)),
        ("fused_layernorm", validate_layernorm),
        ("fused_mlp (SwiGLU)", validate_mlp),
        ("fused_moe", validate_moe),
    ]

    all_pass = True
    for name, fn in tests:
        max_diff, close = fn()
        status = "PASS" if close else "FAIL"
        if not close:
            all_pass = False
        print(f"  {name:30s}  {status}  (max_diff={max_diff:.2e})")

    print()
    return all_pass


def run_benchmarks():
    """Run all kernel benchmarks at multiple sizes."""
    print("-" * 80)
    print("  Performance Benchmarks")
    print("-" * 80)
    print()

    # Define benchmark configurations
    configs = {
        "fused_attention": [
            {"batch": 1, "seq_len": 256, "dim": 768, "num_heads": 12},
            {"batch": 4, "seq_len": 512, "dim": 768, "num_heads": 12},
            {"batch": 4, "seq_len": 1024, "dim": 1024, "num_heads": 16},
            {"batch": 2, "seq_len": 2048, "dim": 1024, "num_heads": 16},
        ],
        "fused_layernorm": [
            {"batch": 4, "seq_len": 512, "embed_dim": 768},
            {"batch": 4, "seq_len": 1024, "embed_dim": 768},
            {"batch": 4, "seq_len": 1024, "embed_dim": 1536},
        ],
        "fused_mlp": [
            {"batch": 4, "seq_len": 512, "embed_dim": 768, "expansion": 4},
            {"batch": 4, "seq_len": 1024, "embed_dim": 768, "expansion": 4},
        ],
        "fused_moe": [
            {"batch": 4, "seq_len": 256, "embed_dim": 768, "num_experts": 5, "top_k": 2},
            {"batch": 4, "seq_len": 512, "embed_dim": 768, "num_experts": 5, "top_k": 2},
        ],
    }

    all_results = []

    # Table header
    print(f"  {'Kernel':<22s} {'Config':<30s} {'PyTorch (ms)':>12s} {'Triton (ms)':>12s} {'Speedup':>8s}")
    print(f"  {'-'*22} {'-'*30} {'-'*12} {'-'*12} {'-'*8}")

    # Attention benchmarks
    for cfg in configs["fused_attention"]:
        res = benchmark_attention(**cfg, warmup=3, iters=10)
        conf_str = f"B={cfg['batch']} T={cfg['seq_len']} D={cfg['dim']}"
        print(f"  {'fused_attention':<22s} {conf_str:<30s} {res['pytorch_ms']:>12.3f} {res['triton_ms']:>12.3f} {res['speedup']:>7.2f}x")
        all_results.append(res)

    # LayerNorm benchmarks
    for cfg in configs["fused_layernorm"]:
        res = benchmark_layernorm(**cfg, warmup=3, iters=10)
        conf_str = f"B={cfg['batch']} T={cfg['seq_len']} D={cfg['embed_dim']}"
        print(f"  {'fused_layernorm':<22s} {conf_str:<30s} {res['pytorch_ms']:>12.3f} {res['triton_ms']:>12.3f} {res['speedup']:>7.2f}x")
        all_results.append(res)

    # MLP benchmarks
    for cfg in configs["fused_mlp"]:
        res = benchmark_mlp(**cfg, warmup=3, iters=10)
        conf_str = f"B={cfg['batch']} T={cfg['seq_len']} D={cfg['embed_dim']}"
        print(f"  {'fused_mlp (SwiGLU)':<22s} {conf_str:<30s} {res['pytorch_ms']:>12.3f} {res['triton_ms']:>12.3f} {res['speedup']:>7.2f}x")
        all_results.append(res)

    # MoE benchmarks
    for cfg in configs["fused_moe"]:
        res = benchmark_moe(**cfg, warmup=3, iters=10)
        conf_str = f"B={cfg['batch']} T={cfg['seq_len']} E={cfg['num_experts']}"
        print(f"  {'fused_moe':<22s} {conf_str:<30s} {res['pytorch_ms']:>12.3f} {res['triton_ms']:>12.3f} {res['speedup']:>7.2f}x")
        all_results.append(res)

    print()
    return all_results


def estimate_model_speedup(results):
    """
    Estimate overall model speedup based on kernel benchmarks.

    Uses approximate time breakdown for a typical ArcisVLM forward pass:
      - Attention: ~40% of compute
      - FFN/MLP: ~30% of compute
      - MoE routing: ~15% of compute
      - LayerNorm: ~5% of compute
      - Other (embeddings, projections): ~10%
    """
    print("-" * 80)
    print("  Estimated Model Speedup")
    print("-" * 80)

    # Get average speedup per kernel type
    kernel_speedups = {}
    for res in results:
        name = res["kernel"]
        if name not in kernel_speedups:
            kernel_speedups[name] = []
        kernel_speedups[name].append(res["speedup"])

    avg_speedups = {k: sum(v) / len(v) for k, v in kernel_speedups.items()}

    # Weighted contribution to overall speedup
    weights = {
        "fused_attention": 0.40,
        "fused_mlp": 0.30,
        "fused_moe": 0.15,
        "fused_layernorm": 0.05,
    }
    other_fraction = 0.10  # unoptimised portion

    print(f"\n  {'Component':<22s} {'Weight':>8s} {'Avg Speedup':>12s} {'Effective':>10s}")
    print(f"  {'-'*22} {'-'*8} {'-'*12} {'-'*10}")

    weighted_time_fraction = other_fraction  # unoptimised stays at 1.0x
    for kernel, weight in weights.items():
        spd = avg_speedups.get(kernel, 1.0)
        effective = weight / spd  # fraction of original time this component now takes
        weighted_time_fraction += effective
        print(f"  {kernel:<22s} {weight:>7.0%} {spd:>11.2f}x {effective:>9.3f}")

    print(f"  {'other (unoptimised)':<22s} {other_fraction:>7.0%} {'1.00':>11s}x {other_fraction:>9.3f}")

    overall_speedup = 1.0 / weighted_time_fraction
    print(f"\n  Overall estimated model speedup: {overall_speedup:.2f}x")

    if not any(r["has_triton"] and r["device"] == "cuda" for r in results):
        print("\n  NOTE: Running on CPU without Triton. Speedup is 1.0x (fallback path).")
        print("        Install Triton and run on a CUDA GPU to see actual kernel speedups.")
        print("        Expected GPU speedups: attention ~2-3x, layernorm ~1.5-2x, mlp ~1.3-1.5x")

    print()


def main():
    print_header()

    all_pass = run_correctness_tests()
    if not all_pass:
        print("  WARNING: Some correctness tests failed!")
        print()

    results = run_benchmarks()
    estimate_model_speedup(results)

    # Final summary
    print("=" * 80)
    correctness_status = "ALL PASS" if all_pass else "SOME FAILURES"
    print(f"  Correctness: {correctness_status}")
    print(f"  Benchmarks completed: {len(results)}")
    print("=" * 80)


if __name__ == "__main__":
    main()