#!/usr/bin/env python3
"""Cross-source training manifest: merge Hitit + overlapping labels from ebl_ocr,
old_babylonian_signs, deepscribe. Only keep records whose unified_label is in Hitit label set.

Output manifest:
- All Hittite records kept as-is (tablet_view_fold preserved)
- Matching Akkadian/Sumerian/OB records added as training fold with `cross_source=True`
- Val (fold 0) remains pure Hittite
"""
import json, argparse
from pathlib import Path
from collections import Counter

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

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--hitit-manifest', required=True)
    ap.add_argument('--output', required=True)
    ap.add_argument('--sources', nargs='+', default=['ebl_ocr', 'old_babylonian_signs', 'deepscribe'],
                    help='Source directory names under datasets/sources/')
    ap.add_argument('--val-fold', type=int, default=0)
    ap.add_argument('--cap-per-source', type=int, default=0,
                    help='Max records per source (0 = unlimited)')
    ap.add_argument('--cap-per-label', type=int, default=500,
                    help='Per-label cap across sources to prevent head inflation')
    args = ap.parse_args()

    # 1. Hitit label set
    h_labels = set()
    n_hitit = 0
    with open(args.output, 'w') as out:
        with open(args.hitit_manifest) as f:
            for line in f:
                r = json.loads(line)
                out.write(line)
                if r.get('task') == 'classification' and r.get('unified_label'):
                    h_labels.add(r['unified_label'])
                    n_hitit += 1
    print(f"Hitit labels: {len(h_labels)}, Hitit records: {n_hitit}")

    # 2. Cross-source label count (for cap)
    per_label = Counter()
    added_per_src = Counter()
    with open(args.output, 'a') as out:
        for src in args.sources:
            src_path = ROOT / f'datasets/sources/{src}/manifest.jsonl'
            if not src_path.exists():
                print(f"SKIP: {src_path} not found"); continue
            n_added = 0
            with open(src_path) as f:
                for line in f:
                    try: r = json.loads(line)
                    except: continue
                    if r.get('task') != 'classification': continue
                    lab = r.get('unified_label')
                    if not lab or lab not in h_labels: continue
                    if not r.get('path') or r.get('storage') != 'fs': continue
                    if r.get('integrity_ok') is False: continue
                    if per_label[lab] >= args.cap_per_label: continue
                    # Mark as cross-source, train-only (fold != val_fold)
                    r['cross_source'] = True
                    r['cross_source_origin'] = src
                    r['tablet_view_fold'] = 1  # train split
                    out.write(json.dumps(r) + '\n')
                    per_label[lab] += 1
                    n_added += 1
                    if args.cap_per_source and n_added >= args.cap_per_source: break
            print(f"  {src}: added {n_added}")
            added_per_src[src] = n_added

    # Stats
    total_cross = sum(added_per_src.values())
    print(f"\nTotal cross-source: {total_cross}")
    print(f"Per-source: {dict(added_per_src)}")
    print(f"Output: {args.output}")

if __name__ == '__main__':
    main()