#!/usr/bin/env python3
"""Exports dragonSwing/xlm-roberta-capu to the ONNX files mobile apps run with ONNX Runtime (Android and iOS).

The model is rebuilt here from its weights (XLM-RoBERTa-base with one added token, $START, and two linear heads),
so the model repo's Python code is never imported. Writes out/repo/, laid out as the Hugging Face repo that gets
published:
  onnx/model.onnx               fp32: input_ids, attention_mask [batch, tokens], input_offsets [batch, words]
  onnx/model_quantized.onnx     8-bit weights (MatMulNBits, int8 embedding): the file mobile apps should download
  unigram.tsv                   tokenizer: one "piece<TAB>score" line per token id, plus $START as the last id
  labels.txt, d_tags.txt        output classes of logits and detect_logits, one per line in index order
  number_words.tsv              "word<TAB>tag<TAB>replacement" for the $TRANSFORM_VERB_* labels (năm -> 5, gờ -> G)
  tokenizer.json                fast tokenizer with $START, for Python users (tokenizers / transformers)
  config.json                   ids and pipeline settings

Usage:
  .venv/bin/python -I export.py
"""

import hashlib
import json
import shutil
import sys
from pathlib import Path

import numpy as np
import onnx
import onnxruntime as ort
import torch
from onnxruntime.quantization import QuantType, quantize_dynamic
from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer
from safetensors.torch import load_file
from torch import nn
from transformers import AutoConfig, AutoTokenizer, XLMRobertaModel

sys.path.insert(0, str(Path(__file__).resolve().parent))  # -I drops the script's own directory
from common import HERE, SRC  # noqa: E402

BASE = "FacebookAI/xlm-roberta-base"
# The weights as safetensors (no pickle, so loading cannot run code): Hugging Face's conversion of pytorch_model.bin,
# open as dragonSwing/xlm-roberta-capu refs/pr/2. Tensor for tensor identical to pytorch_model.bin.
WEIGHTS = "model.safetensors"
WEIGHTS_REVISION = "13f33470adda1f517b487e730cb71e94d3fbca9e"
WEIGHTS_SHA256 = "3e04125e6af567a431c5c216225f61e27b48c65012aa48e1ecac481c575ebca9"
START_TOKEN = "$START"
OPSET = 17


