#!/usr/bin/env python3
"""LM NEURONAL de rescoring (§6.1) — petit Transformer causal au NIVEAU CARACTÈRE,
entraîné UNIQUEMENT sur le texte train WAXAL de la langue (conforme : aucune donnée externe).
Le char-level est robuste sur 14k phrases et colle à la métrique (moitié CER).
Complète le KenLM 5-grammes : dépendances bien plus longues.

Usage : LANG=lin /root/venv/bin/python train_charlm.py
Sortie : /scratch/lm/charlm_<lang>.pt (poids + vocab)
"""
import json, math, os, time
import torch, torch.nn as nn

LANG = os.environ.get("LANG_ASR", "lin")
MANIFEST = "/root/devhard/train_%s_min.jsonl" % LANG
OUT = "/scratch/lm/charlm_%s.pt" % LANG
DIM = int(os.environ.get("DIM", "256"))
NLAYER = int(os.environ.get("NLAYER", "4"))
NHEAD = 4
EPOCHS = int(os.environ.get("EPOCHS", "12"))
BS = 64
MAXLEN = 256


def main():
    rows = [json.loads(l) for l in open(MANIFEST, encoding="utf-8")]
    texts = [r["text"].strip() for r in rows if r.get("text", "").strip()]
    print("%s : %d phrases" % (LANG, len(texts)), flush=True)

    chars = sorted(set("".join(texts)))
    stoi = {c: i + 2 for i, c in enumerate(chars)}   # 0=pad, 1=bos/eos
    V = len(stoi) + 2
    print("vocab caracteres : %d" % V, flush=True)

    def enc(t):
        return [1] + [stoi[c] for c in t if c in stoi][:MAXLEN - 2] + [1]

    seqs = [enc(t) for t in texts]
    # split train/val interne pour early-stop (pas de fuite devhard : manifest = train pur)
    nval = max(200, len(seqs) // 20)
    val, tr = seqs[:nval], seqs[nval:]

    class LM(nn.Module):
        def __init__(self):
            super().__init__()
            self.emb = nn.Embedding(V, DIM, padding_idx=0)
            self.pos = nn.Embedding(MAXLEN, DIM)
            layer = nn.TransformerEncoderLayer(DIM, NHEAD, DIM * 4, dropout=0.1,
                                               batch_first=True, norm_first=True)
            self.tr = nn.TransformerEncoder(layer, NLAYER)
            self.ln = nn.LayerNorm(DIM)
            self.head = nn.Linear(DIM, V)

        def forward(self, x):
            T = x.shape[1]
            h = self.emb(x) + self.pos(torch.arange(T, device=x.device))[None]
            mask = nn.Transformer.generate_square_subsequent_mask(T, device=x.device)
            h = self.tr(h, mask=mask, is_causal=True)
            return self.head(self.ln(h))

    dev = "cuda"
    model = LM().to(dev)
    nparam = sum(p.numel() for p in model.parameters())
    print("parametres : %.2f M" % (nparam / 1e6), flush=True)
    opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
    lossf = nn.CrossEntropyLoss(ignore_index=0)

    def batches(data, bs, shuffle=True):
        idx = torch.randperm(len(data)).tolist() if shuffle else range(len(data))
        buf = []
        for i in idx:
            buf.append(data[i])
            if len(buf) == bs:
                yield buf; buf = []
        if buf:
            yield buf

    def pad(b):
        L = max(len(x) for x in b)
        t = torch.zeros(len(b), L, dtype=torch.long)
        for i, x in enumerate(b):
            t[i, :len(x)] = torch.tensor(x)
        return t.to(dev)

    best = 1e9
    for ep in range(EPOCHS):
        model.train(); tot = n = 0
        for b in batches(tr, BS):
            x = pad(b)
            logits = model(x[:, :-1])
            loss = lossf(logits.reshape(-1, V), x[:, 1:].reshape(-1))
            opt.zero_grad(); loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            tot += float(loss); n += 1
        model.eval(); vt = vn = 0
        with torch.inference_mode():
            for b in batches(val, BS, False):
                x = pad(b)
                vt += float(lossf(model(x[:, :-1]).reshape(-1, V), x[:, 1:].reshape(-1))); vn += 1
        vl = vt / max(vn, 1)
        print("ep %d train %.4f val %.4f (ppl %.2f)" % (ep, tot / max(n, 1), vl, math.exp(vl)), flush=True)
        if vl < best:
            best = vl
            torch.save({"stoi": stoi, "dim": DIM, "nlayer": NLAYER, "V": V,
                        "state": model.state_dict()}, OUT)
    print("CHARLM_DONE %s  best_val %.4f  ppl %.2f -> %s" % (LANG, best, math.exp(best), OUT), flush=True)


if __name__ == "__main__":
    main()