#!/usr/bin/env python3
"""Build a split summary that preserves RAGTruth's own train/test boundary.

RAGTruth ships a train/test split. That boundary is authoritative and is kept
exactly: every row RAGTruth calls `test` becomes our held-out eval set, and
nothing else does. Reshuffling it would make our numbers incomparable to every
published RAGTruth result, including LLM-Check's.

Our protocol additionally needs a validation split, which RAGTruth does not
provide, for checkpoint and threshold selection. It is carved out of RAGTruth's
*train* portion only, grouped by `source_id`, so the eval set is untouched and
no source document appears on both sides of the train/validation line. Several
responses share one article; splitting by row would leak an article's content
between fitting and selection.

Usage:
    python scripts/make_ragtruth_split.py \
        --input ~/Desktop/datasets/hallucinations/ragtruth_llama2_7b_chat.jsonl \
        --output ~/Desktop/datasets/hallucinations/ragtruth_llama2_7b_chat_split.json \
        --validation-fraction 0.15 --seed 42
"""

from __future__ import annotations

import argparse
import collections
import json
from pathlib import Path
import sys

import numpy as np

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from scripts.make_eval_split import (
    NATIVE_SPLIT_SCHEME,
    build_native_split,
    dataset_binding,
)


def parse_args(argv=None):
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("--input", type=Path, required=True)
    p.add_argument("--output", type=Path, required=True)
    p.add_argument(
        "--validation-fraction",
        type=float,
        default=0.15,
        help="Share of RAGTruth's TRAIN groups held out for validation.",
    )
    p.add_argument("--seed", type=int, default=42)
    return p.parse_args(argv)


def main(argv=None):
    args = parse_args(argv)
    rows = [json.loads(l) for l in args.input.open(encoding="utf-8")]
    if not rows:
        raise SystemExit("empty input")

    # Built by the SAME function validate_split_summary() uses to verify it,
    # so the writer and the checker cannot drift apart.
    built = build_native_split(
        rows,
        seed=args.seed,
        validation_fraction=args.validation_fraction,
        native_split_field="metadata.split",
        eval_value="test",
        grouping_key="prompt",
    )["split"]
    train_indices = built["train_indices"]
    validation_indices = built["validation_indices"]
    eval_indices = built["eval_indices"]

    native_test = {i for i, r in enumerate(rows) if r["metadata"]["split"] == "test"}
    assert set(eval_indices) == native_test, "eval must equal RAGTruth test exactly"
    train_sources = {rows[i]["metadata"]["extra"]["ragtruth_source_id"] for i in train_indices}
    val_sources = {rows[i]["metadata"]["extra"]["ragtruth_source_id"] for i in validation_indices}
    assert not (train_sources & val_sources), "source_id leaked across train/validation"

    summary = {
        "scheme": NATIVE_SPLIT_SCHEME,
        "seed": args.seed,
        "n": len(rows),
        "dataset_binding": dataset_binding(rows),
        "test_fraction": len(eval_indices) / len(rows),
        "validation_fraction": args.validation_fraction,
        "native": {
            "split_field": "metadata.split",
            "eval_value": "test",
            "grouping_key": "prompt",
        },
        "provenance": {
            "eval_split": "RAGTruth's own `test` rows, preserved exactly",
            "validation_split": "carved from RAGTruth's `train` rows only, grouped by source_id",
        },
        "split": {
            "train_indices": train_indices,
            "validation_indices": validation_indices,
            "eval_indices": eval_indices,
        },
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8")

    def h_rate(idx):
        tok = sum(len(rows[i]["labels"]) for i in idx)
        h = sum(sum(1 for v in rows[i]["labels"] if v == 1) for i in idx)
        return tok, h / max(tok, 1)

    print(f"wrote {args.output}")
    for name, idx in (("train", train_indices), ("validation", validation_indices), ("eval(test)", eval_indices)):
        tok, rate = h_rate(idx)
        print(f"  {name:<12} {len(idx):>5} rows  {tok:>8,} tok  H {rate:.2%}")
    print(f"  source groups: train {len(train_sources)}, validation {len(val_sources)}, disjoint")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())