"""
OpenMythos vs. vanilla GQA+MoE transformer benchmark.

Compares OpenMythos (Prelude + looped Recurrent Block + Coda with ACT halting,
LTI-stable injection, LoRA depth adapter) against a parameter-matched vanilla
transformer built from the same GQAttention + MoEFFN building blocks stacked
non-recurrently. The baseline reuses OpenMythos primitives so the comparison
isolates the recurrent-depth architecture, not the kernels.

Metrics reported:
    - Parameter counts (total, MoE-active approximation)
    - Prefill latency + throughput at several sequence lengths
    - Decode (autoregressive step) latency with KV cache
    - Peak memory (CUDA only)
    - OpenMythos depth-scaling sweep: latency vs. n_loops

Run:
    python benchmarks/bench_vs_transformer.py                     # small CPU/GPU smoke test
    python benchmarks/bench_vs_transformer.py --size 1b --device cuda
    python benchmarks/bench_vs_transformer.py --seq-lens 128,512,2048 --n-loops 1,4,8,16
"""

from __future__ import annotations

import argparse
import gc
import statistics
import time
from dataclasses import dataclass
from typing import Optional

import torch
import torch.nn as nn

from open_mythos import MythosConfig, OpenMythos, mythos_1b
from open_mythos.main import (
    RMSNorm,
    TransformerBlock,
    precompute_rope_freqs,
)


# ---------------------------------------------------------------------------
# Baseline: non-looped GQA + MoE transformer
# ---------------------------------------------------------------------------


