"""Canvas model: the whole บท is generated at once, rhyme and line as one decision (user requirement).

Architecture
  * one token = one orthographic syllable (pythainlp syllable_tokenize); a poem is a fixed grid of
    8 lines x 10 cells = 80 (user spec; 7-10 syllables per line, the rest <pad>), preceded by the title
  * every cell is embedded as syllable + rime + tone (phoneme channels) + line + position + segment
  * a bidirectional Transformer encoder over title + grid: every cell sees every other cell, so a
    rhyme word is chosen with its own line AND the lines it rhymes with in view
  * trained as a masked LM on 31k human poems with a random mask ratio (Mask-Predict style)

Decoding (iterative refinement, T steps)
  * all cells start masked; each step the model predicts all cells at once
  * each rhyme link (a, b) is ONE variable: the rime family maximising p_a(family) * p_b(family) is
    chosen and both cells are restricted to it; tone classes are restricted at line-final cells
  * the most confident cells are committed, the rest re-masked; repeat. Nothing is fitted afterwards.

Usage: canvas_mlm.py train --epochs 20 --out runs/CANVAS_MLM
       canvas_mlm.py decode --model runs/CANVAS_MLM --prompts data/prompts/test.jsonl --out O.jsonl
"""
import os, re, sys, json, math, time, random, argparse, collections
sys.path.insert(0, ".")
import mlx.core as mx, mlx.nn as nn, mlx.optimizers as optim
from klon8 import prosody as P
from pythainlp.tokenize import syllable_tokenize

L, C, TL = 8, 10, 12                     # lines, cells per line (user spec: 8 x 10 = 80 cells), title length
PAD, MASK, UNK, SEP = 0, 1, 2, 3
EVEN, CHECKED = (0, 3), (1, 2, 4)
# rhyme GROUPS (cells that must share one rime family): w2 rhymes with w3 and x4, w6 with w4, w7 and x8.
# Deciding per link would restrict a shared cell to two different families (empty intersection).
GROUPS = [[(1, -1), (2, 2)], [(2, -1), (3, -1), (4, 2)], [(4, -1), (6, -1), (7, -1), (8, 2)], [(5, -1), (6, 2)]]
TONE_AT = {2: CHECKED, 6: CHECKED, 3: EVEN, 4: EVEN, 7: EVEN, 8: EVEN}
_PH = {}
RIME_OF = TONE_OF = None


def phon(syl):
    """(rime family key, tone) of a syllable string, cached; (None, None) if g2p fails."""
    if syl not in _PH:
        g = P.g2p_line(syl)
        _PH[syl] = ((P._MERGED_RIME.get(g[-1].rime, g[-1].rime), g[-1].tone) if g else (None, None))
    return _PH[syl]


def syls_of(text, limit=None):
    out = [s for s in syllable_tokenize(text) if re.search(r"[ก-๙]", s)]
    return out[:limit] if limit else out


class Vocab:
    def __init__(self, tokens, min_freq=3):
        cnt = collections.Counter(tokens)
        self.itos = ["<pad>", "<mask>", "<unk>", "<sep>"] + sorted(t for t, c in cnt.items() if c >= min_freq)
        self.stoi = {t: i for i, t in enumerate(self.itos)}
        rimes = {}
        self.rime, self.tone, self.fam = [0] * len(self.itos), [5] * len(self.itos), {}
        for i, t in enumerate(self.itos):
            if i < 4: continue
            r, tn = phon(t)
            if r is None: continue
            rimes.setdefault(r, len(rimes) + 1)
            self.rime[i], self.tone[i] = rimes[r], tn
            self.fam.setdefault(rimes[r], []).append(i)
        self.n_rime = len(rimes) + 1

    def enc(self, t):
        return self.stoi.get(t, UNK)


