"""Decode speed of the Ling-mini-2.0 megakernel.

    python bench.py              one chat: 10 launches of 256 greedy tokens after a 32-token prompt, median
    python bench.py --chats 3    three chats at once, each on its share of the SMs: tokens/s each and in total,
                                 and whether each chat wrote the same 128 tokens as alone

The GPU is warmed up first (Windows lowers an idle GPU's clocks): launches run until they stop
getting faster for two seconds.
"""

from __future__ import annotations

import argparse
import statistics
import sys
import time

import torch

from ling_kernel import HERE, LingKernel, State

PROMPTS = [
    "The history of the printing press begins in the fifteenth century, when Johannes Gutenberg combined "
    "movable metal type, oil-based ink and a wooden press adapted from those used to make wine.",
    "def merge_sort(items):\n    \"\"\"Return a new list with the items in ascending order.\"\"\"\n"
    "    if len(items) <= 1:\n        return list(items)\n    middle = len(items) // 2\n",
]


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--chats", type=int, default=1)
    args = ap.parse_args()
    sys.stdout.reconfigure(encoding="utf-8")
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained(HERE)
    mk = LingKernel(max_pos=1024, log=lambda s: print(s, flush=True))
    prompts = [tok(p)["input_ids"][:32] for p in PROMPTS]
    firsts = [mk.forced(p)[-1] for p in prompts]
    best, since, start = float("inf"), time.time(), time.time()
    while time.time() - start < 30:
        s = mk.timed_greedy(firsts[0], len(prompts[0]), 64)
        if s < 0.97 * best:
            best, since = s, time.time()
        elif time.time() - since > 2:
            break
    if args.chats <= 1:
        mk.forced(prompts[0])
        times = [mk.timed_greedy(firsts[0], len(prompts[0]), 256) for _ in range(10)]
        per = statistics.median(times) / 256
        print(f"one chat: {1 / per:.0f} tokens/s (median of 10 launches of 256 tokens; best {256 / min(times):.0f})")
        return
    n = args.chats
    nb = max(16, mk.sms // n)
    states = [State(1024) for _ in range(n)]
    streams = [torch.cuda.Stream() for _ in range(n)]
    pr = [prompts[i % 2] for i in range(n)]
    for i in range(n):                                   # each chat reads its prompt
        mk.launch_async(states[i], pr[i], 0, len(pr[i]), True, nb, streams[i])
    torch.cuda.synchronize()
    first = [states[i].out[len(pr[i]) - 1].item() for i in range(n)]
    alone = []
    for i in range(2):
        mk.forced(prompts[i])
        alone.append(mk.greedy(firsts[i], len(prompts[i]), 128))

    def round_(k: int) -> float:
        torch.cuda.synchronize()
        t = time.perf_counter()
        for i in range(n):
            mk.launch_async(states[i], [first[i]], len(pr[i]), k, False, nb, streams[i])
        torch.cuda.synchronize()
        return time.perf_counter() - t

    round_(128)
    same = [states[i].out[:128].tolist() == alone[i % 2] for i in range(n)]
    t = statistics.median(round_(256) for _ in range(6))
    print(f"{n} chats at once, {nb} SMs each: {256 / t:.0f} tokens/s each, {n * 256 / t:.0f} in total; "
          f"each wrote the same 128 tokens as alone: {same}")


if __name__ == "__main__":
    main()