"""Finnish punctuation and casing for speech recognition output (Mateli-punct, ONNX Runtime).

    python punctuate.py "mä luulen et se on ihan hyvä juttu mut en mä tiiä mitä mieltä sä oot"

    from punctuate import Punctuator
    p = Punctuator(".")                                  # directory with config.json, tokenizer.json, *.onnx
    p.punctuate("mä luulen et se on ihan hyvä juttu")    # -> "Mä luulen, et se on ihan hyvä juttu."
    p.punctuate(words, pauses=[0.05, 0.4, ...])          # pauses[i] = silence after words[i] in seconds

    stream = p.stream()                                  # live use, one word at a time
    for word, pause_before in asr_words:
        print(stream.push(word, pause_before), end=" ")  # words come out once 3 words follow them
    print(stream.flush())

Words are never changed, added or removed: each word gets one label, the punctuation after it
(none , . ? !) and its casing (lower, Capitalised, ALL CAPS). Pauses are optional; they help
mostly with sentence boundaries. Needs numpy, onnxruntime and tokenizers.
"""
from __future__ import annotations

import argparse
import bisect
import json
import math
from pathlib import Path

import numpy as np
import onnxruntime as ort
from tokenizers import Tokenizer

SENTENCE_END = {".", "?", "!"}


class Model:
    """input ids -> logits [seq, labels]."""

    def __init__(self, model_dir: Path, precision: str = "int8", threads: int = 1):
        cfg = json.loads((model_dir / "config.json").read_text(encoding="utf-8"))
        so = ort.SessionOptions()
        so.intra_op_num_threads = threads
        so.log_severity_level = 3
        self.sess = ort.InferenceSession(str(model_dir / cfg["files"][precision]), so,
                                         providers=["CPUExecutionProvider"])
        self.cfg = cfg

    def __call__(self, ids: list[int]) -> np.ndarray:
        x = np.array([ids], np.int64)
        return self.sess.run(["logits"], {"input_ids": x, "attention_mask": np.ones_like(x)})[0][0]


class Punctuator:
    def __init__(self, model_dir=".", precision: str = "int8", threads: int = 1,
                 window: int = 150, overlap: int = 50):
        d = Path(model_dir)
        self.model = Model(d, precision, threads)
        cfg = self.model.cfg
        self.labels = [cfg["id2label"][str(i)] for i in range(len(cfg["id2label"]))]
        self.tok = Tokenizer.from_file(str(d / "tokenizer.json"))
        pm = cfg["pause_markers"]
        self.edges = pm["edges"]
        self.flag_id = self.tok.token_to_id(pm["flag"])
        self.bucket_ids = [self.tok.token_to_id(b) for b in pm["buckets"]]
        self.cls_id, self.sep_id = self.tok.token_to_id("[CLS]"), self.tok.token_to_id("[SEP]")
        self.unk_id = self.tok.token_to_id("[UNK]")
        if window < 1 or not 0 <= overlap < window:
            raise ValueError("need window >= 1 and 0 <= overlap < window")
        self.window, self.overlap = window, overlap

    # ------------------------------------------------------------------ one model call
    def _encode(self, words, pauses):
        """[CLS] (flag) w1 pieces (pause marker) w2 pieces ... [SEP]; returns ids and the word
        index of each position (-1 for special tokens). The flag says pauses are given; a marker
        follows a word whose pause is at least the first bucket edge (unknown pause: no marker)."""
        ids, widx = [self.cls_id], [-1]
        if pauses is not None:
            ids.append(self.flag_id)
            widx.append(-1)
        for w, word in enumerate(words):
            pieces = self.tok.encode(word, add_special_tokens=False).ids or [self.unk_id]
            ids += pieces
            widx += [w] * len(pieces)
            p = None if pauses is None else pauses[w]
            if p is not None and not math.isnan(p):
                b = bisect.bisect_right(self.edges, p)
                if b > 0:
                    ids.append(self.bucket_ids[b - 1])
                    widx.append(-1)
        return ids[:511] + [self.sep_id], widx[:511] + [-1]

    def _call(self, words, pauses) -> list[int]:
        """Label id of each word (the prediction on its first sub-token); -1 if truncated away."""
        ids, widx = self._encode(words, pauses)
        logits = self.model(ids)
        out = [-1] * len(words)
        for i, w in enumerate(widx):
            if w >= 0 and out[w] < 0:
                out[w] = int(np.argmax(logits[i]))
        return out

    # ------------------------------------------------------------------ whole texts
    def predict(self, words: list[str], pauses: list[float] | None = None) -> list[str]:
        """One label per word ("O", ",|U", ".", ...). Long inputs run in overlapping windows of
        `window` words; each word keeps the prediction from the window where it has the most
        context on both sides."""
        n = len(words)
        if pauses is not None and len(pauses) != n:
            raise ValueError("give one pause per word (the silence after it)")
        best = [(-1, -1)] * n                     # (context on the nearer side, label id)

        def run(s):
            chunk = words[s:s + self.window]
            labels = self._call(chunk, None if pauses is None else pauses[s:s + self.window])
            for i, lab in enumerate(labels):
                if lab < 0:                       # cut off by the 512-token limit
                    continue
                left = i if s > 0 else n
                right = len(chunk) - 1 - i if s + len(chunk) < n else n
                if min(left, right) > best[s + i][0]:
                    best[s + i] = (min(left, right), lab)

        stride = self.window - self.overlap
        for s in (range(0, max(1, n - self.overlap), stride) if n > self.window else [0]):
            run(s)
        while (i := next((i for i, (_, lab) in enumerate(best) if lab < 0), None)) is not None:
            run(i)                                # words that very long windows lost to truncation
            if best[i][1] < 0:
                raise RuntimeError(f"word {i} does not fit into the model input")
        return [self.labels[lab] for _, lab in best]

    def punctuate(self, text_or_words, pauses: list[float] | None = None) -> str:
        words = text_or_words.split() if isinstance(text_or_words, str) else list(text_or_words)
        return render(words, self.predict(words, pauses))

    def stream(self, **kw) -> "StreamingPunctuator":
        return StreamingPunctuator(self, **kw)


