"""Split the judged pairs into train / unseen-validation.

The paper holds out a 500-pair validation set and reports on 75 of them.  We hold
out the same *proportion* and, critically, make the split **disjoint by video**:
no clip that appears in a validation pair (as reference or target) may appear
anywhere in training.  A naive random split would leak, because the same clip is
reused across many pairs.
"""

from __future__ import annotations

import argparse
import json
import random
from pathlib import Path


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--pairs", default="data/pairs.jsonl")
    ap.add_argument("--train_out", default="data/pairs_train.jsonl")
    ap.add_argument("--val_out", default="data/pairs_val.jsonl")
    ap.add_argument("--val_frac", type=float, default=0.15)
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()

    rows = [json.loads(l) for l in open(args.pairs)]
    rng = random.Random(args.seed)
    order = list(range(len(rows)))
    rng.shuffle(order)

    # Select validation *pairs* first, then quarantine every clip they touch.
    # (Selecting held-out *videos* first and keeping only pairs fully inside that
    # set is far too restrictive at this dataset size -- it yielded 1 val pair.)
    n_val = max(1, int(len(rows) * args.val_frac))
    val_idx = set(order[:n_val])
    val = [rows[i] for i in sorted(val_idx)]

    val_vids = {r["ref_id"] for r in val} | {r["tgt_id"] for r in val}
    # Train pairs must not touch any held-out clip at all, or physics/appearance
    # from a validation clip leaks into training.
    train = [r for i, r in enumerate(rows)
             if i not in val_idx
             and r["ref_id"] not in val_vids and r["tgt_id"] not in val_vids]
    vids = sorted({r["ref_id"] for r in rows} | {r["tgt_id"] for r in rows})

    for path, data in ((args.train_out, train), (args.val_out, val)):
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w") as f:
            for r in data:
                f.write(json.dumps(r, ensure_ascii=False) + "\n")

    dropped = len(rows) - len(train) - len(val)
    print(f"videos: {len(vids)} ({len(val_vids)} held out)")
    print(f"train pairs: {len(train)}")
    print(f"val pairs:   {len(val)}")
    print(f"dropped (cross-split, would leak): {dropped}")


if __name__ == "__main__":
    main()