#!/usr/bin/env python3
"""Viterbi lattice rescoring with 3-gram LM.

Given per-position top-K classifier (ABZ, logp) lattice, find best path by:
    score(path) = sum_t [ log_p_cls(t) + lambda_lm * log_p_lm(w_t | w_{t-2}, w_{t-1}) ]

Supports both KenLM (binary) and pure-Python JSON n-gram fallback.
"""
import json, math
from pathlib import Path

class _PythonLM:
    def __init__(self, path):
        obj = json.load(open(path))
        self.V = obj['vocab_size']
        self.order = obj['order']
        self.alpha = obj.get('alpha', 1.0)
        self.unigram = obj['unigram']
        self.U = sum(self.unigram.values())
        self.cond = obj['cond']  # prefix_str -> {next: count}

    def _cond_prob(self, ctx, w):
        # Laplace smoothing; back-off to lower order
        while len(ctx) > 0:
            key = ' '.join(ctx)
            if key in self.cond:
                ctr = self.cond[key]
                tot = sum(ctr.values())
                c = ctr.get(w, 0)
                return (c + self.alpha) / (tot + self.alpha * self.V)
            ctx = ctx[1:]
        c = self.unigram.get(w, 0)
        return (c + self.alpha) / (self.U + self.alpha * self.V)

    def score(self, context_tokens, w):
        ctx = context_tokens[-(self.order - 1):]
        return math.log(max(1e-20, self._cond_prob(ctx, w)))

def load_lm(path):
    path = Path(path)
    if path.suffix in ('.binary', '.arpa'):
        try:
            import kenlm
            return kenlm.Model(str(path))
        except Exception as e:
            print(f"kenlm unavailable ({e}); try .json instead"); raise
    # JSON fallback
    return _PythonLM(path)

def lm_score(lm, context, w):
    """Unified log-prob interface."""
    if hasattr(lm, 'BaseScore'):  # kenlm.Model
        # kenlm expects full string — use score with bos/eos False
        s = ' '.join(context[-2:] + [w]) if context else w
        return lm.score(s, bos=False, eos=False) * math.log(10)  # kenlm returns log10
    return lm.score(context, w)

def viterbi_rescore(lattice, lm, lambda_lm=0.3, beam=16):
    """
    lattice: list of [(token, log_p_cls), ...] per position (sorted desc by log_p)
    lm: KenLM or _PythonLM
    Returns: (best_path_tokens, total_score)
    """
    if not lattice: return [], 0.0
    # Beam Viterbi: state = context (last 2 tokens), value = best_score + path
    INF = float('-inf')
    beams = [{(): (0.0, [])}]  # index 0 = before pos 0
    for t, candidates in enumerate(lattice):
        candidates = candidates[:beam]
        new_beam = {}
        for ctx, (s_prev, path_prev) in beams[-1].items():
            for w, lp_cls in candidates:
                lp_lm = lm_score(lm, list(ctx), w)
                new_s = s_prev + lp_cls + lambda_lm * lp_lm
                new_ctx = tuple((list(ctx) + [w])[-2:])
                key = new_ctx
                if key not in new_beam or new_beam[key][0] < new_s:
                    new_beam[key] = (new_s, path_prev + [w])
        # Prune beam to top-N
        new_beam = dict(sorted(new_beam.items(), key=lambda kv: kv[1][0], reverse=True)[:beam])
        beams.append(new_beam)
    # Pick best from final beam
    best = max(beams[-1].values(), key=lambda v: v[0])
    return best[1], best[0]

if __name__ == '__main__':
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument('--lm', required=True)
    ap.add_argument('--lattice', required=True, help='JSONL: one list-of-(token,logp) per line')
    ap.add_argument('--output', required=True)
    ap.add_argument('--lambda-lm', type=float, default=0.3)
    args = ap.parse_args()

    lm = load_lm(args.lm)
    with open(args.lattice) as f, open(args.output, 'w') as g:
        for line in f:
            lat = json.loads(line)
            path, score = viterbi_rescore(lat, lm, lambda_lm=args.lambda_lm)
            g.write(json.dumps({'path': path, 'score': score}) + '\n')
    print(f"Rescored → {args.output}")