class BaselineTransformer(nn.Module):
    """
    Vanilla decoder-only transformer with GQA attention and MoE FFNs, stacked
    non-recurrently. Shares TransformerBlock / GQAttention / MoEFFN kernels
    with OpenMythos so any speed delta is attributable to the recurrent-depth
    architecture rather than the underlying attention/FFN implementation.
    """

    def __init__(self, cfg: MythosConfig, n_layers: int):
        super().__init__()
        self.cfg = cfg
        self.n_layers = n_layers
        self.embed = nn.Embedding(cfg.vocab_size, cfg.dim)
        self.layers = nn.ModuleList(
            [TransformerBlock(cfg, use_moe=True) for _ in range(n_layers)]
        )
        self.norm = RMSNorm(cfg.dim)
        self.head = nn.Linear(cfg.dim, cfg.vocab_size, bias=False)

        head_dim = cfg.dim // cfg.n_heads
        self.register_buffer(
            "freqs_cis",
            precompute_rope_freqs(head_dim, cfg.max_seq_len, cfg.rope_theta),
            persistent=False,
        )

    @staticmethod
    def _causal_mask(T: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
        mask = torch.full((1, 1, T, T), float("-inf"), device=device, dtype=dtype)
        return torch.triu(mask, diagonal=1)

    def forward(
        self,
        input_ids: torch.Tensor,
        kv_cache: Optional[dict] = None,
        start_pos: int = 0,
    ) -> torch.Tensor:
        T = input_ids.shape[1]
        x = self.embed(input_ids)
        freqs_cis = self.freqs_cis[start_pos : start_pos + T]
        mask = self._causal_mask(T, x.device, x.dtype) if T > 1 else None
        for i, layer in enumerate(self.layers):
            x = layer(x, freqs_cis, mask, kv_cache, cache_key=f"layer_{i}")
        return self.head(self.norm(x))


# ---------------------------------------------------------------------------
# Timing utilities
# ---------------------------------------------------------------------------


def _sync(device: torch.device) -> None:
    if device.type == "cuda":
        torch.cuda.synchronize()


def time_fn(fn, device: torch.device, warmup: int = 2, trials: int = 5) -> float:
    """Returns median wall-clock seconds over `trials` after `warmup` runs."""
    for _ in range(warmup):
        fn()
    _sync(device)
    times = []
    for _ in range(trials):
        _sync(device)
        t0 = time.perf_counter()
        fn()
        _sync(device)
        times.append(time.perf_counter() - t0)
    return statistics.median(times)


def peak_mem_mb(device: torch.device) -> float:
    if device.type != "cuda":
        return 0.0
    return torch.cuda.max_memory_allocated() / (1024 * 1024)


def reset_mem(device: torch.device) -> None:
    gc.collect()
    if device.type == "cuda":
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()


# ---------------------------------------------------------------------------
# Parameter counting
# ---------------------------------------------------------------------------


@dataclass
class ParamCounts:
    total: int
    moe_active_est: int  # active per token (shared + top-k routed)


def count_params(model: nn.Module, cfg: MythosConfig) -> ParamCounts:
    total = sum(p.numel() for p in model.parameters())
    # Rough active-per-token count for MoE layers: shared + top-k routed fraction.
    # For simplicity we report total and an estimated activation ratio separately.
    active_ratio = (cfg.n_shared_experts + cfg.n_experts_per_tok) / (
        cfg.n_shared_experts + cfg.n_experts
    )
    # Only FFN parameters shrink under activation; attention + embed/head are always on.
    # This is a coarse lower bound on active params.
    ffn_params = 0
    other_params = 0
    for name, p in model.named_parameters():
        if ".ffn." in name or name.startswith("ffn.") or ".experts." in name:
            ffn_params += p.numel()
        else:
            other_params += p.numel()
    active_est = other_params + int(ffn_params * active_ratio)
    return ParamCounts(total=total, moe_active_est=active_est)


# ---------------------------------------------------------------------------
# Benchmarks
# ---------------------------------------------------------------------------


def bench_prefill(
    model: nn.Module,
    vocab_size: int,
    batch: int,
    seq_len: int,
    device: torch.device,
    n_loops: Optional[int] = None,
) -> tuple[float, float]:
    """Returns (median_seconds, tokens_per_sec)."""
    ids = torch.randint(0, vocab_size, (batch, seq_len), device=device)

    if isinstance(model, OpenMythos):

        def run() -> None:
            with torch.no_grad():
                model(ids, n_loops=n_loops)

    else:

        def run() -> None:
            with torch.no_grad():
                model(ids)

    secs = time_fn(run, device)
    tps = (batch * seq_len) / secs
    return secs, tps


def bench_decode(
    model: nn.Module,
    vocab_size: int,
    batch: int,
    prompt_len: int,
    decode_steps: int,
    device: torch.device,
    n_loops: Optional[int] = None,
) -> tuple[float, float]:
    """
    Prefill a `prompt_len` prompt, then time `decode_steps` single-token decode
    steps with KV cache. Returns (avg_seconds_per_step, decode_tokens_per_sec).
    """
    prompt = torch.randint(0, vocab_size, (batch, prompt_len), device=device)

    def one_run() -> None:
        kv_cache: dict = {}
        with torch.no_grad():
            if isinstance(model, OpenMythos):
                model(prompt, n_loops=n_loops, kv_cache=kv_cache, start_pos=0)
            else:
                model(prompt, kv_cache=kv_cache, start_pos=0)
            for i in range(decode_steps):
                next_tok = torch.randint(0, vocab_size, (batch, 1), device=device)
                if isinstance(model, OpenMythos):
                    model(
                        next_tok,
                        n_loops=n_loops,
                        kv_cache=kv_cache,
                        start_pos=prompt_len + i,
                    )
                else:
                    model(next_tok, kv_cache=kv_cache, start_pos=prompt_len + i)

    secs = time_fn(one_run, device, warmup=1, trials=3)
    per_step = secs / decode_steps
    tps = batch * decode_steps / secs
    return per_step, tps


# ---------------------------------------------------------------------------
# Config helpers
# ---------------------------------------------------------------------------


def small_cfg() -> MythosConfig:
    """Tiny config for smoke tests — runs on CPU in seconds."""
    return MythosConfig(
        vocab_size=1024,
        dim=256,
        n_heads=8,
        n_kv_heads=2,
        max_seq_len=1024,
        max_loop_iters=4,
        prelude_layers=1,
        coda_layers=1,
        attn_type="gqa",
        n_experts=8,
        n_shared_experts=1,
        n_experts_per_tok=2,
        expert_dim=128,
        lora_rank=4,
        dropout=0.0,
    )


def get_cfg(size: str) -> MythosConfig:
    size = size.lower()
    if size == "small":
        return small_cfg()
    if size == "1b":
        cfg = mythos_1b()
        # GQA for apples-to-apples; MLA changes KV shape semantics.
        cfg.attn_type = "gqa"
        return cfg
    raise ValueError(f"unknown size: {size!r} (use 'small' or '1b')")


# ---------------------------------------------------------------------------
# Reporting
# ---------------------------------------------------------------------------


def fmt_count(n: int) -> str:
    for unit in ("", "K", "M", "B", "T"):
        if abs(n) < 1000:
            return f"{n:.2f}{unit}"
        n /= 1000
    return f"{n:.2f}P"


def print_header(title: str) -> None:
    bar = "=" * 72
    print(f"\n{bar}\n{title}\n{bar}")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("--size", default="small", choices=["small", "1b"])
    p.add_argument(
        "--device",
        default="cuda" if torch.cuda.is_available() else "cpu",
        choices=["cuda", "cpu"],
    )
    p.add_argument(
        "--dtype",
        default="auto",
        choices=["auto", "fp32", "bf16", "fp16"],
        help="'auto' picks fp32 on CPU and bf16 on CUDA",
    )
    p.add_argument("--batch", type=int, default=1)
    p.add_argument(
        "--seq-lens",
        default="128,512",
        help="comma-separated prefill sequence lengths",
    )
    p.add_argument(
        "--n-loops",
        default="1,4,8",
        help="comma-separated loop counts to sweep (OpenMythos only)",
    )
    p.add_argument(
        "--decode-steps",
        type=int,
        default=32,
        help="number of autoregressive decode steps after prefill",
    )
    p.add_argument(
        "--decode-prompt-len",
        type=int,
        default=128,
        help="prefill length before decode",
    )
    return p.parse_args()


def main() -> None:
    args = parse_args()
    device = torch.device(args.device)
    dtype_arg = args.dtype
    if dtype_arg == "auto":
        dtype_arg = "bf16" if device.type == "cuda" else "fp32"
    dtype = {
        "fp32": torch.float32,
        "bf16": torch.bfloat16,
        "fp16": torch.float16,
    }[dtype_arg]

    seq_lens = [int(s) for s in args.seq_lens.split(",") if s.strip()]
    n_loops_sweep = [int(s) for s in args.n_loops.split(",") if s.strip()]

    cfg = get_cfg(args.size)
    print_header(f"Config: size={args.size}  device={device}  dtype={dtype_arg}")
    print(
        f"  dim={cfg.dim}  n_heads={cfg.n_heads}  n_kv_heads={cfg.n_kv_heads}  "
        f"prelude={cfg.prelude_layers}  coda={cfg.coda_layers}  "
        f"max_loop_iters={cfg.max_loop_iters}\n"
        f"  experts={cfg.n_experts}  shared={cfg.n_shared_experts}  "
        f"top_k={cfg.n_experts_per_tok}  expert_dim={cfg.expert_dim}"
    )

    # Build models. Baseline depth = prelude + 1 (one unique recurrent block) + coda
    # to match the unique-parameter depth of OpenMythos (parameter-matched baseline).
    baseline_n_layers = cfg.prelude_layers + 1 + cfg.coda_layers

    torch.manual_seed(0)
    mythos = OpenMythos(cfg).to(device=device, dtype=dtype).eval()
    torch.manual_seed(0)
    baseline = (
        BaselineTransformer(cfg, n_layers=baseline_n_layers)
        .to(device=device, dtype=dtype)
        .eval()
    )

    m_params = count_params(mythos, cfg)
    b_params = count_params(baseline, cfg)
    print_header(
        "Parameters (block-matched: baseline depth = prelude + 1 recurrent + coda)"
    )
    print(
        f"  OpenMythos : total={fmt_count(m_params.total):>10}   "
        f"active/tok≈{fmt_count(m_params.moe_active_est):>10}"
    )
    print(
        f"  Baseline   : total={fmt_count(b_params.total):>10}   "
        f"active/tok≈{fmt_count(b_params.moe_active_est):>10}"
    )
    print(
        f"  Baseline unique layers = {baseline_n_layers}  "
        f"(Mythos total runtime depth at max_loops = "
        f"{cfg.prelude_layers + cfg.max_loop_iters + cfg.coda_layers})"
    )

    # ---- Prefill ----
    print_header("Prefill latency (batch={batch})".format(batch=args.batch))
    header = f"  {'model':<26} {'seq':>6} {'sec':>10} {'tok/s':>12} {'peak MB':>10}"
    print(header)
    for seq_len in seq_lens:
        if seq_len > cfg.max_seq_len:
            print(f"  skip seq_len={seq_len} (> max_seq_len={cfg.max_seq_len})")
            continue

        reset_mem(device)
        secs, tps = bench_prefill(baseline, cfg.vocab_size, args.batch, seq_len, device)
        mem = peak_mem_mb(device)
        print(
            f"  {'Baseline (stacked)':<26} {seq_len:>6} "
            f"{secs*1000:>9.2f}ms {tps:>12,.0f} {mem:>10.1f}"
        )

        for nl in n_loops_sweep:
            reset_mem(device)
            secs, tps = bench_prefill(
                mythos, cfg.vocab_size, args.batch, seq_len, device, n_loops=nl
            )
            mem = peak_mem_mb(device)
            print(
                f"  {'OpenMythos (loops=' + str(nl) + ')':<26} {seq_len:>6} "
                f"{secs*1000:>9.2f}ms {tps:>12,.0f} {mem:>10.1f}"
            )

    # ---- Decode ----
    print_header(
        f"Decode latency (prefill {args.decode_prompt_len} tokens + "
        f"{args.decode_steps} decode steps, batch={args.batch})"
    )
    print(f"  {'model':<26} {'sec/step':>12} {'decode tok/s':>14}")

    reset_mem(device)
    per_step, tps = bench_decode(
        baseline,
        cfg.vocab_size,
        args.batch,
        args.decode_prompt_len,
        args.decode_steps,
        device,
    )
    print(f"  {'Baseline (stacked)':<26} {per_step*1000:>10.2f}ms {tps:>14,.1f}")

    for nl in n_loops_sweep:
        reset_mem(device)
        per_step, tps = bench_decode(
            mythos,
            cfg.vocab_size,
            args.batch,
            args.decode_prompt_len,
            args.decode_steps,
            device,
            n_loops=nl,
        )
        print(
            f"  {'OpenMythos (loops=' + str(nl) + ')':<26} "
            f"{per_step*1000:>10.2f}ms {tps:>14,.1f}"
        )

    # ---- Depth scaling ----
    print_header(
        "OpenMythos depth scaling (fixed seq={}, batch={})".format(
            seq_lens[0], args.batch
        )
    )
    print(f"  {'n_loops':>8} {'sec':>10} {'tok/s':>12} {'Δ vs loops=1':>14}")
    base_secs = None
    for nl in n_loops_sweep:
        reset_mem(device)
        secs, tps = bench_prefill(
            mythos, cfg.vocab_size, args.batch, seq_lens[0], device, n_loops=nl
        )
        if base_secs is None:
            base_secs = secs
            delta = "1.00x"
        else:
            delta = f"{secs / base_secs:.2f}x"
        print(f"  {nl:>8} {secs*1000:>9.2f}ms {tps:>12,.0f} {delta:>14}")

    print("\nDone.")


if __name__ == "__main__":
    main()