"""显式 --run 才执行测试时采样实验；先固定方案，再生成全部 32 个回答。"""

import argparse
from datetime import datetime, timezone
from importlib.metadata import version
import json
from pathlib import Path
import platform
import time

import test_time_compute as lab


def experiment(directory):
    directory.mkdir(parents=True, exist_ok=False)
    plan = lab.make_plan()
    lab.save_new(directory / "plan.json", plan)
    import torch
    from huggingface_hub import snapshot_download
    from transformers import AutoModelForCausalLM, AutoTokenizer
    settings = plan["generation"]
    torch.set_num_threads(settings["threads"])
    snapshot = Path(snapshot_download(
        plan["model"], revision=plan["revision"],
        allow_patterns=["*.json", "*.safetensors", "merges.txt", "README.md", "LICENSE"]))
    tokenizer = AutoTokenizer.from_pretrained(snapshot, trust_remote_code=False)
    model = AutoModelForCausalLM.from_pretrained(
        snapshot, dtype=torch.float32, attn_implementation=settings["attention"],
        trust_remote_code=False).eval()
    report = dict(
        complete=False, started_at=datetime.now(timezone.utc).isoformat(),
        model=plan["model"], revision=plan["revision"], generation=settings,
        plan_sha256=lab.fingerprint(plan),
        source_sha256={path.name: lab.digest(path) for path in (Path(lab.__file__), Path(__file__))},
        model_files_sha256={path.name: lab.digest(path) for path in snapshot.iterdir()
                            if path.is_file() and (path.suffix in {".json", ".safetensors"}
                                                   or path.name == "merges.txt")},
        environment=dict(python=platform.python_version(), platform=platform.platform(),
                         packages={name: version(name) for name in
                                   ("torch", "transformers", "huggingface-hub",
                                    "tokenizers", "safetensors")}),
        rows=[],
    )
    target = directory / "model.json"

    def checkpoint():
        temporary = directory / "model.pending.json"
        temporary.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n")
        temporary.replace(target)

    checkpoint()
    eos = model.generation_config.eos_token_id
    eos = [eos] if isinstance(eos, int) else eos
    for problem_index, problem in enumerate(plan["problems"]):
        conversation = lab.messages(problem)
        rendered = tokenizer.apply_chat_template(
            conversation, tokenize=False, add_generation_prompt=True)
        inputs = tokenizer.apply_chat_template(
            conversation, tokenize=True, add_generation_prompt=True,
            return_dict=True, return_tensors="pt")
        length = inputs["input_ids"].shape[1]
        if length + settings["max_new_tokens"] > settings["input_plus_output_cap"]:
            raise ValueError("超过固定预算，不静默截断")
        for sample in range(settings["samples"]):
            seed = settings["first_seed"] + problem_index * settings["samples"] + sample
            torch.manual_seed(seed)
            started = time.perf_counter()
            with torch.inference_mode():
                output = model.generate(
                    **inputs, max_new_tokens=settings["max_new_tokens"],
                    do_sample=True, temperature=settings["temperature"],
                    top_p=settings["top_p"], top_k=settings["top_k"],
                    pad_token_id=tokenizer.eos_token_id)
            seconds = time.perf_counter() - started
            ids = output[0, length:].tolist()
            text = tokenizer.decode(ids, skip_special_tokens=True)
            report["rows"].append(dict(
                id=problem["id"], sample=sample, seed=seed,
                messages=conversation, rendered_prompt=rendered,
                input_ids=inputs["input_ids"][0].tolist(), input_tokens=length,
                output_ids=ids, output_tokens=len(ids), text=text,
                stopped_on_eos=bool(ids and ids[-1] in eos), seconds=seconds,
                parsed=lab.parse_answer(text)))
            checkpoint()
            print(f"{problem['id']}/{sample}: answer={lab.parse_answer(text)}, "
                  f"output={len(ids)}, eos={report['rows'][-1]['stopped_on_eos']}", flush=True)
    report["complete"] = True
    report["finished_at"] = datetime.now(timezone.utc).isoformat()
    checkpoint()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--run", action="store_true")
    parser.add_argument("--out-dir", type=Path)
    args = parser.parse_args()
    if not args.run:
        print("显式传入 --run --out-dir 新目录，才执行固定的 32 次真实模型采样。")
        return
    if args.out_dir is None:
        parser.error("请指定尚不存在的输出目录")
    experiment(args.out_dir)


if __name__ == "__main__":
    main()