#!/usr/bin/env python
"""Restauration des majuscules et de la ponctuation pour la sortie ASR malgache.

    from punctuate import restore
    restore("lanonana ara-panambadiana iray ahitana karazan'olona mandihy")
    -> "Lanonana ara-panambadiana iray ahitana karazan'olona mandihy."

Garantie : normalize(sortie).split() == normalize(entree).split() (assert).
Les mots ne sont jamais changes ; seules la casse et la ponctuation qui suit
chaque mot sont predites (plus, en option, la remise des lettres accentuees
qu'efface normalize(), ex. "tan na" -> "tanàna", qui respecte la meme garantie).

Modele : XLM-R base (Davlan/xlm-roberta-base-finetuned-malagasy) + une tete de
classification de jetons combinee (4 casses x 7 ponctuations), etiquette lue
sur le premier sous-jeton de chaque mot. Textes longs : fenetres glissantes.

CLI :
  python punctuate.py "texte en minuscules"            # restaure
  python punctuate.py --bench 100 --threads 2          # latence CPU
"""
import argparse
import json
import os
import sys
import time

import numpy as np
import torch

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
from common import CASES, PUNCTS, normalize, render  # noqa: E402

DEFAULT_MODEL = os.environ.get("MG_PUNCT_DIR", HERE)
NP = len(PUNCTS)


def windows(n, size, step):
    """Fenetres [a, b) couvrant n mots ; chaque mot sera lu dans la fenetre ou il est le plus central."""
    if n <= size:
        return [(0, n)]
    out, a = [], 0
    while True:
        b = min(a + size, n)
        out.append((a, b))
        if b == n:
            return out
        a += step


class Restorer:
    def __init__(self, model_dir=DEFAULT_MODEL, device="cpu", model=None, tokenizer=None, use_lex=True, use_bias=True, quantize=False,
                 win_words=150, win_step=100, max_tokens=512):
        from transformers import AutoModelForTokenClassification, AutoTokenizer

        self.device = device
        self.tok = tokenizer or AutoTokenizer.from_pretrained(model_dir)
        if model is None:
            model = AutoModelForTokenClassification.from_pretrained(model_dir)
            model.to(device).eval()
            if quantize and device == "cpu":
                # int8 dynamique sur les couches lineaires (option, CPU seulement).
                model = torch.ao.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)
        self.model = model
        tables = {"mixed": {}, "lex": {}}
        tpath = os.path.join(model_dir, "tables.json") if model_dir else ""
        if tpath and os.path.exists(tpath):
            tables = json.load(open(tpath))
        self.mixed = tables.get("mixed", {})
        self.lex = tables.get("lex", {}) if use_lex else {}
        self.win_words, self.win_step, self.max_tokens = win_words, win_step, max_tokens
        # Biais additif par signe sur log P(ponct.), regle sur la validation (tune_bias.py).
        self.bias = np.zeros(NP, dtype=np.float32)
        bpath = os.path.join(model_dir, "punct_bias.json") if model_dir else ""
        if use_bias and bpath and os.path.exists(bpath):
            b = json.load(open(bpath))
            self.bias = np.array([b.get(p, 0.0) for p in PUNCTS], dtype=np.float32)

    def predict_words(self, batch_words, batch_size=32):
        """Liste de listes de mots -> liste de (casses, ponctuations). Argmax des marginales."""
        out = []
        for w, (pc, pp, seen) in zip(batch_words, self.predict_probs(batch_words, batch_size)):
            cases = [CASES[i] for i in pc.argmax(1)] if len(w) else []
            puncts = [PUNCTS[i] for i in (np.log(pp + 1e-9) + self.bias).argmax(1)] if len(w) else []
            # Mot jamais vu (tronque) : pas de prediction -> minuscule, sans signe.
            for g in range(len(w)):
                if not seen[g]:
                    cases[g], puncts[g] = "L", "O"
            out.append((cases, puncts))
        return out

    @torch.no_grad()
    def predict_probs(self, batch_words, batch_size=32):
        """Liste de listes de mots -> [(P(casse) [n,4], P(ponct.) [n,7], vu [n])] par exemple."""
        jobs = []  # (index exemple, a, b)
        for k, w in enumerate(batch_words):
            for a, b in windows(len(w), self.win_words, self.win_step):
                jobs.append((k, a, b))
        probs = [np.zeros((len(w), len(CASES) * NP), dtype=np.float32) for w in batch_words]
        best = [np.full(len(w), -1.0) for w in batch_words]
        jobs.sort(key=lambda j: j[2] - j[1])
        for s in range(0, len(jobs), batch_size):
            chunk = jobs[s : s + batch_size]
            enc = self.tok([batch_words[k][a:b] for k, a, b in chunk], is_split_into_words=True,
                           truncation=True, max_length=self.max_tokens, padding=True, return_tensors="pt")
            logits = self.model(input_ids=enc["input_ids"].to(self.device),
                                attention_mask=enc["attention_mask"].to(self.device)).logits.float()
            p = torch.softmax(logits, -1).cpu().numpy()
            for j, (k, a, b) in enumerate(chunk):
                prev = None
                for t, wid in enumerate(enc.word_ids(j)):
                    if wid is None or wid == prev:
                        prev = wid
                        continue
                    prev = wid
                    g = a + wid
                    # centralite : distance au bord de fenetre (les bords de texte ne comptent pas)
                    left = wid if a > 0 else 10**6
                    right = (b - a - 1 - wid) if b < len(batch_words[k]) else 10**6
                    c = min(left, right)
                    if c > best[k][g]:
                        best[k][g] = c
                        probs[k][g] = p[j, t]
        out = []
        for k, w in enumerate(batch_words):
            pr = probs[k].reshape(len(w), len(CASES), NP)
            out.append((pr.sum(2), pr.sum(1), best[k] >= 0))
        return out

    def restore_many(self, texts, batch_size=32):
        words = [normalize(t).split() for t in texts]
        preds = self.predict_words(words, batch_size=batch_size)
        outs = []
        for w, (c, p) in zip(words, preds):
            o = render(w, c, p, self.mixed, self.lex) if w else ""
            assert normalize(o).split() == w, (w, o)
            outs.append(o)
        return outs

    def restore(self, text):
        return self.restore_many([text])[0]