def load_poems(path, vocab_min=3):
    rows = [json.loads(l) for l in open(path, encoding="utf-8") if l.strip()]
    poems = []
    for r in rows:
        lines = r["poem"].split("\n")
        if len(lines) != 8: continue
        grid = [syls_of(ln) for ln in lines]
        if all(7 <= len(g) <= C for g in grid):
            poems.append({"title": syls_of(r.get("title", "") or "", TL), "grid": grid})
    return poems


def build_arrays(poems, V):
    """token ids (N, TL+L*C), plus static line/pos/segment ids."""
    N = len(poems); T = TL + L * C
    tok = [[PAD] * T for _ in range(N)]
    for n, p in enumerate(poems):
        for j, s in enumerate(p["title"]): tok[n][j] = V.enc(s)
        for i, line in enumerate(p["grid"]):
            for j, s in enumerate(line): tok[n][TL + i * C + j] = V.enc(s)
    line_id = [0] * TL + [1 + i for i in range(L) for _ in range(C)]
    pos_id = list(range(TL)) + [j for _ in range(L) for j in range(C)]
    seg_id = [0] * TL + [1] * (L * C)
    return mx.array(tok), mx.array(line_id), mx.array(pos_id), mx.array(seg_id)


class Canvas(nn.Module):
    def __init__(self, V, d=256, h=8, layers=6):
        super().__init__()
        n = len(V.itos)
        self.tok = nn.Embedding(n, d); self.rime = nn.Embedding(V.n_rime, d); self.tone = nn.Embedding(6, d)
        self.line = nn.Embedding(L + 1, d); self.pos = nn.Embedding(max(TL, C), d); self.seg = nn.Embedding(2, d)
        self.enc = nn.TransformerEncoder(layers, d, h, mlp_dims=4 * d, dropout=0.1)
        self.norm = nn.LayerNorm(d); self.out = nn.Linear(d, n)

    def __call__(self, tok, line_id, pos_id, seg_id):
        known = (tok != MASK)
        # phoneme lookups are fixed tables, kept outside the module so they are not parameters
        r = mx.where(known, RIME_OF[tok], 0); t = mx.where(known, TONE_OF[tok], 5)
        x = self.tok(tok) + self.rime(r) + self.tone(t) + self.line(line_id)[None] + self.pos(pos_id)[None] + self.seg(seg_id)[None]
        x = self.enc(x, None)
        return self.out(self.norm(x))


def mask_batch(tok, seg_id, rng):
    """Mask a random share (15-90%) of poem cells that are not <pad>; title never masked."""
    B, T = tok.shape
    ratio = mx.array([rng.uniform(0.15, 0.9) for _ in range(B)])[:, None]
    u = mx.random.uniform(shape=(B, T))
    maskable = (seg_id[None, :] == 1) & (tok != PAD)
    m = maskable & (u < ratio)
    return mx.where(m, MASK, tok), m


