#!/usr/bin/env python3
"""Random stratified 80/20 split (literature standard: DeepScribe, CuReD).

Replaces tablet-level fold with per-class stratified random split.
Each class: 20% go to val fold (0), 80% train (fold=1).
"""
import json, argparse, random
from collections import defaultdict
from pathlib import Path

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--manifest', required=True)
    ap.add_argument('--output', required=True)
    ap.add_argument('--val-frac', type=float, default=0.20)
    ap.add_argument('--seed', type=int, default=42)
    args = ap.parse_args()

    rng = random.Random(args.seed)
    records = [json.loads(l) for l in open(args.manifest)]

    # Group by class
    per_class = defaultdict(list)
    other = []
    for i, r in enumerate(records):
        if r.get('task') != 'classification' or not r.get('unified_label'):
            other.append(i); continue
        per_class[r['unified_label']].append(i)

    # Shuffle + split each class
    val_indices = set()
    for cls, idxs in per_class.items():
        rng.shuffle(idxs)
        n_val = max(1, int(len(idxs) * args.val_frac))
        if len(idxs) == 1:
            # 1-sample classes go to train (else no train data)
            continue
        val_indices.update(idxs[:n_val])

    # Write with new random_stratified_fold field (0=val, 1=train)
    n_val = n_train = 0
    with open(args.output, 'w') as f:
        for i, r in enumerate(records):
            if r.get('task') == 'classification' and r.get('unified_label'):
                if i in val_indices:
                    r['random_stratified_fold'] = 0
                    r['tablet_view_fold'] = 0   # reuse field so downstream code works
                    n_val += 1
                else:
                    r['random_stratified_fold'] = 1
                    r['tablet_view_fold'] = 1
                    n_train += 1
            f.write(json.dumps(r) + '\n')
    print(f"Random stratified: train={n_train} val={n_val} (val_frac={args.val_frac})")
    print(f"Output: {args.output}")

if __name__ == '__main__':
    main()