#!/usr/bin/env python3
"""Combine all text/transliteration sources into a single sign-sequence corpus
for BERT-style Masked Sign Language Model pretraining.

Sources aggregated:
  - tlhdig/corpus.jsonl (Hittite transliteration, primary)
  - cuneiml/manifest.jsonl (Sumerian/Akkadian, larger but different script era)
  - transliterated_fragments/manifest.jsonl (mixed period fragments)

Output: JSONL of {source, tablet_id, signs: [str,...], lang, period}
Records whose sign sequence has no token in label_to_idx are skipped.
"""
import json, argparse, re
from pathlib import Path
from collections import Counter

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


def _norm_tlh_token(t):
    """Strip determinatives/numerics from TLHdig transliteration tokens.
    'URUne-ri' → 'ne'; strip leading 'URU/DINGIR/GIŠ/KUR/GÍŠ' + subscripts.
    Splits on hyphens and returns list of normalized tokens."""
    if not t: return []
    # Remove digit subscripts and diacritic marks for matching
    cleaned = re.sub(r'[₀-₉0-9]', '', t)
    # Split on hyphens → each cuneiform sign
    parts = [p.strip() for p in cleaned.split('-') if p.strip()]
    # Drop bracketed determinatives like <D>, [x], etc.
    parts = [p for p in parts if not re.match(r'^[<\[\]]', p)]
    # Strip leading determinatives (URU, DINGIR, GIŠ, KUR prefixes merged with syllable)
    out = []
    for p in parts:
        m = re.match(r'^(URU|DINGIR|GIŠ|KUR|GÍŠ|MUNUS|LÚ)(.+)$', p)
        if m: out.append(m.group(2))
        else: out.append(p)
    return out


def iter_tlhdig(path, l2i):
    with open(path) as f:
        for line in f:
            try: r = json.loads(line)
            except: continue
            words = r.get('words') or []
            signs = []
            for w in words:
                if not isinstance(w, dict): continue
                for tok in _norm_tlh_token(w.get('text') or ''):
                    if tok in l2i:
                        signs.append(tok)
            if len(signs) < 3: continue
            yield {
                'source': 'tlhdig',
                'tablet_id': r.get('tablet') or '',
                'signs': signs,
                'lang': r.get('lang') or 'hit',
                'period': 'Hittite',
            }


def iter_cuneiml(path, l2i):
    """cuneiml stores sign sequences inside extra.text.{obverse,reverse}[*].sign."""
    with open(path) as f:
        for line in f:
            try: r = json.loads(line)
            except: continue
            extra = r.get('extra') or {}
            text = extra.get('text') or {}
            all_signs = []
            for side in ('obverse', 'reverse'):
                for row in text.get(side) or []:
                    if not isinstance(row, dict): continue
                    for s in row.get('sign') or []:
                        if not isinstance(s, str): continue
                        for c in s:
                            if c in l2i:
                                all_signs.append(c)
            if len(all_signs) < 3: continue
            yield {
                'source': 'cuneiml',
                'tablet_id': r.get('tablet_id') or '',
                'signs': all_signs,
                'lang': r.get('language') or 'sum',
                'period': r.get('period') or 'Ur-III',
            }


def iter_transliterated_fragments(path, l2i):
    with open(path) as f:
        for line in f:
            try: r = json.loads(line)
            except: continue
            extra = r.get('extra') or {}
            # fragments use similar structure — try multiple shapes
            sign_seq = extra.get('signs') or extra.get('sign_sequence') or []
            if not sign_seq:
                cun = extra.get('cuneiform') or r.get('cuneiform') or ''
                sign_seq = [c for c in cun if c in l2i]
            else:
                sign_seq = [s for s in sign_seq if s in l2i]
            if len(sign_seq) < 3: continue
            yield {
                'source': 'transliterated_fragments',
                'tablet_id': r.get('tablet_id') or '',
                'signs': sign_seq,
                'lang': r.get('language') or 'akk',
                'period': r.get('period') or 'mixed',
            }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--label-to-idx', required=True)
    ap.add_argument('--output', required=True)
    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"Hittite label vocab: {len(l2i)}")

    sources = [
        ('tlhdig', ROOT / 'datasets/sources/tlhdig/corpus.jsonl', iter_tlhdig),
        ('cuneiml', ROOT / 'datasets/sources/cuneiml/manifest.jsonl', iter_cuneiml),
        ('transliterated_fragments',
         ROOT / 'datasets/sources/transliterated_fragments/manifest.jsonl',
         iter_transliterated_fragments),
    ]

    total = 0
    stats = Counter()
    with open(args.output, 'w') as out:
        for name, path, iterfn in sources:
            if not path.exists():
                print(f"  {name}: missing — skip")
                continue
            n = 0
            for rec in iterfn(path, l2i):
                out.write(json.dumps(rec, ensure_ascii=False) + '\n')
                n += 1
            stats[name] = n
            total += n
            print(f"  {name}: {n} records")

    print(f"\nTotal text-corpus records: {total}")
    print(f"Saved → {args.output}")


if __name__ == '__main__':
    main()