def render(words: list[str], labels: list[str], capitalise_first: bool = True) -> str:
    """Apply labels; also capitalise the first word and every word after . ? !"""
    if len(words) != len(labels):
        raise ValueError(f"{len(words)} words but {len(labels)} labels")
    out, prev = [], "." if capitalise_first else "O"
    for w, lab in zip(words, labels):
        punct, _, case = lab.partition("|")
        if case == "A":
            w = w.upper()
        elif case == "U" or prev in SENTENCE_END:
            w = w[:1].upper() + w[1:]
        out.append(w if punct == "O" else w + punct)
        prev = punct
    return " ".join(out)


class StreamingPunctuator:
    """Live punctuation, one word at a time. A word is committed once `lookahead` later words are
    known; the model sees up to `context` committed words before the pending ones and runs at
    most every `every` words. With the defaults a word appears about 2.4 s after it is spoken
    (at a typical speaking rate) and boundary F1 is within 2 points of whole-text punctuation."""

    def __init__(self, p: Punctuator, lookahead: int = 3, context: int = 64, every: int = 5):
        if lookahead < 0 or context < 0 or every < 1:
            raise ValueError("need lookahead >= 0, context >= 0 and every >= 1")
        self.p, self.k, self.L, self.every = p, lookahead, context, every
        self.context, self.pending, self.pauses = [], [], []    # pauses: after each context+pending word
        self.new, self.prev = 0, "."

    def push(self, word: str, pause_before: float | None = None) -> str:
        """Add a word; pause_before is the silence since the previous word (None = no pauses).
        Returns the newly committed text ("" if nothing is committed yet)."""
        if self.pauses:
            self.pauses[-1] = math.nan if pause_before is None else pause_before
        self.pauses.append(math.nan)                 # the newest word's pause is not known yet
        self.pending.append(word)
        self.new += 1
        if len(self.pending) <= self.k or self.new < self.every:
            return ""
        return self._run(final=False)

    def flush(self, pause_after: float | None = None) -> str:
        """Commit everything still pending; pause_after is the silence after the last word, if known."""
        if self.pending and pause_after is not None:
            self.pauses[-1] = pause_after
        parts = []
        while self.pending:
            text = self._run(final=True)
            if not text:                          # the context left no room: drop it and retry
                if not self.context:
                    raise RuntimeError("cannot fit the next word into the model input")
                self.pauses = self.pauses[len(self.context):]
                self.context = []
                continue
            parts.append(text)
        return " ".join(parts)

    def _run(self, final: bool) -> str:
        words = self.context + self.pending
        use_pauses = any(not math.isnan(x) for x in self.pauses)
        labels = self.p._call(words, self.pauses if use_pauses else None)[len(self.context):]
        covered = next((i for i, lab in enumerate(labels) if lab < 0), len(labels))   # < len if truncated
        n = covered if final else max(0, covered - self.k)
        committed = self.pending[:n]
        text = render(committed, [self.p.labels[lab] for lab in labels[:n]], capitalise_first=False)
        if n:      # the casing rule across calls: capitalise after a sentence end
            if self.prev in SENTENCE_END:
                text = text[:1].upper() + text[1:]
            self.prev = self.p.labels[labels[n - 1]].partition("|")[0]
        self.context = (self.context + committed)[-self.L:] if self.L else []
        self.pending = self.pending[n:]
        self.pauses = self.pauses[len(self.pauses) - len(self.context) - len(self.pending):]
        self.new = 0
        return text


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("text", help="lowercase words without punctuation")
    ap.add_argument("--model-dir", default=str(Path(__file__).resolve().parent))
    ap.add_argument("--fp32", action="store_true", help="use model.onnx instead of model.int8.onnx")
    a = ap.parse_args()
    print(Punctuator(a.model_dir, "fp32" if a.fp32 else "int8").punctuate(a.text))


if __name__ == "__main__":
    main()