class Seq2Labels(nn.Module):
    """Seq2LabelsModel from the model repo, inference only: labels and error-detection logits for each word.

    input_offsets holds, for each word, the index of its first sub-word token; index 0 is the $START token.
    """

    def __init__(self, vocab_size: int, num_labels: int, num_detect: int):
        super().__init__()
        config = AutoConfig.from_pretrained(BASE)
        config.vocab_size = vocab_size
        self.bert = XLMRobertaModel(config, add_pooling_layer=False)
        self.classifier = nn.Linear(config.hidden_size, num_labels)
        self.detector = nn.Linear(config.hidden_size, num_detect)

    def forward(self, input_ids, attention_mask, input_offsets):
        hidden = self.bert(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
        index = input_offsets.unsqueeze(-1).expand(-1, -1, hidden.size(-1))
        words = torch.gather(hidden, 1, index)
        return self.classifier(words), self.detector(words)


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def load_model() -> Seq2Labels:
    config = json.loads((SRC / "config.json").read_text())
    weights = SRC / WEIGHTS
    if not weights.exists() or sha256(weights) != WEIGHTS_SHA256:
        sys.exit(f"error: {weights} missing or not the pinned file; see README.md for the download command")
    state = load_file(weights)
    state = {k: v for k, v in state.items() if not k.startswith("bert.pooler.")}
    vocab_size = state["bert.embeddings.word_embeddings.weight"].shape[0]
    model = Seq2Labels(vocab_size, config["vocab_size"], config["num_detect_classes"])
    # position_ids is a buffer that some transformers versions save and others do not.
    missing, unexpected = model.load_state_dict(state, strict=False)
    missing = [k for k in missing if not k.endswith("position_ids")]
    unexpected = [k for k in unexpected if not k.endswith("position_ids")]
    if missing or unexpected:
        sys.exit(f"error: weights do not match the model: missing {missing}, unexpected {unexpected}")
    return model.eval()


def sample(tokenizer) -> dict[str, torch.Tensor]:
    words = [START_TOKEN] + "theo đó thủ tướng dự kiến tiếp bộ trưởng nông nghiệp mỹ".split()
    batch = tokenizer([words, words[:5]], is_split_into_words=True, add_special_tokens=False,
                      padding=True, return_tensors="pt")
    offsets = []
    for i in range(2):
        ids = batch.word_ids(i)
        offsets.append([0] + [j for j in range(1, len(ids)) if ids[j] is not None and ids[j] != ids[j - 1]])
    width = max(map(len, offsets))
    return {
        "input_ids": batch["input_ids"],
        "attention_mask": batch["attention_mask"],
        "input_offsets": torch.tensor([o + [0] * (width - len(o)) for o in offsets]),
    }


def export(model: Seq2Labels, inputs: dict[str, torch.Tensor], path: Path) -> None:
    torch.onnx.export(
        model, (inputs["input_ids"], inputs["attention_mask"], inputs["input_offsets"]), str(path),
        input_names=["input_ids", "attention_mask", "input_offsets"],
        output_names=["logits", "detect_logits"],
        dynamic_axes={
            "input_ids": {0: "batch", 1: "tokens"},
            "attention_mask": {0: "batch", 1: "tokens"},
            "input_offsets": {0: "batch", 1: "words"},
            "logits": {0: "batch", 1: "words"},
            "detect_logits": {0: "batch", 1: "words"},
        },
        opset_version=OPSET,
        dynamo=False,
    )
    onnx.checker.check_model(str(path))
    print(f"wrote {path}")


def check(path: Path, model: Seq2Labels, inputs: dict[str, torch.Tensor]) -> None:
    with torch.no_grad():
        expected = [t.numpy() for t in model(**inputs)]
    session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
    actual = session.run(None, {k: v.numpy() for k, v in inputs.items()})
    for name, e, a in zip(("logits", "detect_logits"), expected, actual):
        diff = float(np.abs(e - a).max())
        same = float((e.argmax(-1) == a.argmax(-1)).mean())
        print(f"  {path.name} {name}: max abs diff {diff:.2e}, argmax agreement {same:.0%}")


def quantize(src: Path, dst: Path) -> None:
    """8-bit weights, activations left in float: as accurate as fp32 on the test set, a quarter of the size.

    Linear layers become MatMulNBits (8-bit, symmetric, blocks of 128 along the input dimension;
    accuracy_level 4 lets the CPU kernels compute with int8 activations quantized per block). Plain dynamic int8
    (MatMulInteger, one scale per tensor for activations) lost 2 points of punctuation F1. The 250k-token
    embedding, two thirds of the weights, becomes an int8 Gather with one scale.
    """
    tmp = dst.with_name("matmul_nbits.onnx")
    quantizer = MatMulNBitsQuantizer(onnx.load(str(src)), bits=8, block_size=128, is_symmetric=True,
                                     accuracy_level=4, op_types_to_quantize=("MatMul",))
    quantizer.process()
    quantizer.model.save_model_to_file(str(tmp), False)
    quantize_dynamic(tmp, dst, op_types_to_quantize=["Gather"], weight_type=QuantType.QInt8,
                     per_channel=False, reduce_range=False)
    tmp.unlink()
    onnx.checker.check_model(str(dst))
    print(f"wrote {dst}")


def write_tokenizer(tokenizer, repo: Path) -> None:
    tokenizer.save_pretrained(repo)
    # Only tokenizer.json (which includes $START) is published; sentencepiece.bpe.model lacks $START.
    for name in ("special_tokens_map.json", "tokenizer_config.json", "added_tokens.json", "sentencepiece.bpe.model"):
        (repo / name).unlink(missing_ok=True)
    tokenizer_json = json.loads((repo / "tokenizer.json").read_text(encoding="utf-8"))
    pieces = tokenizer_json["model"]["vocab"]
    if any(c in piece for piece, _ in pieces for c in "\t\n\r"):
        sys.exit("error: a token contains a tab or newline and cannot go in unigram.tsv")
    start_id = tokenizer.convert_tokens_to_ids(START_TOKEN)
    if start_id != len(pieces):
        sys.exit(f"error: {START_TOKEN} has id {start_id}, expected {len(pieces)}")
    # $START gets the lowest score so the Viterbi search never picks it for ordinary text.
    lines = [f"{piece}\t{score!r}" for piece, score in pieces] + [f"{START_TOKEN}\t-1e9"]
    (repo / "unigram.tsv").write_text("\n".join(lines) + "\n", encoding="utf-8")
    print(f"wrote {repo / 'unigram.tsv'} ({len(lines)} tokens)")


def write_vocab(repo: Path) -> None:
    for name in ("labels.txt", "d_tags.txt"):
        shutil.copy2(SRC / "vocabulary" / name, repo / name)
    # utils.get_verb_form_dicts: "word1_word2:tag1_tag2" lines, the first replacement for (word1, tag1, tag2) wins.
    seen, lines = set(), []
    for line in (SRC / "verb-form-vocab.txt").read_text(encoding="utf-8").splitlines():
        words, tags = line.split(":")
        word1, word2 = words.split("_")
        tag1, tag2 = tags.split("_")
        key = (word1, f"{tag1}_{tag2.strip()}")
        if key not in seen:
            seen.add(key)
            lines.append(f"{key[0]}\t{key[1]}\t{word2}")
    (repo / "number_words.tsv").write_text("\n".join(lines) + "\n", encoding="utf-8")
    print(f"wrote {repo / 'number_words.tsv'} ({len(lines)} entries)")


def write_config(model: Seq2Labels, repo: Path) -> None:
    labels = (repo / "labels.txt").read_text(encoding="utf-8").splitlines()
    config = {
        "base_model": "dragonSwing/xlm-roberta-capu",
        "architecture": "xlm-roberta-base + Seq2Labels heads",
        "vocab_size": model.bert.config.vocab_size,
        "start_token": START_TOKEN,
        "start_token_id": model.bert.config.vocab_size - 1,
        "pad_token_id": 1,
        "unk_token_id": 3,
        "max_position_embeddings": model.bert.config.max_position_embeddings,
        "num_labels": len(labels),
        "keep_label_id": labels.index("$KEEP"),
        "pipeline": {
            "max_words": 64,
            "min_words": 3,
            "iterations": 3,
            "chunk_size": 48,
            "overlap_size": 12,
            "min_words_cut": 6,
            "punctuation": [".", ",", ":", "?"],
        },
    }
    (repo / "config.json").write_text(json.dumps(config, indent=2) + "\n", encoding="utf-8")


def main() -> None:
    repo = HERE / "out" / "repo"
    onnx_dir = repo / "onnx"
    onnx_dir.mkdir(parents=True, exist_ok=True)

    model = load_model()
    tokenizer = AutoTokenizer.from_pretrained(BASE)
    tokenizer.add_tokens([START_TOKEN])
    if len(tokenizer) != model.bert.config.vocab_size:
        sys.exit(f"error: tokenizer has {len(tokenizer)} tokens, model {model.bert.config.vocab_size}")
    inputs = sample(tokenizer)

    fp32, int8 = onnx_dir / "model.onnx", onnx_dir / "model_quantized.onnx"
    export(model, inputs, fp32)
    quantize(fp32, int8)
    for path in (fp32, int8):
        check(path, model, inputs)

    write_tokenizer(tokenizer, repo)
    write_vocab(repo)
    write_config(model, repo)


if __name__ == "__main__":
    main()