#!/usr/bin/env python3
"""Vietnamese capitalization and punctuation restoration with the ONNX export of dragonSwing/xlm-roberta-capu.

The reference for the Kotlin and Swift ports: it needs only numpy, onnxruntime and the files of the published repo
(onnx/*.onnx, unigram.tsv, labels.txt, number_words.tsv), and does every step itself, the way an app must.
It is a port of the model repo's GecBERTModel (gec_model.py, utils.py) with the model card's settings
(split_chunk=True): same chunking, the same up-to-3 correction rounds, the same chunk merging, so it gives the same
text (validate.py checks this).

  capu = Capu("path/to/repo", quantized=True)
  capu("theo đó thủ tướng dự kiến tiếp bộ trưởng nông nghiệp mỹ")  # "Theo đó, Thủ tướng dự kiến tiếp Bộ trưởng Nông nghiệp Mỹ."

Usage:
  python -I capu.py path/to/repo "text to correct" [--fp32]
"""

import re
import sys
import unicodedata
from pathlib import Path

import numpy as np
import onnxruntime as ort

START = "$START"
UNK, PAD = "@@UNKNOWN@@", "@@PADDING@@"
PUNCTUATION = (":", ".", ",", "?")
MAX_WORDS, MIN_WORDS, ITERATIONS = 64, 3, 3
CHUNK_SIZE, OVERLAP_SIZE, MIN_WORDS_CUT = 48, 12, 6
MAX_TOKENS = 512  # XLM-RoBERTa's position embeddings


class Tokenizer:
    """XLM-RoBERTa's SentencePiece unigram model (as in tokenizers' Unigram), one word at a time.

    Each word is NFKC-normalized, prefixed with "▁" and split by Viterbi search for the pieces with the highest
    total log-probability. A character that no piece covers becomes <unk>, with consecutive <unk>s fused.
    """

    def __init__(self, path: Path):
        self.ids: dict[str, int] = {}
        self.scores: list[float] = []
        with open(path, encoding="utf-8") as f:
            for i, line in enumerate(f):
                piece, score = line.rstrip("\n").split("\t")
                self.ids.setdefault(piece, i)
                self.scores.append(float(score))
        self.start_id = self.ids[START]
        self.unk_id = self.ids["<unk>"]
        self.unk_score = min(self.scores[4:self.start_id]) - 10.0  # tokenizers: min score - kUnkPenalty
        self.max_len = max(len(p) for p in self.ids)

    def encode_word(self, word: str) -> list[int]:
        text = "▁" + unicodedata.normalize("NFKC", word)
        n = len(text)
        best = [-np.inf] * (n + 1)  # best score of text[:end]
        back: list[tuple[int, int]] = [(0, 0)] * (n + 1)  # (start, token id) of the last piece of that path
        best[0] = 0.0
        for start in range(n):
            if best[start] == -np.inf:
                continue
            single = False
            for end in range(start + 1, min(n, start + self.max_len) + 1):
                token = self.ids.get(text[start:end])
                if token is None or token == self.start_id:
                    continue
                single |= end == start + 1
                score = best[start] + self.scores[token]
                if score > best[end]:
                    best[end], back[end] = score, (start, token)
            if not single and best[start] + self.unk_score > best[start + 1]:
                best[start + 1], back[start + 1] = best[start] + self.unk_score, (start, self.unk_id)
        tokens, end = [], n
        while end > 0:
            start, token = back[end]
            if not (token == self.unk_id and tokens and tokens[-1] == self.unk_id):
                tokens.append(token)
            end = start
        return tokens[::-1]


def capitalize(s: str) -> str:
    return s[:1].upper() + s[1:].lower()


def transform_case(token: str, label: str) -> str:
    # utils.convert_using_case; the endswith order matters ("$..._CAPITAL_1" does not end with "CAPITAL").
    if label.endswith("LOWER"):
        return token.lower()
    if label.endswith("UPPER"):
        return token.upper()
    if label.endswith("CAPITAL"):
        return capitalize(token)
    if label.endswith("CAPITAL_1"):
        return token[0] + capitalize(token[1:])
    if label.endswith("UPPER_-1"):
        return token[:-1].upper() + token[-1]
    return token


