#!/usr/bin/env python3
"""Select the deterministic 1000-prompt Predictor v4 training split."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import random
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_SOURCE = ROOT / "prompts" / "vidprom_filtered_extended.txt"
DEFAULT_HOLDOUT = ROOT / "prompts" / "MovieGenVideoBench_extended.txt"
DEFAULT_OUTPUT = Path(
    "/mnt/local_nvme/zoubin/cz/self_forcing_predictor_v4_1000_seed0"
)


def read_nonempty(path: Path) -> list[tuple[int, str]]:
    with path.open("r", encoding="utf-8") as handle:
        return [
            (line_number, line.strip())
            for line_number, line in enumerate(handle, start=1)
            if line.strip()
        ]


def atomic_write(path: Path, text: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}")
    with temporary.open("w", encoding="utf-8") as handle:
        handle.write(text)
        handle.flush()
        os.fsync(handle.fileno())
    os.replace(temporary, path)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", type=Path, default=DEFAULT_SOURCE)
    parser.add_argument("--holdout", type=Path, default=DEFAULT_HOLDOUT)
    parser.add_argument("--output_root", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--num_prompts", type=int, default=1000)
    parser.add_argument("--sampling_seed", type=int, default=0)
    parser.add_argument("--inference_seed", type=int, default=0)
    args = parser.parse_args()

    source = args.source.resolve()
    holdout = args.holdout.resolve()
    output_root = args.output_root.resolve()
    if args.num_prompts <= 0:
        raise ValueError("--num_prompts must be positive")
    if args.sampling_seed != 0:
        raise ValueError("the Predictor v4 split is fixed to random.Random(0)")
    if args.inference_seed != 0:
        raise ValueError("every Predictor v4 case must use inference seed 0")
    for path in (source, holdout):
        if not path.is_file():
            raise FileNotFoundError(path)

    source_rows = read_nonempty(source)
    holdout_rows = read_nonempty(holdout)
    holdout_first100 = {prompt for _, prompt in holdout_rows[:100]}
    eligible = [
        (line_number, prompt)
        for line_number, prompt in source_rows
        if prompt not in holdout_first100
    ]
    if len(eligible) < args.num_prompts:
        raise ValueError(
            f"only {len(eligible)} eligible prompts remain after holdout exclusion"
        )
    selected = random.Random(0).sample(eligible, args.num_prompts)
    cases = []
    for case_id, (source_line_number, prompt) in enumerate(selected):
        cases.append(
            {
                "case_id": case_id,
                "prompt": prompt,
                "prompt_sha256": hashlib.sha256(prompt.encode("utf-8")).hexdigest(),
                "source_line_number": source_line_number,
                "seed": 0,
            }
        )

    cases_text = "".join(
        json.dumps(item, ensure_ascii=False, sort_keys=True) + "\n"
        for item in cases
    )
    prompts_text = "".join(f"{item['prompt']}\n" for item in cases)
    config = {
        "schema_version": "self_forcing_predictor_v4_bf16_v1",
        "selection": {
            "algorithm": "random.Random(0).sample",
            "sampling_seed": 0,
            "num_prompts": args.num_prompts,
            "source": str(source),
            "source_nonempty_count": len(source_rows),
            "holdout": str(holdout),
            "holdout_nonempty_prefix_count": min(100, len(holdout_rows)),
            "holdout_exact_prompt_exclusion_count": len(holdout_first100),
            "eligible_count": len(eligible),
        },
        "inference": {
            "seed": 0,
            "reset_seed_per_case": True,
            "num_chunks": 7,
            "frames_per_chunk": 3,
            "num_steps": 4,
            "candidate_block_ids": [0, 1, 28, 29],
            "teacher_checkpoint": "checkpoints/self_forcing_dmd.pt",
            "teacher_checkpoint_key": "generator_ema",
        },
    }
    atomic_write(output_root / "cases.jsonl", cases_text)
    atomic_write(output_root / "selected_prompts.txt", prompts_text)
    atomic_write(
        output_root / "dataset_config.json",
        json.dumps(config, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
    )
    print(
        f"Selected {len(cases)} prompts into {output_root}; "
        f"excluded exact matches against the first {len(holdout_rows[:100])} "
        "non-empty MovieGen prompts."
    )


if __name__ == "__main__":
    main()