#!/usr/bin/env python3
"""Build one split per RAGTruth task type for a train/test/transfer experiment.

RAGTruth carries three genuinely different generation settings -- answering from
retrieved passages (QA), summarising news (Summary), and writing prose from
structured data (Data2txt). Training on one and evaluating on all three is the
design Lookback Lens used to show that hidden-state probes overfit their
training distribution: 100.0 train / 97.4 in-domain test / 58.3 transfer, versus
88.3 / 87.1 / 66.1 for attention ratios.

Our RAGTruth numbers so far are entirely in-domain, which is the column where a
hidden-state probe is expected to look strongest and where its weakness is
invisible. This makes the other columns measurable.

Each split restricts train/validation to one task; the eval set stays whole, so
one scoring pass yields both in-domain and transfer numbers once predictions are
sliced by `extra.task_type`. Nothing is re-extracted -- the feature cache is
keyed on row content, not on the split.

Usage:
    python scripts/make_ragtruth_task_splits.py \
        --input ~/Desktop/datasets/hallucinations/ragtruth_llama2_7b_chat.jsonl \
        --out-dir ~/Desktop/datasets/hallucinations/task_splits
"""

from __future__ import annotations

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

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_DOMAIN_SPLIT_SCHEME,
    build_native_split,
    dataset_binding,
)

TASK_FIELD = "metadata.extra.task_type"


def parse_args(argv=None):
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("--input", type=Path, required=True)
    p.add_argument("--out-dir", type=Path, required=True)
    p.add_argument("--validation-fraction", type=float, default=0.15)
    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")]
    tasks = sorted({r["metadata"]["extra"]["task_type"] for r in rows})
    args.out_dir.mkdir(parents=True, exist_ok=True)

    test_by_task = collections.Counter(
        r["metadata"]["extra"]["task_type"]
        for r in rows
        if r["metadata"]["split"] == "test"
    )
    print(f"tasks: {tasks}")
    print(f"test rows per task (the transfer targets): {dict(test_by_task)}\n")

    for task in tasks:
        built = build_native_split(
            rows,
            seed=args.seed,
            validation_fraction=args.validation_fraction,
            native_split_field="metadata.split",
            eval_value="test",
            grouping_key="prompt",
            train_domain_field=TASK_FIELD,
            train_domain_value=task,
        )["split"]

        # Guards: training must be single-task, evaluation must stay whole.
        train_tasks = {
            rows[i]["metadata"]["extra"]["task_type"] for i in built["train_indices"]
        }
        assert train_tasks == {task}, f"train leaked other tasks: {train_tasks}"
        eval_tasks = {
            rows[i]["metadata"]["extra"]["task_type"] for i in built["eval_indices"]
        }
        assert eval_tasks == set(tasks), "eval must span every task for transfer"

        summary = {
            "scheme": NATIVE_DOMAIN_SPLIT_SCHEME,
            "seed": args.seed,
            "n": len(rows),
            "dataset_binding": dataset_binding(rows),
            "test_fraction": len(built["eval_indices"]) / len(rows),
            "validation_fraction": args.validation_fraction,
            "native": {
                "split_field": "metadata.split",
                "eval_value": "test",
                "grouping_key": "prompt",
                "train_domain_field": TASK_FIELD,
                "train_domain_value": task,
            },
            "provenance": {
                "design": "train on one task, evaluate on all -- slice predictions by task_type",
                "train_domain": task,
            },
            "split": built,
        }
        out = args.out_dir / f"train_{task.lower()}.json"
        out.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8")
        print(
            f"  train={task:<9} {len(built['train_indices']):>5} train "
            f"{len(built['validation_indices']):>4} val "
            f"{len(built['eval_indices']):>4} eval(all tasks)  -> {out.name}"
        )
    return 0


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