class Capu:
    def __init__(self, repo: str | Path, quantized: bool = True, threads: int = 4):
        repo = Path(repo)
        opts = ort.SessionOptions()
        opts.intra_op_num_threads = threads
        name = "model_quantized.onnx" if quantized else "model.onnx"
        self.session = ort.InferenceSession(str(repo / "onnx" / name), opts, providers=["CPUExecutionProvider"])
        self.tokenizer = Tokenizer(repo / "unigram.tsv")
        self.labels = (repo / "labels.txt").read_text(encoding="utf-8").splitlines()
        self.number_words: dict[tuple[str, str], str] = {}
        for line in (repo / "number_words.tsv").read_text(encoding="utf-8").splitlines():
            word, tag, replacement = line.split("\t")
            self.number_words[(word, tag)] = replacement

    def __call__(self, text: str) -> str:
        chunks = split_chunks(text.split())
        merged: list[str] = []
        for chunk in chunks:
            chunk = self.correct(chunk)
            merged = chunk if len(chunks) == 1 else merge_chunks(merged, chunk)
        return re.sub(r"\s+([:.,?])", r"\1", " ".join(merged))

    def correct(self, tokens: list[str]) -> list[str]:
        """GecBERTModel.handle_batch for one chunk: up to 3 rounds, stopping when the text stops changing."""
        if len(tokens) < MIN_WORDS:
            return tokens
        seen = [tokens]
        for _ in range(ITERATIONS):
            new = self.apply(tokens, self.predict(tokens))
            if new == tokens:
                break
            tokens = new
            if new in seen:
                break
            seen.append(new)
        return tokens

    def predict(self, tokens: list[str]) -> list[int]:
        """Label id for $START and each of the first MAX_WORDS tokens."""
        ids, offsets = [self.tokenizer.start_id], [0]
        for word in tokens[:MAX_WORDS]:
            offsets.append(len(ids))
            ids += self.tokenizer.encode_word(word)
        if len(ids) > MAX_TOKENS:
            # Only happens with very long "words" (URLs); words starting past the limit keep $KEEP.
            words = sum(o < MAX_TOKENS for o in offsets)
            ids, offsets = ids[:MAX_TOKENS], offsets[:words]
        logits = self.session.run(["logits"], {
            "input_ids": np.array([ids], np.int64),
            "attention_mask": np.ones((1, len(ids)), np.int64),
            "input_offsets": np.array([offsets], np.int64),
        })[0][0]
        labels = logits.argmax(-1).tolist()
        return labels + [0] * (min(len(tokens), MAX_WORDS) + 1 - len(labels))

    def apply(self, tokens: list[str], labels: list[int]) -> list[str]:
        """GecBERTModel.postprocess_batch and utils.get_target_sent_by_edits."""
        if max(labels) == 0:
            return tokens
        edits = []  # (position, label) in the order of the words
        for i in range(min(len(tokens), MAX_WORDS) + 1):
            label = self.labels[labels[i]]
            if label in (UNK, PAD, "$KEEP"):
                continue
            if label.startswith("$APPEND_"):
                edits.append((i, i, label[len("$APPEND_"):]))
            elif label.startswith("$MERGE_"):
                edits.append((i, i, label))
            else:  # $TRANSFORM_*; i - 1 because position 0 is $START
                edits.append((i - 1, i, label))

        target, shift = tokens[:], 0
        for start, end, label in edits:
            pos = start + shift
            if start < 0:
                continue
            source = target[pos] if pos < len(target) else ""
            if start == end:  # append a punctuation mark or a $MERGE_ marker after the word
                if (pos < len(target) and target[pos] == label) or (pos > 0 and target[pos - 1] == label):
                    continue
                target.insert(pos, label)
                shift += 1
            elif label.startswith("$TRANSFORM_CASE"):
                target[pos] = transform_case(source, label)
            elif label.startswith("$TRANSFORM_VERB_"):
                target[pos] = self.number_words.get((source, label[len("$TRANSFORM_VERB_"):]), source)
        return replace_merges(target)


def replace_merges(tokens: list[str]) -> list[str]:
    """utils.replace_merge_transforms: "5 $MERGE_SPACE G" -> "5G"."""
    if not any(t.startswith("$MERGE_") for t in tokens):
        return tokens
    if tokens[0].startswith("$MERGE_"):
        tokens = tokens[1:]
    if tokens and tokens[-1].startswith("$MERGE_"):
        tokens = tokens[:-1]
    line = " ".join(tokens).replace(" $MERGE_HYPHEN ", "-").replace(" $MERGE_SPACE ", "")
    return re.sub(r"([.,?:]\s+)+", r"\1", line).split()


