#!/usr/bin/env python3
"""
Grammar-aware beam search decoder for tablet-level sign sequences.

Given a tablet with per-sign top-K candidates (from 4-model ensemble) and their
softmax probs, pick a full-sequence decoding that maximizes:
    score = Σᵢ [log p_cls(sᵢ | xᵢ)  +  λ · log p_LM(sᵢ | s_{i-4..i-1})]

This runs after classification: input is the structured tablet's per-sign top5,
output is the LM-rescored sequence.

Also applies Hittite phonotactic hard constraints (optional, strict mode):
  - Consonant-only sequences are downweighted (Hittite has open syllables)
  - Certain determinatives must precede specific sign categories
"""
import argparse
import json
import math
import pickle
import sys
from collections import defaultdict, Counter
from pathlib import Path


def load_lm(path):
    with open(path, 'rb') as f:
        d = pickle.load(f)
    counts = {k: Counter(v) for k, v in d['counts'].items()}
    ctypes = {k: defaultdict(set, {kk: set(vv) for kk, vv in v.items()})
              for k, v in d['context_types'].items()}
    return dict(counts=counts, ctypes=ctypes, vocab=d['vocab'],
                order=d['order'], D=d['discount'])


class LM:
    def __init__(self, lm):
        self.__dict__.update(lm)
        self.ctx_total = {n: defaultdict(int) for n in range(1, self.order + 1)}
        for n, cnt in self.counts.items():
            for ng, c in cnt.items():
                if n > 1:
                    self.ctx_total[n][ng[:-1]] += c
        self.cont_count = defaultdict(int)
        for ng in self.counts[2]:
            self.cont_count[ng[-1]] += 1
        self.total_cont = sum(self.cont_count.values())

    def cont_prob(self, w):
        return (self.cont_count.get(w, 0) + 1) / (self.total_cont + len(self.vocab) + 1)

    def prob(self, context, word):
        n = min(len(context) + 1, self.order)
        while n > 1:
            ctx = tuple(context[-(n - 1):])
            if self.ctx_total[n].get(ctx, 0) > 0:
                break
            n -= 1
        if n == 1:
            return self.cont_prob(word)
        ctx = tuple(context[-(n - 1):])
        c_ng = self.counts[n].get(ctx + (word,), 0)
        c_ctx = self.ctx_total[n][ctx]
        types = len(self.ctypes[n - 1].get(ctx, ())) if (n - 1) in self.ctypes else 0
        alpha = max(c_ng - self.D, 0) / c_ctx if c_ctx else 0
        lam = (self.D * types) / c_ctx if c_ctx else 1.0
        lower = self.prob(context[-(n - 2):] if n > 2 else (), word)
        return alpha + lam * lower


def beam_decode_line(candidates_per_pos, lm, beam_size=8, lam=0.3):
    """candidates_per_pos: list of list of (sign, prob).
    Returns best sequence (list of signs) and its log-score.
    Beam search: each beam is (tokens, score).
    """
    beams = [([], 0.0)]
    for cand_list in candidates_per_pos:
        new_beams = []
        for tokens, score in beams:
            for sign, prob in cand_list:
                cls_logp = math.log(max(prob, 1e-12))
                # Note: lm.vocab tokens may or may not match sign exactly
                # Use lm token if found, else skip LM contribution
                lm_logp = math.log(max(lm.prob(tuple(tokens[-4:]), sign), 1e-12))
                new_beams.append((tokens + [sign], score + cls_logp + lam * lm_logp))
        # Keep top-B
        new_beams.sort(key=lambda b: -b[1])
        beams = new_beams[:beam_size]
    return beams[0]


def decode_tablet(tablet, lm, beam_size=8, lam=0.3, topk=5):
    """Apply beam decoding within each line of a tablet structure with top5 fields."""
    for L in tablet['lines']:
        cand_list = []
        for s in L['signs']:
            top5 = s.get('top5') or [{'sign': s['sign'], 'prob': s.get('conf', 1.0)}]
            cand_list.append([(t['sign'], t['prob']) for t in top5[:topk]])
        if not cand_list:
            continue
        seq, score = beam_decode_line(cand_list, lm, beam_size=beam_size, lam=lam)
        # Overwrite signs with decoded versions
        for i, s in enumerate(L['signs']):
            if i < len(seq):
                s['sign_before_beam'] = s['sign']
                s['sign'] = seq[i]
        L['beam_score'] = score
    return tablet


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--structure', required=True, help='Tablet JSONL with top5 per sign')
    ap.add_argument('--lm', required=True, help='KN LM pickle')
    ap.add_argument('--output', required=True)
    ap.add_argument('--beam-size', type=int, default=8)
    ap.add_argument('--lambda-lm', type=float, default=0.3)
    ap.add_argument('--topk', type=int, default=5)
    args = ap.parse_args()

    print(f"[load] LM {args.lm}")
    lm = LM(load_lm(args.lm))
    print(f"  vocab={len(lm.vocab)}, order={lm.order}")

    n = 0
    with open(args.structure) as f, open(args.output, 'w') as out:
        for line in f:
            t = json.loads(line)
            decode_tablet(t, lm, beam_size=args.beam_size,
                           lam=args.lambda_lm, topk=args.topk)
            out.write(json.dumps(t, ensure_ascii=False) + "\n")
            n += 1
            if n % 10 == 0:
                print(f"  [{n}]")
    print(f"\nDONE: {n} tablets → {args.output}")


if __name__ == '__main__':
    main()