"""Fair batch-one decode comparison: HF eager, torch.compile, and CISM.

All baselines run the SAME model checkpoint in FP32 with a KV cache, greedy
decoding, one torch thread, warmup discarded. torch.compile uses a manual
decode loop with StaticCache (its recommended fast path), not slow
use_cache=False generate(). CISM numbers come from the same machine via
cism bench. Output: JSON to stdout. Exit 0 on success.
"""

import argparse
import json
import time

import torch


def timed(fn, warmup_tokens: int, measured_tokens: int):
    fn(warmup_tokens)
    start = time.perf_counter()
    generated = fn(measured_tokens)
    elapsed = time.perf_counter() - start
    return {
        "tokens": generated,
        "seconds": elapsed,
        "tokens_per_second": generated / elapsed if elapsed > 0 else 0,
    }


def run_eager(model, input_ids, eos_pad, count, warmup):
    def loop(tokens):
        out = model.generate(
            input_ids, attention_mask=None, max_new_tokens=tokens,
            min_new_tokens=tokens, do_sample=False, num_beams=1,
            use_cache=True, pad_token_id=eos_pad,
        )
        return tokens

    result = timed(loop, warmup, count)
    result["method"] = "transformers generate(use_cache=True) eager"
    return result


def run_manual_eager(model, input_ids, count, warmup):
    """Raw decode loop without generate()'s framework overhead."""
    from transformers import StaticCache

    cache = StaticCache(config=model.config, max_batch_size=1,
                        max_cache_len=warmup + count + len(input_ids[0]),
                        device="cpu", dtype=torch.float32)

    def loop(tokens):
        with torch.inference_mode():
            cache.reset()
            generated = 0
            out = model(input_ids, past_key_values=cache, use_cache=True)
            token = int(out.logits[0, -1].argmax())
            generated += 1
            one = torch.tensor([[token]])
            while generated < tokens:
                out = model(one, past_key_values=cache, use_cache=True)
                token = int(out.logits[0, -1].argmax())
                generated += 1
        return generated

    result = timed(loop, warmup, count)
    result["method"] = "manual eager decode loop + StaticCache (no generate overhead)"
    return result


def _dynamo_snapshot():
    import torch._dynamo

    return {name: sum(counts.values()) for name, counts in torch._dynamo.utils.counters.items()}


def _dynamo_delta(before, after):
    keys = set(before) | set(after)
    return {key: after.get(key, 0) - before.get(key, 0) for key in keys if after.get(key, 0) != before.get(key, 0)}


def run_compiled(model, input_ids, count, warmup, dynamic):
    from transformers import StaticCache

    cache = StaticCache(config=model.config, max_batch_size=1,
                        max_cache_len=warmup + count + len(input_ids[0]),
                        device="cpu", dtype=torch.float32)
    compiled = torch.compile(model, dynamic=dynamic)

    def loop(tokens):
        with torch.inference_mode():
            cache.reset()
            generated = 0
            out = compiled(input_ids, past_key_values=cache, use_cache=True)
            token = int(out.logits[0, -1].argmax())
            generated += 1
            one = torch.tensor([[token]])
            while generated < tokens:
                out = compiled(one, past_key_values=cache, use_cache=True)
                token = int(out.logits[0, -1].argmax())
                generated += 1
        return generated

    # Warmup must absorb ALL compilation and recompilation; counters are reset
    # afterwards so the measured window only reports steady-state guard hits.
    loop(warmup)
    loop(warmup)
    before = _dynamo_snapshot()
    generated = timed(loop, warmup, count)
    after = _dynamo_snapshot()
    generated["dynamo_counter_delta_during_measurement"] = _dynamo_delta(before, after)
    generated["method"] = f"torch.compile(model, dynamic={dynamic}) + StaticCache decode loop"
    return generated


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("model")
    parser.add_argument("--revision", default=None)
    parser.add_argument("--threads", type=int, default=1)
    parser.add_argument("--tokens", type=int, default=64)
    parser.add_argument("--warmup", type=int, default=32)
    parser.add_argument("--skip-compile", action="store_true")
    args = parser.parse_args()

    torch.set_num_threads(args.threads)
    torch.set_num_interop_threads(1)

    from transformers import AutoModelForCausalLM, AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(args.model, revision=args.revision)
    model = AutoModelForCausalLM.from_pretrained(
        args.model, revision=args.revision, dtype=torch.float32,
        attn_implementation="eager",
    ).to("cpu").eval()
    config = {"max_position_embeddings": model.config.max_position_embeddings}
    ids = tokenizer("The future of artificial intelligence is", return_tensors="pt").input_ids
    eos_pad = tokenizer.eos_token_id or 0

    report = {
        "model": args.model, "torch": torch.__version__, "threads": args.threads,
        "precision": "fp32", "greedy": True, "warmup_tokens": args.warmup,
        "measured_tokens": args.tokens,
    }
    report["eager_generate"] = run_eager(model, ids, eos_pad, args.tokens, args.warmup)
    report["eager_manual_loop"] = run_manual_eager(model, ids, args.tokens, args.warmup)
    if args.skip_compile:
        report["compiled_dynamic_false"] = {"skipped": True}
        report["compiled_dynamic_true"] = {"skipped": True}
    else:
        for dynamic in (False, True):
            key = f"compiled_dynamic_{str(dynamic).lower()}"
            try:
                report[key] = run_compiled(model, ids, args.tokens, args.warmup, dynamic)
            except Exception as error:
                report[key] = {"failed": f"{type(error).__name__}: {error}"}
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()