"""Bug fixes for the canvas data path (runs/CSEL/FIXED_BASE_PREREG.md), imported by the patched scripts/canvas_bot.py.

1. ๆ / ฯ (data audit 2026-09-28: 1,137+ training poems dropped). klon8.prosody deletes ๆ, ฯ and ๏ before G2P, but
   the segmenter keeps them as cells (or glues ๆ to its word), so the cell count never matched the pronounced count
   and the poem was filtered out - while a poet counts ไม้ยมก as a repeated syllable (เก่งๆ = เก่ง เก่ง). prep_line
   drops ฯ / ๏ (end marks, not pronounced); expand_yamok writes ๆ out as a copy of the syllable before it, so the
   cells, the joined text the checker reads, and the pronunciation agree. Approximation (stated): ๆ repeats the
   preceding SYLLABLE; for a multi-syllable word it repeats only its last syllable.
2. Rime / tone labels in context (data audit: canvas_mlm.phon labels each vocabulary type by G2P of the syllable
   ALONE - ส of สมอง gets the letter name, วั no label, มอ of เสมอ the wrong rime; 3.9% of cells mislabelled,
   2.1% rime / 2.8% tone at line ends where the rhyme and tone masks act). ctx_phon_table labels each training cell
   by G2P of its whole line (cells align 1:1 with the line's G2P syllables: load_poems keeps only lines where they
   do) and gives every type its MAJORITY in-context (rime, tone); types never seen in a poem line keep the old
   out-of-context label. apply_phon rebuilds V.rime / V.tone / V.fam / V.n_rime from that table deterministically.
"""
import collections, sys
if "." not in sys.path: sys.path.insert(0, ".")   # run as a script from the repo root
from klon8 import prosody as P


def prep_line(line):
    """ฯ: 896 of 1,130 in the train split follow a word (mostly line/poem-end marks after names and เทอญ; ข้าฯ
    ~34), 225 stand alone, 9 are ฯลฯ. Choice (stated in the prereg): drop ฯ and read the short form (ข้าฯ = ข้า,
    กรุงเทพฯ = กรุงเทพ), NOT ข้าพเจ้า and not a token of its own - klon8.prosody pronounces the short form, and
    the count filter keeps only poems whose metre agrees with that reading. ฯลฯ is dropped whole."""
    return line.replace("ฯลฯ", "").replace("ฯ", "").replace("๏", "")


def expand_yamok(cells):
    """ๆ repeats the syllable before it, wherever it sits in a cell: ดีๆ -> ดี ดี; ๆคน after ดี -> ดี คน
    (review: a leading ๆ repeats the PREVIOUS cell, not the rest of its own cell); ดีๆๆ -> ดี ดี ดี."""
    out = []
    for c in cells:
        for j, part in enumerate(c.split("ๆ")):
            if j and out: out.append(out[-1])
            if part: out.append(part)
    return out


def contract_yamok(line, seg_fn, words=None):
    """Output side (review): write reduplication back as ๆ for the judge, since the human anchors use ๆ.
    Adjacent equal cells X X -> Xๆ (X X X -> Xๆๆ), the exact inverse of expand_yamok, except where XX is itself
    a dictionary word written without ๆ (นานา, จะจะ; pythainlp thai_words). The pipeline measures stutters on the
    written-out text BEFORE contraction, so a contracted stutter is still counted by the gates."""
    if words is None:
        from pythainlp.corpus import thai_words
        words = thai_words()
    out = []
    for c in seg_fn(line):
        if out and out[-1].rstrip("ๆ") == c and c + c not in words: out[-1] += "ๆ"
        else: out.append(c)
    return "".join(out)


def redup_syllables(path="data/modern_bot_train.jsonl"):
    """Syllables the training poems write with ไม้ยมก (the last aligned2 syllable of the word before ๆ)."""
    import re, json, syl_align_v2
    out = set()
    for l in open(path, encoding="utf-8"):
        for w in re.findall(r"([ก-ฮะ-ฺเ-ๅ็-์]+)\s*ๆ", json.loads(l)["poem"]):
            s = syl_align_v2.syls_aligned(w, 12)
            if s: out.add(s[-1])
    return out


