"""
postprocess_v2 — tầng hậu xử lý nâng cấp cho VOCR pipeline.

Cải tiến so với CharNgramLM (trigram + add-epsilon) trong vocr_pipeline.py:
  1. Char 5-gram với stupid backoff (5→4→3→2→1), điểm CHUẨN HÓA THEO ĐỘ DÀI
     (per-char avg logprob) — không còn thiên vị candidate ngắn.
  2. Kiểm tra ÂM TIẾT tiếng Việt: âm tiết không có trong từ vựng corpus
     ("qưân", "hoc", "trưòng") bị phạt — đúng loại lỗi dấu thanh phổ biến nhất.
  3. Bigram mức âm tiết (word-level) — phân định dấu thanh theo ngữ cảnh
     ("chữa bệnh" vs "chửa bệnh").
  4. Re-rank CÓ ĐIỀU KIỆN: khi CTC đã rất chắc chắn (top1 >= 0.90 và bỏ xa
     top2) thì giữ nguyên — tránh over-correction.

API tương thích chỗ gọi cũ: ViLM().train_from_file(path) ; .rerank(candidates, alpha)
Corpus nhận cả 2 định dạng dòng: "path<TAB>text" hoặc text thuần.
"""
import math
import unicodedata
from collections import Counter

_PUNC = set(".,;:!?-–—_()[]{}\"'“”‘’/\\@#$%&*+=<>~`|«»…0123456789")


def _nfc(s):
    return unicodedata.normalize('NFC', s)


def _syllables(text):
    """Tách chuỗi thành âm tiết (token chữ), bỏ số và dấu câu."""
    out = []
    for tok in text.lower().split():
        tok = tok.strip(''.join(_PUNC & set(tok)))
        if tok and not any(c in _PUNC for c in tok):
            out.append(tok)
    return out


class ViLM:
    """LM ký tự + âm tiết cho tiếng Việt, huấn luyện từ corpus text."""

    def __init__(self, n=5, min_syllable_freq=20):
        self.n = n
        self.min_syllable_freq = min_syllable_freq
        self.char_counts = [Counter() for _ in range(n + 1)]  # index = order
        self.syl_freq = Counter()
        self.syl_bigram = Counter()
        self.vocab = set()
        self.trained = False

    # ---------------- huấn luyện ----------------
    def train_from_file(self, filepath):
        n_lines = 0
        with open(filepath, 'r', encoding='utf-8') as f:
            for line in f:
                line = line.rstrip('\n')
                if not line:
                    continue
                text = line.split('\t', 1)[1] if '\t' in line else line
                self._add_line(_nfc(text))
                n_lines += 1
        self.vocab = {s for s, c in self.syl_freq.items()
                      if c >= self.min_syllable_freq}
        self.trained = n_lines > 0
        return n_lines

    def _add_line(self, text):
        t = '\x02' * (self.n - 1) + text + '\x03'
        for order in range(1, self.n + 1):
            cnt = self.char_counts[order]
            for i in range(len(t) - order + 1):
                cnt[t[i:i + order]] += 1
        syls = _syllables(text)
        self.syl_freq.update(syls)
        for a, b in zip(syls, syls[1:]):
            self.syl_bigram[(a, b)] += 1

    # ---------------- chấm điểm ----------------
    def char_logprob(self, text):
        """Per-char avg logprob, stupid backoff (hệ số 0.4 mỗi bậc lùi)."""
        if not self.trained or not text:
            return 0.0
        t = '\x02' * (self.n - 1) + _nfc(text) + '\x03'
        total_chars = self.char_counts[1]
        n_uni = sum(total_chars.values()) or 1
        lp, n_scored = 0.0, 0
        for i in range(self.n - 1, len(t)):
            score = None
            backoff = 1.0
            for order in range(self.n, 0, -1):
                ctx_start = i - order + 1
                gram = t[ctx_start:i + 1]
                c_gram = self.char_counts[order].get(gram, 0)
                if c_gram > 0:
                    if order == 1:
                        score = backoff * (c_gram / n_uni)
                    else:
                        c_ctx = self.char_counts[order - 1].get(gram[:-1], 0)
                        if c_ctx > 0:
                            score = backoff * (c_gram / c_ctx)
                    if score:
                        break
                backoff *= 0.4
            lp += math.log(score if score else 1e-9)
            n_scored += 1
        return lp / max(n_scored, 1)

    def syllable_score(self, text):
        """(tỷ lệ âm tiết hợp lệ, per-syllable bigram logprob)."""
        syls = _syllables(text)
        if not syls:
            return 1.0, 0.0
        valid = sum(1 for s in syls if s in self.vocab or len(s) <= 1)
        bg = 0.0
        for a, b in zip(syls, syls[1:]):
            c_bg = self.syl_bigram.get((a, b), 0)
            c_a = self.syl_freq.get(a, 0)
            if c_bg > 0 and c_a > 0:
                bg += math.log(c_bg / c_a)
            elif self.syl_freq.get(b, 0) > 0:
                n_syl = sum(self.syl_freq.values()) or 1
                bg += math.log(0.4 * self.syl_freq[b] / n_syl)
            else:
                bg += math.log(1e-7)
        n_bg = max(len(syls) - 1, 1)
        return valid / len(syls), bg / n_bg

    def lm_score(self, text):
        """Điểm LM tổng hợp, đã chuẩn hóa độ dài. Thang ~[-12, 0]."""
        if not text:
            return -12.0
        char_lp = self.char_logprob(text)              # ~[-9, 0]
        valid_frac, syl_bg = self.syllable_score(text)  # [0,1], ~[-16, 0]
        return 0.45 * char_lp + 0.25 * (syl_bg / 2.0) + 3.0 * (valid_frac - 1.0)

    # ---------------- re-rank ----------------
    def rerank(self, candidates, alpha=0.7,
               skip_top1_score=0.90, skip_margin=0.30):
        """Re-rank top-K của CTC. Giữ schema output cũ (ctc_score/lm_score/combined).

        Gating: nếu CTC top1 >= skip_top1_score và bỏ xa top2 >= skip_margin
        thì tin CTC, không cho LM can thiệp (chống over-correction).
        """
        if not candidates or not self.trained:
            return candidates
        cands = sorted(candidates, key=lambda c: c.get('score', 0), reverse=True)
        top1 = cands[0].get('score', 0)
        top2 = cands[1].get('score', 0) if len(cands) > 1 else 0.0
        gated = top1 >= skip_top1_score and (top1 - top2) >= skip_margin

        # softmax LM trong nhóm candidate (temperature 0.25) — so sánh tương đối,
        # tín hiệu LM không bị nén như thang tuyệt đối
        TAU = 0.25
        lms = [self.lm_score(c.get('text', '')) if c.get('text') else -12.0
               for c in cands]
        m = max(lms)
        exps = [math.exp((v - m) / TAU) for v in lms]
        z = sum(exps) or 1.0
        lm_soft = [e / z for e in exps]

        scored = []
        for c, lm01 in zip(cands, lm_soft):
            combined = alpha * c.get('score', 0) + (1 - alpha) * lm01
            scored.append({'text': c.get('text', ''),
                           'ctc_score': round(c.get('score', 0), 4),
                           'lm_score': round(lm01, 4),
                           'combined': round(combined, 4)})
        if not gated:
            scored.sort(key=lambda x: x['combined'], reverse=True)
        return scored