#!/usr/bin/env python3
"""v12 ULTIMATE cross-source manifest.

Integrates EVERY valid source:
  1. Hitit original (17k) + aug (5k)
  2. Pseudo from unsup cluster (36k)
  3. Pseudo iter2 from ensemble (3k)
  4. eBL classification overlap (43k, cap 500/cls)
  5. OB zip after extraction (55k → capped)
  6. DeepScribe classification overlap (3k)
  7. eBL LMU 158k metadata (new!) — filter for Hitit-label overlap
  8. compvis bbox extracted — mzl→abz mapping where possible
"""
import json, argparse, os, re
from pathlib import Path
from collections import Counter

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


def load_mzl_abz_map(path=None):
    """Load MZL→ABZ mapping JSON built from Oracc OSL.

    Returns (mzl_to_sign, mzl_to_abz) where:
      - mzl_to_sign: "839" → "KUR" (canonical sign name, Hitit-compatible)
      - mzl_to_abz:  "839" → "470" (ABZ/ABZL number)
    """
    if path is None:
        path = ROOT / 'hitit_ocr/data/mzl_abz_map.json'
    if not Path(path).exists(): return {}, {}
    with open(path) as f:
        d = json.load(f)
    return d.get('mzl_to_sign', {}), d.get('mzl_to_abz', {})


def extract_compvis_bboxes(label_to_idx, cap_per_label=200, mzl_map=None):
    """Extract sign bboxes from compvis manifest to per-bbox classification records.

    Uses the MZL→ABZ sign mapping built from Oracc OSL to translate
    compvis `mzl_label` numerical codes (e.g. "839") into canonical sign names
    (e.g. "KUR"). Only records whose mapped sign name exists in `label_to_idx`
    are kept.
    """
    mzl_to_sign = mzl_map or {}
    src_path = ROOT / 'datasets/sources/compvis/manifest.jsonl'
    if not src_path.exists(): return []
    out = []
    per_label = Counter()
    n_mapped = 0; n_unmapped = 0
    for line in open(src_path):
        r = json.loads(line)
        if r.get('task') != 'detection': continue
        extra = r.get('extra') or {}
        bboxes = extra.get('bboxes', [])
        # Find image path. compvis images live in cdli/<tablet_id>.jpg;
        # manifest stores image_name like "P334926_Obv" (strip _Obv/_Rev).
        img_path = r.get('path')
        if img_path is None or not os.path.exists(img_path):
            iname = r.get('image_name') or ''
            tid = r.get('tablet_id') or ''
            base = re.sub(r'_(Obv|Rev|obverse|reverse)$', '', iname)
            cand_names = [base, tid, iname]
            img_path = None
            for name in cand_names:
                if not name: continue
                for sub in ('cdli', 'saa05', 'test'):
                    c = ROOT / f'datasets/sources/compvis/cuneiform-sign-detection-dataset/images/{sub}/{name}.jpg'
                    if c.exists():
                        img_path = str(c); break
                if img_path: break
        if not img_path or not os.path.exists(img_path): continue

        for b in bboxes:
            mzl = str(b.get('mzl_label', '')).strip()
            if not mzl: continue
            # Try (1) direct match, (2) MZL-map translation to sign name
            if mzl in label_to_idx:
                lab = mzl
            else:
                lab = mzl_to_sign.get(mzl) or mzl_to_sign.get(mzl.lstrip('0'))
                if not lab or lab not in label_to_idx:
                    n_unmapped += 1; continue
                n_mapped += 1
            if per_label[lab] >= cap_per_label: continue
            bb = b.get('bbox') or b.get('relative_bbox')
            if not bb or len(bb) != 4: continue
            out.append({
                'task': 'classification',
                'storage': 'fs',
                'path': img_path,
                'bbox': bb,
                'unified_label': lab,
                'cross_source': True,
                'cross_source_origin': 'compvis_bbox',
                'tablet_view_fold': 1,
                'integrity_ok': True,
            })
            per_label[lab] += 1
    print(f"  [compvis] bbox records: {len(out)} (mzl-mapped: {n_mapped}, unmapped: {n_unmapped})")
    return out


