#!/usr/bin/env python3
"""Train KenLM 3-gram LM on TLHdig Hitit cuneiform corpus.

Output: tlhdig.arpa + tlhdig.binary (if build_binary available).
Fallback: pure-Python token n-gram counts if KenLM not installed.
"""
import json, argparse, subprocess, shutil, sys
from pathlib import Path
from collections import Counter, defaultdict

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

def iter_corpus(jsonl_path, field='cuneiform', tokenize='char'):
    with open(jsonl_path) as f:
        for line in f:
            try: r = json.loads(line)
            except Exception: continue
            txt = r.get(field)
            if not txt: continue
            if tokenize == 'char':
                # Each cuneiform sign as a token; skip control chars like ▒
                toks = [c for c in txt if c.strip() and c not in ('▒', '…', '·')]
            else:
                toks = txt.split()
            if toks: yield ' '.join(toks)

def try_kenlm(text_path, out_arpa, order=3):
    """Try system kenlm (lmplz binary)."""
    for cand in ['lmplz', '/arf/sw/comp/kenlm/bin/lmplz', shutil.which('lmplz')]:
        if not cand: continue
        try:
            with open(text_path) as fin, open(out_arpa, 'w') as fout:
                r = subprocess.run([cand, '-o', str(order), '--discount_fallback'],
                                   stdin=fin, stdout=fout, timeout=600,
                                   stderr=subprocess.PIPE)
            if r.returncode == 0:
                print(f"KenLM OK via {cand}: {out_arpa}")
                return True
            print(f"KenLM {cand} failed: {r.stderr.decode()[:300]}")
        except Exception as e:
            print(f"KenLM {cand} error: {e}")
    return False

def pure_python_ngram(text_path, out_path, order=3):
    """Fallback: count + save as JSON."""
    ngrams = defaultdict(Counter)
    unigrams = Counter()
    n_lines = 0
    with open(text_path) as f:
        for line in f:
            toks = line.split()
            if not toks: continue
            toks = ['<s>'] + toks + ['</s>']
            unigrams.update(toks)
            for n in range(2, order+1):
                for i in range(len(toks)-n+1):
                    prefix = ' '.join(toks[i:i+n-1])
                    nxt = toks[i+n-1]
                    ngrams[prefix][nxt] += 1
            n_lines += 1
    V = len(unigrams)
    # Export compact n-gram probability table (Laplace smoothing α=1.0)
    probs = {'vocab_size': V, 'unigram': dict(unigrams.most_common()),
             'order': order, 'alpha': 1.0,
             'cond': {pfx: dict(ctr) for pfx, ctr in ngrams.items()}}
    with open(out_path, 'w') as f:
        json.dump(probs, f)
    print(f"Pure-Python {order}-gram: {len(ngrams)} contexts, V={V}, lines={n_lines} → {out_path}")

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--corpus', default=str(ROOT / 'datasets/sources/tlhdig/corpus.jsonl'))
    ap.add_argument('--field', default='cuneiform')
    ap.add_argument('--order', type=int, default=3)
    ap.add_argument('--out-dir', default=str(ROOT / 'hitit_ocr/runs/h100/lm'))
    args = ap.parse_args()

    out_dir = Path(args.out_dir); out_dir.mkdir(parents=True, exist_ok=True)
    text_path = out_dir / 'tlhdig_text.txt'
    with open(text_path, 'w') as g:
        n = 0
        for line in iter_corpus(args.corpus, field=args.field):
            g.write(line + '\n'); n += 1
    print(f"Corpus text: {n} lines → {text_path}")

    arpa = out_dir / f'tlhdig_{args.order}gram.arpa'
    ok = try_kenlm(text_path, arpa, order=args.order)
    if ok:
        # Try build_binary
        for cand in ['build_binary', '/arf/sw/comp/kenlm/bin/build_binary', shutil.which('build_binary')]:
            if not cand: continue
            try:
                bin_path = out_dir / f'tlhdig_{args.order}gram.binary'
                r = subprocess.run([cand, str(arpa), str(bin_path)],
                                   timeout=120, capture_output=True)
                if r.returncode == 0:
                    print(f"build_binary OK: {bin_path}")
                    break
            except Exception: continue
    else:
        print("Falling back to pure-Python n-gram counts")
        pure_python_ngram(text_path, out_dir / f'tlhdig_{args.order}gram.json', order=args.order)

if __name__ == '__main__':
    main()