"""Sanity-check script for the ASD custom generation method.

Run from anywhere:
    python scripts/example.py \
        --model Qwen/Qwen2.5-1.5B-Instruct \
        --assistant-model Qwen/Qwen2.5-0.5B-Instruct \
        --budget 2.0 --local-ratio 0.25 --max-mismatches 2

Checks:
1. ASD runs and produces text.
2. `assistant_asd_budget=0` is token-identical to native strict assisted decoding
   (the gold-standard correctness check).
3. No ASD knob set is also token-identical to native (lossless default).
"""

import argparse
import os

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", required=True)
    parser.add_argument("--assistant-model", required=True)
    parser.add_argument("--budget", type=float, default=2.0)
    parser.add_argument("--local-ratio", type=float, default=0.25)
    parser.add_argument("--max-mismatches", type=int, default=2)
    parser.add_argument("--max-new-tokens", type=int, default=64)
    parser.add_argument("--prompt", default="Give me a short introduction to large language models.")
    parser.add_argument("--repo-path", default=REPO_ROOT,
                        help="Directory containing custom_generate/ (default: repo root)")
    args = parser.parse_args()

    tokenizer = AutoTokenizer.from_pretrained(args.model)
    model = AutoModelForCausalLM.from_pretrained(args.model, dtype="auto", device_map="auto")
    assistant_model = AutoModelForCausalLM.from_pretrained(
        args.assistant_model, dtype="auto", device_map="auto"
    )

    inputs = tokenizer(args.prompt, return_tensors="pt").to(model.device)
    gen_kwargs = dict(max_new_tokens=args.max_new_tokens, do_sample=False)

    # 1. Native strict assisted decoding baseline.
    strict_out = model.generate(**inputs, assistant_model=assistant_model, **gen_kwargs)
    print("Strict      :", tokenizer.decode(strict_out[0], skip_special_tokens=True))

    # 2. ASD relaxed.
    asd_out = model.generate(
        **inputs,
        assistant_model=assistant_model,
        custom_generate=args.repo_path,
        trust_remote_code=True,
        assistant_asd_budget=args.budget,
        assistant_asd_local_ratio=args.local_ratio,
        assistant_asd_max_mismatches=args.max_mismatches,
        **gen_kwargs,
    )
    print("ASD         :", tokenizer.decode(asd_out[0], skip_special_tokens=True))

    # 3. Gold check: budget=0 must be token-identical to native strict decoding.
    asd_b0_out = model.generate(
        **inputs,
        assistant_model=assistant_model,
        custom_generate=args.repo_path,
        trust_remote_code=True,
        assistant_asd_budget=0.0,
        **gen_kwargs,
    )
    assert torch.equal(strict_out, asd_b0_out), "ASD(budget=0) differs from strict decoding!"
    print("OK: assistant_asd_budget=0 is token-identical to strict verification.")

    # 4. Lossless default: no ASD knobs must also be token-identical.
    asd_default_out = model.generate(
        **inputs,
        assistant_model=assistant_model,
        custom_generate=args.repo_path,
        trust_remote_code=True,
        **gen_kwargs,
    )
    assert torch.equal(strict_out, asd_default_out), "ASD (no knobs) differs from strict decoding!"
    print("OK: no ASD knobs is token-identical to strict verification.")


if __name__ == "__main__":
    main()