def train(a):
    mx.set_default_device(mx.cpu if a.cpu else mx.gpu)
    rng = random.Random(0); mx.random.seed(0)
    poems = load_poems(a.data)
    rng.shuffle(poems)
    V = Vocab([s for p in poems for line in p["grid"] for s in line] + [s for p in poems for s in p["title"]])
    os.makedirs(a.out, exist_ok=True)
    json.dump({"itos": V.itos}, open(os.path.join(a.out, "vocab.json"), "w", encoding="utf-8"), ensure_ascii=False)
    valid, train_p = poems[:500], poems[500:]
    tok, line_id, pos_id, seg_id = build_arrays(train_p, V)
    vtok, _, _, _ = build_arrays(valid, V)
    global RIME_OF, TONE_OF
    RIME_OF, TONE_OF = mx.array(V.rime), mx.array(V.tone)
    model = Canvas(V, d=a.dim, h=8, layers=a.layers)
    mx.eval(model.parameters())
    n_par = sum(v.size for _, v in nn.utils.tree_flatten(model.parameters()))
    print(f"poems {len(poems)} (train {len(train_p)}), vocab {len(V.itos)}, rime families {V.n_rime-1}, params {n_par/1e6:.1f}M", flush=True)
    opt = optim.AdamW(learning_rate=a.lr, weight_decay=0.01)

    def loss_fn(model, x, y, m):
        logits = model(x, line_id, pos_id, seg_id)
        ce = nn.losses.cross_entropy(logits.reshape(-1, logits.shape[-1]), y.reshape(-1), reduction="none").reshape(y.shape)
        return (ce * m).sum() / mx.maximum(m.sum(), 1)
    step = nn.value_and_grad(model, loss_fn)
    B = a.batch; t0 = time.time(); it = 0
    for ep in range(a.epochs):
        perm = mx.array(rng.sample(range(tok.shape[0]), tok.shape[0]))
        model.train()
        for i in range(0, tok.shape[0] - B + 1, B):
            y = tok[perm[i:i + B]]
            x, m = mask_batch(y, seg_id, rng)
            loss, grads = step(model, x, y, m.astype(mx.float32))
            opt.update(model, grads); mx.eval(model.parameters(), opt.state); it += 1
            if it % 50 == 0:
                print(f"ep {ep} it {it} loss {loss.item():.3f} {time.time()-t0:.0f}s", flush=True)
        model.eval()
        x, m = mask_batch(vtok, seg_id, random.Random(1))
        vl = loss_fn(model, x, vtok, m.astype(mx.float32)).item()
        print(f"== epoch {ep} valid masked-loss {vl:.3f}", flush=True)
        mx.save_safetensors(os.path.join(a.out, "model.safetensors"), dict(nn.utils.tree_flatten(model.parameters())))
        json.dump({"dim": a.dim, "layers": a.layers}, open(os.path.join(a.out, "config.json"), "w"))


STOP_SYL = set("ที่ ก็ จะ ได้ ไป มา ให้ อยู่ แล้ว และ กับ ว่า ไม่ นี้ นั้น ใน ของ เป็น คน ๆ".split())