_DEFAULT = None


def restore(text: str) -> str:
    """Restaure majuscules et ponctuation (CPU). Le modele est charge au premier appel."""
    global _DEFAULT
    if _DEFAULT is None:
        _DEFAULT = Restorer()
    return _DEFAULT.restore(text)


def bench(n, threads, model_dir, quantize=False):
    """Latence CPU, une phrase a la fois, sur n hypotheses ASR de la VALIDATION."""
    torch.set_num_threads(threads)
    r = Restorer(model_dir, quantize=quantize)
    path = os.path.join(HERE, "asr_valid", "valid_lm_transcripts.jsonl")
    if os.path.exists(path):
        texts = [json.loads(l)["hyp"] for l in open(path)]
        import random
        texts = random.Random(0).sample(texts, min(n, len(texts)))  # le fichier est trie par duree
    else:
        texts = [" ".join(json.loads(l)["words"]) for l in open(os.path.join(HERE, "data", "valid.jsonl"))][:n]
    for t in texts[:5]:
        r.restore(t)
    lat = []
    for t in texts:
        t0 = time.perf_counter()
        r.restore(t)
        lat.append(1000 * (time.perf_counter() - t0))
    lat = np.array(lat)
    words = np.mean([len(t.split()) for t in texts])
    res = {"n": len(texts), "threads": threads, "int8": quantize, "mots_moyens": round(float(words), 1),
           "median_ms": round(float(np.median(lat)), 1), "p90_ms": round(float(np.percentile(lat, 90)), 1),
           "max_ms": round(float(lat.max()), 1), "moyenne_ms": round(float(lat.mean()), 1)}
    print(json.dumps(res, ensure_ascii=False))
    return res


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("text", nargs="?")
    ap.add_argument("--model", default=DEFAULT_MODEL)
    ap.add_argument("--bench", type=int, default=0)
    ap.add_argument("--threads", type=int, default=2)
    ap.add_argument("--out", default=None)
    ap.add_argument("--int8", action="store_true", help="quantification dynamique int8 (CPU)")
    a = ap.parse_args()
    if a.bench:
        res = bench(a.bench, a.threads, a.model, a.int8)
        if a.out:
            json.dump(res, open(a.out, "w"), ensure_ascii=False, indent=2)
    else:
        torch.set_num_threads(a.threads)
        print(Restorer(a.model, quantize=a.int8).restore(a.text or sys.stdin.read()))