#!/usr/bin/env python3
"""Rebuild tablet_view_fold with class-stratification.

Problem: current folds are tablet-level; some classes end up in val
with 0 train samples. This caps val accuracy.

Solution: hash(tablet_id) still drives fold to avoid tablet leakage, but
we re-balance so every class appears in train (by promoting isolated val
samples to train). For classes with ≥5 samples we keep fold; for <5 we
force all samples to fold != 0 (i.e., into train).
"""
import json, argparse, shutil
from collections import defaultdict, Counter
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-fold', type=int, default=0)
    ap.add_argument('--min-per-class-train', type=int, default=3,
                    help='Force classes with < N train samples to have 0 val samples')
    args = ap.parse_args()

    records = []
    with open(args.manifest) as f:
        for line in f:
            r = json.loads(line)
            records.append(r)

    # Classify
    per_class_all = defaultdict(list)
    for r in records:
        if r.get('task') != 'classification': continue
        if not r.get('unified_label'): continue
        per_class_all[r['unified_label']].append(r)

    # Current train/val distribution
    cls_train_count = Counter()
    cls_val_count = Counter()
    for label, recs in per_class_all.items():
        for r in recs:
            if r.get('tablet_view_fold', 0) == args.val_fold:
                cls_val_count[label] += 1
            else:
                cls_train_count[label] += 1

    val_only = [l for l in per_class_all if cls_train_count[l] == 0]
    thin_train = [l for l in per_class_all
                  if cls_train_count[l] < args.min_per_class_train]

    print(f"Classes total: {len(per_class_all)}")
    print(f"Val-only (0 train): {len(val_only)}")
    print(f"Thin train (<{args.min_per_class_train}): {len(thin_train)}")

    # Promote: for thin classes, flip all records to non-val-fold (fold=1)
    flipped = 0
    promote = set(thin_train + val_only)
    for r in records:
        if r.get('task') != 'classification': continue
        if r.get('unified_label') in promote and r.get('tablet_view_fold', 0) == args.val_fold:
            r['tablet_view_fold'] = 1
            r['stratified_promoted'] = True
            flipped += 1
    print(f"Flipped val→train: {flipped}")

    # Write
    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
    with open(args.output, 'w') as f:
        for r in records:
            f.write(json.dumps(r) + '\n')

    # Report post
    cls_train_post = Counter()
    cls_val_post = Counter()
    for r in records:
        if r.get('task') != 'classification': continue
        if not r.get('unified_label'): continue
        if r.get('tablet_view_fold', 0) == args.val_fold:
            cls_val_post[r['unified_label']] += 1
        else:
            cls_train_post[r['unified_label']] += 1
    post_val_only = [l for l in per_class_all if cls_train_post[l] == 0]
    print(f"\nAfter stratification:")
    print(f"  Val-only: {len(post_val_only)}")
    print(f"  Val size: {sum(cls_val_post.values())}")
    print(f"  Train size: {sum(cls_train_post.values())}")
    print(f"Output: {args.output}")

if __name__ == '__main__':
    main()