def clauses_of(prose, n=4):
    """Cut the prose into n clauses (spaces/punctuation, else equal word chunks), one per บาท."""
    parts = [x for x in re.split(r"\s+|[,;.!?]", prose) if x.strip()]
    if len(parts) < n:
        sy = syls_of(prose)
        if len(sy) <= 8: return [prose.strip()] * n  # a short topic steers every บาท whole
        k = max(1, len(sy) // n)
        parts = ["".join(sy[i:i + k]) for i in range(0, len(sy), k)]
    while len(parts) > n:
        i = min(range(len(parts) - 1), key=lambda j: len(parts[j]) + len(parts[j + 1]))
        parts[i:i + 2] = [parts[i] + parts[i + 1]]
    return parts


def decode_one(model, V, title_syls, line_id, pos_id, seg_id, T=12, temp=0.8, rng=None, line_len=8, fam_temp=0.5, anchors=None, lengths=None, topic=None, topic_boost=0.0, schedule="stanza"):
    """anchors: {(line, cell): syllable} fixed from an LLM draft (topic content); lengths: syllables per line."""
    n = len(V.itos)
    lens = lengths or [line_len] * L
    tok = [PAD] * (TL + L * C)
    for j, s in enumerate(title_syls[:TL]): tok[j] = V.enc(s)
    cells = [(i, j) for i in range(L) for j in range(lens[i])]
    idx = {(i, j): TL + i * C + j for i in range(L) for j in range(C)}
    for (i, j) in cells: tok[idx[(i, j)]] = MASK
    committed = set()
    for cell, syl in (anchors or {}).items():
        if cell in idx and cell[1] < lens[cell[0]] and V.enc(syl) != UNK:
            tok[idx[cell]] = V.enc(syl); committed.add(idx[cell])
    # vocabulary masks: never emit specials or <pad> inside a line; tone classes at line-final cells
    base_bad = mx.array([1.0 if k < 4 else 0.0 for k in range(n)]) * -1e9
    tone_mask = {}
    for ln, allowed in TONE_AT.items():
        tone_mask[ln] = mx.array([0.0 if V.tone[k] in allowed and k >= 4 else -1e9 for k in range(n)])
    fam_vec = {f: mx.array([0.0 if k in set(ks) else -1e9 for k in range(n)]) for f, ks in V.fam.items()}
    groups = [[(l - 1, lens[l - 1] - 1 if c == -1 else c) for (l, c) in g] for g in GROUPS]
    group_family = {}
    # topical steering: each บาท (2 lines) is nudged toward the syllables of its own prompt clause
    topic_vec = None
    if topic and topic_boost:
        topic_vec = mx.zeros((L, n))
        rows_ = []
        for bt, clause in enumerate(topic[:4]):
            ids = [V.enc(sy) for sy in syls_of(clause) if sy not in STOP_SYL and V.enc(sy) != UNK and len(sy) > 1]
            v = [0.0] * n
            for k in ids: v[k] = topic_boost
            rows_ += [v, v]
        while len(rows_) < L: rows_.append([0.0] * n)
        topic_vec = mx.array(rows_[:L])
    for t in range(T):
        x = mx.array([tok])
        logits = model(x, line_id, pos_id, seg_id)[0].astype(mx.float32) + base_bad[None, :]
        for ln, tm in tone_mask.items():
            k = idx[(ln - 1, lens[ln - 1] - 1)]
            logits[k] = logits[k] + tm
        if topic_vec is not None:
            # a prompt syllable earns its bonus once: after it is committed anywhere, the bonus is
            # withdrawn (otherwise the boost outweighs the repetition penalty: แม่แม่แม่แม่)
            tv = topic_vec
            if committed:
                usedk = mx.array(sorted({tok[k] for k in committed if tok[k] >= 4}))
                tv = mx.array(topic_vec); tv[:, usedk] = 0.0
            for i in range(L):
                for j in range(lens[i]):
                    logits[idx[(i, j)]] = logits[idx[(i, j)]] + tv[i]
        lp = nn.log_softmax(logits, axis=-1)
        for gi, cells_g in enumerate(groups):
            if gi in group_family: continue
            done = [idx[c] for c in cells_g if idx[c] in committed]
            if done:
                f = V.rime[tok[done[0]]]
                if f: group_family[gi] = f
                continue
            scores, fams = [], []
            used_f = collections.Counter(group_family.values())
            for f, ks in V.fam.items():
                if len(ks) < 20: continue
                ks_a = mx.array(ks)
                sc = sum(float(mx.logsumexp(lp[idx[c]][ks_a]).item()) for c in cells_g) - 1.0 * used_f.get(f, 0)
                scores.append(sc); fams.append(f)
            if fams:
                z = mx.array(scores) / max(fam_temp, 1e-6)
                group_family[gi] = fams[int(mx.random.categorical(z).item())] if fam_temp > 0 else fams[max(range(len(fams)), key=lambda i: scores[i])]
        for gi, cells_g in enumerate(groups):
            f = group_family.get(gi)
            if f and f in fam_vec:
                for cell in cells_g:
                    if idx[cell] not in committed:
                        logits[idx[cell]] = logits[idx[cell]] + fam_vec[f]
        # no repeated syllable inside a rhyme group
        for cells_g in groups:
            for cell in cells_g:
                if idx[cell] in committed: continue
                for other in cells_g:
                    if other != cell and idx[other] in committed:
                        logits[idx[cell], tok[idx[other]]] = -1e9
        if committed:
            rep = mx.array(sorted({tok[k] for k in committed if tok[k] >= 4}))
            logits[:, rep] = logits[:, rep] - 2.0                  # repetition inside the poem
        probs = mx.softmax(logits / temp, axis=-1)
        conf = mx.max(probs, axis=-1).tolist()
        # commit the most confident open cells: rhyme cells first in the first pass
        if schedule == "line":
            # วรรค-recurrent: only the current line is open; it is decided jointly (all its cells) over
            # T_line steps with every earlier line fixed, then the next line
            cur_line = min(L - 1, t // max(1, T // L))
            open_cells = [idx[c] for c in cells if idx[c] not in committed and c[0] <= cur_line]
            per_line = max(1, T // L)
            target_done = sum(lens[i] for i in range(cur_line + 1)) * min(1.0, ((t % per_line) + 1) / per_line)
            n_commit = max(1, int(math.ceil(target_done)) - len([k for k in committed if k >= TL]))
        else:
            open_cells = [idx[c] for c in cells if idx[c] not in committed]
            n_commit = max(1, math.ceil(len(cells) * (t + 1) / T) - len(committed))
        order = sorted(open_cells, key=lambda k: -conf[k])
        if t == 0:
            rh = [idx[c] for g in groups for c in g if idx[c] not in committed]
            order = list(dict.fromkeys(rh + order))
        for k in order[:n_commit]:
            choice = int(mx.random.categorical(mx.log(probs[k] + 1e-12) / 1.0).item()) if temp > 0 else int(mx.argmax(probs[k]).item())
            tok[k] = choice; committed.add(k)
        for k in open_cells:
            if k not in committed: tok[k] = MASK
    lines = ["".join(V.itos[tok[idx[(i, j)]]] for j in range(lens[i]) if tok[idx[(i, j)]] >= 4) for i in range(L)]
    return "\n".join(lines)


def decode(a):
    mx.set_default_device(mx.cpu if a.cpu else mx.gpu)
    V = Vocab.__new__(Vocab)
    V.itos = json.load(open(os.path.join(a.model, "vocab.json"), encoding="utf-8"))["itos"]
    V.stoi = {t: i for i, t in enumerate(V.itos)}
    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 = 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
    cfg = json.load(open(os.path.join(a.model, "config.json")))
    global RIME_OF, TONE_OF
    RIME_OF, TONE_OF = mx.array(V.rime), mx.array(V.tone)
    model = Canvas(V, d=cfg["dim"], h=8, layers=cfg["layers"])
    model.load_weights(os.path.join(a.model, "model.safetensors"))
    model.eval(); mx.eval(model.parameters())
    _, line_id, pos_id, seg_id = build_arrays([{"title": [], "grid": [[] for _ in range(L)]}], V)
    rows = [json.loads(l) for l in open(a.prompts, encoding="utf-8") if l.strip()]
    if a.limit: rows = rows[:a.limit]
    drafts = {}
    if a.draft:
        for r in map(json.loads, open(a.draft, encoding="utf-8")):
            if r.get("poem") and r["id"] not in drafts:
                drafts[r["id"]] = r["poem"]
    os.makedirs(os.path.dirname(os.path.abspath(a.out)), exist_ok=True)
    mx.random.seed(a.seed)
    with open(a.out, "w", encoding="utf-8") as f:
        for i, r in enumerate(rows, 1):
            title = syls_of(re.split(r"\s+|[,;.!?]", r["text"])[0], TL)
            anchors, lengths = None, None
            if r["id"] in drafts:
                dl = [syls_of(ln) for ln in drafts[r["id"]].split("\n")][:L]
                if len(dl) == L and all(1 <= len(x) <= C for x in dl):
                    lengths = [min(max(len(x), 7), C) for x in dl]
                    prompt_syl = set(syls_of(r["text"]))
                    rhyme_cells = {(l - 1, lengths[l - 1] - 1 if c == -1 else c) for g in GROUPS for (l, c) in g}
                    anchors = {}
                    for li_, line in enumerate(dl):
                        for j, syl in enumerate(line[:lengths[li_]]):
                            if syl in prompt_syl and (li_, j) not in rhyme_cells and len(syl) > 1:
                                anchors[(li_, j)] = syl
                    if a.free_window is not None:            # keep the whole draft except a window around each rhyme cell
                        free = set()
                        for (li_, cj) in rhyme_cells:
                            for dj in range(-a.free_window, a.free_window + 1):
                                free.add((li_, cj + dj))
                        for li_, line in enumerate(dl):
                            for j, syl in enumerate(line[:lengths[li_]]):
                                if (li_, j) not in free and (li_, j) not in rhyme_cells:
                                    anchors[(li_, j)] = syl
                    if a.keep_frac:                          # also keep a share of the draft's other cells
                        pool = [(li_, j, syl) for li_, line in enumerate(dl) for j, syl in enumerate(line[:lengths[li_]])
                                if (li_, j) not in anchors and (li_, j) not in rhyme_cells]
                        random.seed(a.seed + i); random.shuffle(pool)
                        for (li_, j, syl) in pool[: int(a.keep_frac * len(pool))]:
                            anchors[(li_, j)] = syl
            best = None
            for attempt in range(a.retries):                       # joint decoding is cheap: retry until the checker agrees
                poem = decode_one(model, V, title, line_id, pos_id, seg_id, T=a.steps, temp=a.temp, fam_temp=a.fam_temp, anchors=anchors, lengths=lengths,
                                  topic=clauses_of(r["text"]) if a.topic_boost else None, topic_boost=a.topic_boost, schedule=a.schedule)
                s = P.score(poem, n_bot_expected=2)
                key = (s["all_pass"], s["ext_rhyme"] + s["tone"] + 0.5 * s["internal"])
                if best is None or key > best[0]:
                    best = (key, poem, s)
                if s["all_pass"]:
                    break
            _, poem, s = best
            f.write(json.dumps({"id": r["id"], "prompt": r["text"], "sample": 0, "poem": poem,
                                "all_pass": s["all_pass"], "structure": s["structure"], "ext_rhyme": s["ext_rhyme"],
                                "tone": s["tone"], "internal": s["internal"]}, ensure_ascii=False) + "\n"); f.flush()
            print(f"{i}/{len(rows)} all_pass={s['all_pass']} rhyme={s['ext_rhyme']:.2f}", file=sys.stderr, flush=True)
    P.save_cache()


if __name__ == "__main__":
    ap = argparse.ArgumentParser(); sub = ap.add_subparsers(dest="cmd")
    t = sub.add_parser("train"); t.add_argument("--data", default="data/modern_klon8.jsonl"); t.add_argument("--out", default="runs/CANVAS_MLM")
    t.add_argument("--epochs", type=int, default=20); t.add_argument("--batch", type=int, default=64); t.add_argument("--lr", type=float, default=5e-4)
    t.add_argument("--dim", type=int, default=256); t.add_argument("--layers", type=int, default=6); t.add_argument("--cpu", action="store_true")
    d = sub.add_parser("decode"); d.add_argument("--model", default="runs/CANVAS_MLM"); d.add_argument("--prompts", required=True)
    d.add_argument("--out", required=True); d.add_argument("--limit", type=int, default=0); d.add_argument("--steps", type=int, default=12)
    d.add_argument("--temp", type=float, default=0.8); d.add_argument("--seed", type=int, default=0); d.add_argument("--cpu", action="store_true")
    d.add_argument("--fam-temp", type=float, default=0.5, help="temperature for sampling each rhyme link's rime family (0 = argmax)")
    d.add_argument("--retries", type=int, default=4, help="re-decode a prompt until the checker passes (keeps the best attempt)")
    d.add_argument("--draft", default=None, help="jsonl of LLM poems (id, poem): their prompt-word syllables become fixed cells")
    d.add_argument("--keep-frac", type=float, default=0.0, help="also keep this share of the draft's other non-rhyme cells")
    d.add_argument("--free-window", type=int, default=None, help="anchor the whole draft except +-W cells around each rhyme cell")
    d.add_argument("--schedule", default="stanza", choices=["stanza", "line"], help="line = decide one วรรค at a time with earlier วรรค fixed (recurrent)")
    d.add_argument("--topic-boost", type=float, default=0.0, help="logit bonus for each บาท's own prompt-clause syllables (spreads the topic over all 8 lines)")
    a = ap.parse_args()
    train(a) if a.cmd == "train" else decode(a)