#!/usr/bin/env python3
"""5-fold CV wrapper around eval_ensemble_v2.

For each fold (0..4), invoke eval_ensemble_v2 with that as val. Collect
means + stds. Fold 4 is held as test-set; reported separately.
"""
import json, argparse, subprocess, sys, time
from pathlib import Path
import statistics

ROOT = Path("/arf/scratch/stakan/hitit-proje")

def log(m): print(f"[{time.strftime('%H:%M:%S')}] {m}", flush=True)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--ckpts', nargs='+', required=True)
    ap.add_argument('--manifest', required=True)
    ap.add_argument('--folds', nargs='+', type=int, default=[0, 1, 2, 3])
    ap.add_argument('--test-fold', type=int, default=4)
    ap.add_argument('--scales', nargs='+', type=int, default=[224, 320, 384])
    ap.add_argument('--output-dir', required=True)
    args = ap.parse_args()

    out_dir = Path(args.output_dir); out_dir.mkdir(parents=True, exist_ok=True)

    fold_results = {}
    for f in args.folds + [args.test_fold]:
        out_file = out_dir / f"fold_{f}.json"
        cmd = [sys.executable, str(ROOT / 'hitit_ocr/src/eval_ensemble_v2.py'),
                '--ckpts'] + args.ckpts + [
                '--manifest', args.manifest,
                '--val-fold', str(f),
                '--scales'] + [str(s) for s in args.scales] + [
                '--ensemble-space', 'logit',
                '--temperature-scale',
                '--optimize-weights',
                '--output', str(out_file)]
        log(f"=== Fold {f} ===")
        try: subprocess.run(cmd, check=True, cwd=str(ROOT))
        except Exception as e:
            log(f"Fold {f} failed: {e}"); continue
        fold_results[f] = json.load(open(out_file))

    # Aggregate
    def _collect(key):
        return [fold_results[f].get(key) for f in args.folds
                if f in fold_results and fold_results[f].get(key) is not None]

    val_top1 = _collect('ensemble_top1')
    summary = {
        'cv_folds': args.folds,
        'test_fold': args.test_fold,
        'val_top1_mean': statistics.mean(val_top1) if val_top1 else None,
        'val_top1_std': statistics.stdev(val_top1) if len(val_top1) > 1 else 0,
        'val_top1_per_fold': {f: fold_results[f].get('ensemble_top1')
                               for f in args.folds if f in fold_results},
        'test_top1': fold_results.get(args.test_fold, {}).get('ensemble_top1'),
        'val_optimized_top1_mean': statistics.mean(
            [x for x in _collect('optimized_top1') if x is not None]) if _collect('optimized_top1') else None,
    }
    json.dump(summary, open(out_dir / 'summary.json', 'w'), indent=2)
    log(f"CV summary → {out_dir / 'summary.json'}")
    log(f"  val mean: {summary['val_top1_mean']:.4f} ± {summary['val_top1_std']:.4f}")
    log(f"  test (hold-out): {summary['test_top1']:.4f}")

if __name__ == '__main__':
    main()