#!/usr/bin/env python3
"""Build v8 ultimate training manifest.

Pipeline:
  1. Start from stratified manifest
  2. Apply confusion-merge map (192 classes)
  3. Drop cleanlab-flagged noisy records (from train only; val untouched)
  4. Append tail-aug classical records
  5. Append unsup-cluster 36k pseudo-labels
"""
import json, argparse
from pathlib import Path

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--base-manifest', required=True)
    ap.add_argument('--merge-map', required=True)
    ap.add_argument('--cleanlab-issues', required=True)
    ap.add_argument('--aug-manifest', default=None,
                    help='Optional: tail-aug manifest to append')
    ap.add_argument('--pseudo-manifest', default=None,
                    help='Optional: pseudo-labels jsonl to append')
    ap.add_argument('--val-fold', type=int, default=0)
    ap.add_argument('--output', required=True)
    args = ap.parse_args()

    # Load merge map
    mm = json.load(open(args.merge_map))['merge_map']
    # Load noisy paths
    cl = json.load(open(args.cleanlab_issues))
    noisy_paths = set(r['path'] for r in cl.get('noisy_records', []))

    def apply_merge(r):
        lab = r.get('unified_label')
        if lab and lab in mm:
            r['unified_label'] = mm[lab]
            r['original_label'] = lab
        return r

    n_total = n_dropped = 0
    with open(args.output, 'w') as out:
        # 1. Base manifest, with merge + cleanlab drop for train
        with open(args.base_manifest) as f:
            for line in f:
                r = json.loads(line)
                n_total += 1
                is_val = r.get('tablet_view_fold', 0) == args.val_fold
                # Drop noisy training samples (keep val intact)
                if not is_val and r.get('path') in noisy_paths:
                    n_dropped += 1; continue
                r = apply_merge(r)
                out.write(json.dumps(r) + '\n')
        # 2. Aug records
        if args.aug_manifest and Path(args.aug_manifest).exists():
            n_aug = 0
            with open(args.aug_manifest) as f:
                for line in f:
                    r = json.loads(line)
                    if not r.get('synthetic'): continue  # only new aug rows
                    r = apply_merge(r)
                    out.write(json.dumps(r) + '\n'); n_aug += 1
            print(f"Appended aug: {n_aug}")
        # 3. Pseudo-labels
        if args.pseudo_manifest and Path(args.pseudo_manifest).exists():
            n_pseudo = 0
            with open(args.pseudo_manifest) as f:
                for line in f:
                    r = json.loads(line)
                    if not r.get('pseudo_label'): continue
                    r = apply_merge(r)
                    out.write(json.dumps(r) + '\n'); n_pseudo += 1
            print(f"Appended pseudo: {n_pseudo}")

    print(f"Base total: {n_total}, dropped as noisy: {n_dropped}")
    print(f"Output: {args.output}")

if __name__ == '__main__':
    main()