def extract_yeni_veri_bboxes(label_to_idx, cap_per_label=500):
    """yeni_veri: 38 Hittite tablets with YOLO bboxes + class_name (ABZ-compatible).
    Each bbox becomes a classification record (native Hittite)."""
    src_path = ROOT / 'datasets/sources/yeni_veri/manifest.jsonl'
    if not src_path.exists(): return []
    out = []; per_label = Counter()
    for line in open(src_path):
        try: r = json.loads(line)
        except: continue
        if r.get('task') != 'detection': continue
        img_path = r.get('path') or ''
        if img_path and not os.path.isabs(img_path):
            img_path = str(ROOT / img_path)
        if not os.path.exists(img_path): continue
        extra = r.get('extra') or {}
        for b in extra.get('bboxes', []):
            lab = b.get('class_name')
            if not lab or lab not in label_to_idx: continue
            if per_label[lab] >= cap_per_label: continue
            yb = b.get('yolo_bbox')
            if not yb or len(yb) != 4: continue
            out.append({
                'task': 'classification',
                'storage': 'fs',
                'path': img_path,
                'bbox': yb,
                'bbox_format': 'yolo',
                'unified_label': lab,
                'cross_source': True,
                'cross_source_origin': 'yeni_veri',
                'tablet_view_fold': 1,
                'integrity_ok': True,
                'period': 'Hittite',
            })
            per_label[lab] += 1
    print(f"  [yeni_veri] bbox records: {len(out)}")
    return out


def extract_maicubeda(label_to_idx, cap_per_label=300):
    """maicubeda: 27.5k classification crops (char-level).
    Filter by label overlap, keep fs storage only."""
    src_path = ROOT / 'datasets/sources/maicubeda/manifest_classification.jsonl'
    if not src_path.exists():
        src_path = ROOT / 'datasets/sources/maicubeda/manifest.jsonl'
    if not src_path.exists(): return []
    out = []; per_label = Counter()
    for line in open(src_path):
        try: r = json.loads(line)
        except: continue
        if r.get('task') != 'classification': continue
        if r.get('storage') != 'fs': continue
        if r.get('integrity_ok') is False: continue
        lab = r.get('unified_label')
        if not lab or lab not in label_to_idx: continue
        if per_label[lab] >= cap_per_label: continue
        p = r.get('path')
        if not p or not os.path.exists(p): continue
        out.append({
            'task': 'classification',
            'storage': 'fs',
            'path': p,
            'unified_label': lab,
            'cross_source': True,
            'cross_source_origin': 'maicubeda',
            'tablet_view_fold': 1,
            'integrity_ok': True,
        })
        per_label[lab] += 1
    print(f"  [maicubeda] records: {len(out)}")
    return out