def title_cells(seg_fn, text, tl):
    """Title / topic cells for a --yamok model: same ๆ/ฯ handling as the poem lines."""
    return expand_yamok(seg_fn(prep_line(text), tl))[:tl]


def _key(s):
    return (P._MERGED_RIME.get(s.rime, s.rime), s.tone)


def _majority(votes, fallback):
    """Most frequent label; a tie keeps the out-of-context label if it is among the tied, else the smallest by
    repr - so the table does not depend on the order of the poems (train() shuffles them)."""
    top = max(votes.values()); tied = [k for k, n in votes.items() if n == top]
    return fallback if fallback in tied else min(tied, key=repr)


def ctx_phon_table(poems, phon):
    """{syllable: (rime_key, tone)} = majority in-context label; also returns the measured label error rates."""
    votes = collections.defaultdict(collections.Counter)
    cells = []                      # (syllable, in-context label) for the rate measurement
    for p in poems:
        for g in p["grid"]:
            syl = P.g2p_line("".join(g))
            if len(syl) != len(g): continue
            for c, s in zip(g, syl):
                k = _key(s); votes[c][k] += 1; cells.append((c, k))
    table = {c: _majority(v, phon(c)) for c, v in votes.items()}
    old = sum(1 for c, k in cells if phon(c) != k)
    new = sum(1 for c, k in cells if table[c] != k)
    return table, {"cells": len(cells), "out_of_context_mismatch": old / max(len(cells), 1),
                   "type_majority_mismatch": new / max(len(cells), 1), "types": len(table)}


def apply_phon(V, table, phon):
    """V.rime/V.tone/V.fam/V.n_rime from table (fallback: out-of-context phon). Same id order as load_model."""
    rimes = {}; V.rime, V.tone, V.fam = [0] * len(V.itos), [5] * len(V.itos), {}
    for i, t in enumerate(V.itos):
        if i < 4: continue
        r, tn = table.get(t) or phon(t)
        if r is None: continue
        rimes.setdefault(r, len(rimes) + 1); V.rime[i], V.tone[i] = rimes[r], tn
        V.fam.setdefault(rimes[r], []).append(i)
    V.n_rime = len(rimes) + 1


def dump_table(V, table, phon):
    """JSON-safe table over the whole vocabulary (rime keys are tuples)."""
    out = {}
    for t in V.itos[4:]:
        r, tn = table.get(t) or phon(t)
        out[t] = None if r is None else [list(r) if isinstance(r, tuple) else r, tn]
    return out


def load_table(d):
    return {t: ((tuple(v[0]) if isinstance(v[0], list) else v[0]), v[1]) for t, v in d.items() if v is not None}


def freq_table(poems):
    """{"total", "count"} over the poem cells, in the same format as data/syl_freq_train.json (which counts the
    v1 segmenter's units, so most aligned2 syllables would fall back to the floor in --lm-head-init's bias)."""
    cnt = collections.Counter(c for p in poems for g in p["grid"] for c in g)
    return {"total": sum(cnt.values()), "count": dict(cnt)}


if __name__ == "__main__":
    # fixbase_data.py freq DATA OUT: unigram table of the FIXED training cells (aligned2 + ๆ/ฯ), from the patched
    # scripts/canvas_bot.py load_poems, so the table and the model's units cannot disagree
    import json
    assert sys.argv[1] == "freq", __doc__
    sys.path[:0] = [".", "scripts"]
    import canvas_bot as CB
    assert "clean_titles" in CB.load_poems.__code__.co_varnames, "canvas_bot.py is not patched with patch_fixed_base.py"
    d = freq_table(CB.load_poems(sys.argv[2], seg="aligned2", clean_titles=True, yamok=True))
    json.dump(d, open(sys.argv[3], "w", encoding="utf-8"), ensure_ascii=False)
    print(f"  {sys.argv[3]}: {d['total']:,} cells, {len(d['count']):,} types")