"""从 HuggingFace DeepCoder-Preview-Dataset 采样当前 Code demo 使用的训练集。

生成:
  data/ablation_taco.jsonl  — TACO 子集采样，约 4300 样本

为复现实验记录中的 TACO 采样结果，脚本保留历史采样流程的随机状态。

用法:
    python scripts/sample_ablation_subsets.py [--target-size 4300] [--seed 42]
"""
from __future__ import annotations

import argparse
import json
import random
import sys
from pathlib import Path

from datasets import load_dataset

REPO_ROOT = Path(__file__).resolve().parent.parent
DATA_DIR = REPO_ROOT / "data"

NUM_SHORTEST = 3
NUM_STRATIFIED = 3
PER_CASE_CHAR_BUDGET = 6000
TOTAL_CHAR_BUDGET = 18000
HISTORICAL_COMPATIBILITY_SHUFFLE_COUNT = 14338
DATASET_REVISION = "177913a7bd43791646ef6a43645caa3c871ab3db"


def _case_size(c: dict) -> int:
    return len(c.get("input", "")) + len(c.get("output", ""))


def _split_public(ground_truth: dict) -> str:
    """从 ground_truth 中抽取用于 rollout 的 public test cases。"""
    inputs = ground_truth.get("inputs", [])
    outputs = ground_truth.get("outputs", [])
    n = min(len(inputs), len(outputs))
    if n == 0:
        return "[]"

    cases = [
        {"input": inputs[i], "output": outputs[i], "testtype": "stdin"}
        for i in range(n)
    ]
    cases_sorted = sorted(cases, key=_case_size)

    selected: list[dict] = []
    total_chars = 0
    half_end = n // 2

    for c in cases_sorted[:min(NUM_SHORTEST, half_end)]:
        selected.append(c)
        total_chars += _case_size(c)

    pool = [c for c in cases_sorted[NUM_SHORTEST:half_end]
            if _case_size(c) <= PER_CASE_CHAR_BUDGET]
    if pool:
        if len(pool) <= NUM_STRATIFIED:
            picks = pool
        else:
            step = len(pool) / NUM_STRATIFIED
            picks = [pool[int(i * step)] for i in range(NUM_STRATIFIED)]
        for c in picks:
            if total_chars + _case_size(c) > TOTAL_CHAR_BUDGET:
                break
            selected.append(c)
            total_chars += _case_size(c)

    return json.dumps(selected, ensure_ascii=False)


def normalize_taco(row: dict) -> dict | None:
    """将 HF TACO 格式转为统一训练格式。

    HF TACO tests: {"inputs": [...], "outputs": [...]}
    """
    problem = row["problem"]
    tests_raw = row["tests"]
    ground_truth = json.loads(tests_raw) if isinstance(tests_raw, str) else tests_raw

    if not isinstance(ground_truth, dict):
        return None
    if len(ground_truth.get("inputs", [])) < 5:
        return None
    if "<image>" in problem:
        return None

    prompt_content = problem if "Solve the following" in problem else (
        f"Solve the following coding problem using the programming language python:\n\n{problem}"
    )

    return {
        "messages": [{"content": prompt_content, "role": "user"}],
        "rollout_extra": {
            "test_cases": json.dumps(ground_truth, ensure_ascii=False),
            "public_test_cases": _split_public(ground_truth),
            "dataset": "deepcoder",
            "data_source": "taco",
            "difficulty": "none",
        },
    }


def write_jsonl(records: list[dict], path: Path) -> int:
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        for r in records:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    return len(records)


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--target-size", type=int, default=4300,
                    help="TACO 训练集的目标样本数")
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    random.seed(args.seed)
    target = args.target_size

    print("Loading taco subset from HuggingFace...", file=sys.stderr)
    ds_taco = load_dataset(
        "agentica-org/DeepCoder-Preview-Dataset",
        "taco",
        split="train",
        revision=DATASET_REVISION,
    )
    print(f"  raw rows: {len(ds_taco)}", file=sys.stderr)

    print("Normalizing taco...", file=sys.stderr)
    taco_records = []
    for row in ds_taco:
        rec = normalize_taco(row)
        if rec is not None:
            taco_records.append(rec)
    print(f"  valid records: {len(taco_records)}", file=sys.stderr)

    compatibility_rng_state = list(range(HISTORICAL_COMPATIBILITY_SHUFFLE_COUNT))
    random.shuffle(compatibility_rng_state)
    random.shuffle(taco_records)
    taco_subset = taco_records[:target]
    n = write_jsonl(taco_subset, DATA_DIR / "ablation_taco.jsonl")

    print(f"\n{'='*60}")
    print(f"目标规模: {target} samples")
    print(f"TACO 可用: {len(taco_records)} (采样 {n})")
    print(f"输出文件: {DATA_DIR / 'ablation_taco.jsonl'}")
    print("已生成当前实验使用的训练文件。")
    print("\n该文件可直接用于 submit_job.py 的 training_file 参数。")


if __name__ == "__main__":
    main()