def extract_ebl_lmu(label_to_idx, cap_per_label=200, max_take=80000):
    """Extract Hitit-overlapping records from eBL LMU 158k dataset.

    After tar extraction the jpeg files sit directly in
    datasets/sources/ebl_lmu/<_id>.jpeg (no snippets/ subdir). We also
    normalize signName: eBL uses 'ŠU₂' with Unicode subscripts while Hittite
    labels are plain ASCII ('ŠU2'), so we try both forms.
    """
    meta_path = ROOT / 'datasets/sources/ebl_lmu/metadata.json'
    if not meta_path.exists(): return []
    archive_dirs = [
        ROOT / 'datasets/sources/ebl_lmu',
        ROOT / 'datasets/sources/ebl_lmu/snippets',
        ROOT / 'datasets/sources/ebl_ocr/v2_20251219/snippets',
    ]
    archive_dir = next((d for d in archive_dirs if d.exists()), None)
    if archive_dir is None: return []

    # Unicode subscript → ASCII digit map for signName normalization
    sub_map = str.maketrans('₀₁₂₃₄₅₆₇₈₉', '0123456789')

    with open(meta_path) as f:
        meta = json.load(f)
    if not isinstance(meta, list): return []

    out = []; per_label = Counter()
    n_label_miss = 0; n_file_miss = 0
    for r in meta:
        raw_sn = r.get('signName') or ''
        # Try (1) raw, (2) ASCII-subscript form, (3) uppercase
        for cand in (raw_sn, raw_sn.translate(sub_map), raw_sn.upper().translate(sub_map)):
            if cand in label_to_idx:
                sn = cand; break
        else:
            n_label_miss += 1; continue
        if per_label[sn] >= cap_per_label: continue
        if len(out) >= max_take: break
        sid = r.get('_id', '')
        # Expected file: {archive_dir}/{_id}.jpeg (tar stored them flat)
        p_candidates = [archive_dir / f'{sid}.jpeg', archive_dir / f'{sid}.jpg']
        p = next((x for x in p_candidates if x.exists()), None)
        if p is None:
            n_file_miss += 1; continue
        out.append({
            'task': 'classification',
            'storage': 'fs',
            'path': str(p),
            'unified_label': sn,
            'cross_source': True,
            'cross_source_origin': 'ebl_lmu',
            'tablet_view_fold': 1,
            'integrity_ok': True,
            'period': r.get('script') or 'Mesopotamian',
        })
        per_label[sn] += 1
    print(f"  [ebl_lmu] records: {len(out)} (label_miss={n_label_miss}, file_miss={n_file_miss})")
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--base-manifest', required=True,
                    help='Base Hitit manifest to extend')
    ap.add_argument('--label-to-idx', required=True, help='v4 ckpt for label set')
    ap.add_argument('--output', required=True)
    ap.add_argument('--include-pseudo', action='store_true', default=True)
    ap.add_argument('--cap-per-label', type=int, default=500)
    ap.add_argument('--mzl-map', default=str(ROOT / 'hitit_ocr/data/mzl_abz_map.json'))
    args = ap.parse_args()

    import torch
    ck = torch.load(args.label_to_idx, map_location='cpu', weights_only=False)
    l2i = ck['label_to_idx']
    print(f"Hitit classes: {len(l2i)}")

    mzl_to_sign, _ = load_mzl_abz_map(args.mzl_map)
    print(f"MZL→sign mappings: {len(mzl_to_sign)}")

    added = Counter()
    with open(args.output, 'w') as out:
        # Base
        with open(args.base_manifest) as f:
            for line in f:
                out.write(line); added['base'] += 1
        # Already integrated pseudo manifests (check they exist and not already in base)
        for src_name, src_path in [
            ('unsup_cluster', ROOT / 'hitit_ocr/runs/h100/unsup_cluster_v4/pseudo_labels.jsonl'),
            ('pseudo_iter2_ebl', ROOT / 'datasets/sources/hitit_local/manifest_pseudo_iter2_ebl.jsonl'),
        ]:
            if not src_path.exists(): continue
            with open(src_path) as f:
                for line in f:
                    try: r = json.loads(line)
                    except: continue
                    if r.get('unified_label') not in l2i: continue
                    if r.get('tablet_view_fold', 0) == 0: r['tablet_view_fold'] = 1
                    out.write(json.dumps(r) + '\n'); added[src_name] += 1

        # Cross-source from previously-built v11 (Hitit + eBL + OB + DS)
        # Reuse build logic: read each source, filter by overlap
        for src in ['ebl_ocr', 'old_babylonian_signs', 'deepscribe']:
            # Check both regular and fs-rewritten manifest
            src_mf = ROOT / f'datasets/sources/{src}/manifest_fs.jsonl'
            if not src_mf.exists():
                src_mf = ROOT / f'datasets/sources/{src}/manifest.jsonl'
            if not src_mf.exists(): continue
            per_label = Counter()
            with open(src_mf) 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 l2i: continue
                    if r.get('storage') != 'fs' or not r.get('path'): continue
                    if r.get('integrity_ok') is False: continue
                    if per_label[lab] >= args.cap_per_label: continue
                    r['cross_source'] = True; r['cross_source_origin'] = src
                    r['tablet_view_fold'] = 1
                    out.write(json.dumps(r) + '\n'); added[src] += 1
                    per_label[lab] += 1

        # eBL LMU 158k
        ebl_rows = extract_ebl_lmu(l2i, cap_per_label=args.cap_per_label, max_take=30000)
        for r in ebl_rows:
            out.write(json.dumps(r) + '\n')
        added['ebl_lmu'] = len(ebl_rows)

        # compvis bbox with MZL→ABZ map
        cv_rows = extract_compvis_bboxes(l2i, cap_per_label=args.cap_per_label,
                                          mzl_map=mzl_to_sign)
        for r in cv_rows:
            out.write(json.dumps(r) + '\n')
        added['compvis_bbox'] = len(cv_rows)

        # yeni_veri: Hitit tablet bboxes with native ABZ class_name
        yv_rows = extract_yeni_veri_bboxes(l2i, cap_per_label=args.cap_per_label)
        for r in yv_rows:
            out.write(json.dumps(r) + '\n')
        added['yeni_veri'] = len(yv_rows)

        # maicubeda: char-level classification crops
        mai_rows = extract_maicubeda(l2i, cap_per_label=args.cap_per_label)
        for r in mai_rows:
            out.write(json.dumps(r) + '\n')
        added['maicubeda'] = len(mai_rows)

    print(f"\nManifest v12 breakdown:")
    total = 0
    for k, v in added.items():
        print(f"  {k}: {v}")
        total += v
    print(f"  TOTAL: {total}")
    print(f"Output: {args.output}")

if __name__ == '__main__':
    main()