def split_chunks(words: list[str]) -> list[list[str]]:
    """GecBERTModel.split_chunks: overlapping windows of CHUNK_SIZE words."""
    n = len(words)
    if n <= CHUNK_SIZE:
        return [words]
    if n < CHUNK_SIZE * 2 - OVERLAP_SIZE:
        split = (n + OVERLAP_SIZE + 1) // 2
        return [words[:split], words[split - OVERLAP_SIZE:]]
    return [words[i:i + CHUNK_SIZE] for i in range(0, n - OVERLAP_SIZE, CHUNK_SIZE - OVERLAP_SIZE)]


def merge_chunks(tokens: list[str], following: list[str]) -> list[str]:
    """GecBERTModel.apply_chunk_merging: joins two corrected chunks inside their overlap.

    Aligns the last OVERLAP_SIZE words of one chunk with the first OVERLAP_SIZE of the next (punctuation ignored,
    lower-cased) and cuts at the first common run at least MIN_WORDS_CUT words into the overlap. Where the original
    finds no cut it raises, and its caller drops the following chunk; this keeps that behavior.
    """
    if not tokens:
        return following
    source_idx, source = [], []
    i = 0
    while len(source_idx) < OVERLAP_SIZE and -i < len(tokens):
        i -= 1
        if tokens[i] not in PUNCTUATION:
            source_idx.insert(0, i)
            source.insert(0, tokens[i].lower())
    target_idx, target = [], []
    i = 0
    while len(target_idx) < OVERLAP_SIZE and i < len(following):
        if following[i] not in PUNCTUATION:
            target_idx.append(i)
            target.append(following[i].lower())
        i += 1

    keep = OVERLAP_SIZE - MIN_WORDS_CUT
    for tag, i1, i2, j1, j2 in opcodes(source, target):
        if tag == "equal":
            if i1 >= keep:
                return tokens[:source_idx[i1]] + following[target_idx[j1]:]
            if i2 > keep:
                j = j2 - i2 + keep
                if j >= len(target_idx):
                    return tokens  # IndexError in the original
                return tokens[:source_idx[keep]] + following[target_idx[j]:]
        elif tag == "delete" and i1 == 0:
            keep += i2 // 2
    return tokens


def matching_blocks(a: list[str], b: list[str]) -> list[tuple[int, int, int]]:
    """difflib.SequenceMatcher(None, a, b).get_matching_blocks() for short lists (no junk, no autojunk)."""

    def longest(alo, ahi, blo, bhi):
        # Longest common run; ties go to the smallest i, then the smallest j, as in difflib.
        besti, bestj, bestsize = alo, blo, 0
        lengths: dict[int, int] = {}
        for i in range(alo, ahi):
            current = {}
            for j in range(blo, bhi):
                if a[i] == b[j]:
                    k = current[j] = lengths.get(j - 1, 0) + 1
                    if k > bestsize:
                        besti, bestj, bestsize = i - k + 1, j - k + 1, k
            lengths = current
        return besti, bestj, bestsize

    queue, blocks = [(0, len(a), 0, len(b))], []
    while queue:
        alo, ahi, blo, bhi = queue.pop()
        i, j, k = longest(alo, ahi, blo, bhi)
        if k:
            blocks.append((i, j, k))
            if alo < i and blo < j:
                queue.append((alo, i, blo, j))
            if i + k < ahi and j + k < bhi:
                queue.append((i + k, ahi, j + k, bhi))
    blocks.sort()
    merged = []
    for i, j, k in blocks:  # join adjacent blocks
        if merged and merged[-1][0] + merged[-1][2] == i and merged[-1][1] + merged[-1][2] == j:
            merged[-1] = (merged[-1][0], merged[-1][1], merged[-1][2] + k)
        else:
            merged.append((i, j, k))
    merged.append((len(a), len(b), 0))
    return merged


def opcodes(a: list[str], b: list[str]) -> list[tuple[str, int, int, int, int]]:
    """difflib.SequenceMatcher(None, a, b).get_opcodes()."""
    i = j = 0
    result = []
    for ai, bj, size in matching_blocks(a, b):
        tag = "replace" if i < ai and j < bj else "delete" if i < ai else "insert" if j < bj else ""
        if tag:
            result.append((tag, i, ai, j, bj))
        i, j = ai + size, bj + size
        if size:
            result.append(("equal", ai, i, bj, j))
    return result


if __name__ == "__main__":
    if len(sys.argv) < 3:
        sys.exit(f"usage: {sys.argv[0]} path/to/repo 'text' [--fp32]")
    print(Capu(sys.argv[1], quantized="--fp32" not in sys.argv)(sys.argv[2]))