#!/usr/bin/env python3
"""Exports vinai/PhoWhisper-<size> to the ONNX files mobile apps run with ONNX Runtime (Android and iOS).

Writes out/<size>/repo/, laid out as the Hugging Face repo that gets published:
  onnx/encoder_model.onnx                     fp32 encoder: input_features [1, n_mels, 3000] -> last_hidden_state
  onnx/decoder_model_merged.onnx              fp32 decoder with KV cache (use_cache_branch selects the first step)
  onnx/encoder_model_quantized.onnx           int8 (dynamic, MatMul weights) versions of the two above:
  onnx/decoder_model_merged_quantized.onnx    the files mobile apps should download
  config.json, generation_config.json, preprocessor_config.json, tokenizer files

Weights over 2 GB (medium, large) go to a sidecar <name>.onnx_data next to the graph.

Usage:
  .venv/bin/python -I export.py tiny|base|small|medium|large
"""

import shutil
import sys
from pathlib import Path

import numpy as np
import onnx
from onnx import helper, numpy_helper
from onnxruntime.quantization import QuantType, quantize_dynamic

HERE = Path(__file__).resolve().parent
SIZES = ("tiny", "base", "small", "medium", "large")
KEEP = ("encoder_model", "decoder_model_merged")
COPIED = (
    "config.json",
    "generation_config.json",
    "preprocessor_config.json",
    "tokenizer.json",
    "tokenizer_config.json",
    "vocab.json",
    "merges.txt",
    "added_tokens.json",
    "special_tokens_map.json",
    "normalizer.json",
)
# Protobuf messages cannot exceed 2 GB; above this the weights go to a sidecar file.
EXTERNAL_DATA_THRESHOLD = int(1.8 * 1024**3)


def check_and_save_large(original):
    """Wraps Optimum's check_and_save_model so it also saves models over 2 GB (PhoWhisper-large).

    It starts with model.ByteSize(), which current protobuf refuses to compute over 2 GB, so a
    large model goes straight to an external-data save here instead.
    """

    def check_and_save(model: onnx.ModelProto, save_path) -> None:
        if sum(len(t.raw_data) for t in model.graph.initializer) < EXTERNAL_DATA_THRESHOLD:
            return original(model, save_path)
        path = Path(save_path)
        save(model, path)
        onnx.checker.check_model(str(path))

    return check_and_save


def export(size: str, raw: Path) -> None:
    """Runs the Optimum export (encoder, decoder, decoder with past, merged decoder) into raw/."""
    if (raw / "decoder_model_merged.onnx").exists():
        print(f"reusing {raw}")
        return
    shutil.rmtree(raw, ignore_errors=True)
    from optimum.exporters.onnx import convert, main_export
    from optimum.onnx import graph_transformations

    patched = check_and_save_large(graph_transformations.check_and_save_model)
    graph_transformations.check_and_save_model = patched
    convert.check_and_save_model = patched
    main_export(
        f"vinai/PhoWhisper-{size}",
        output=raw,
        task="automatic-speech-recognition-with-past",
        opset=17,
    )


def save(model: onnx.ModelProto, path: Path) -> None:
    for stale in (path, path.with_name(path.name + "_data")):
        stale.unlink(missing_ok=True)
    weights = sum(len(t.raw_data) for t in model.graph.initializer)
    if weights > EXTERNAL_DATA_THRESHOLD:
        onnx.save_model(
            model, str(path),
            save_as_external_data=True, all_tensors_to_one_file=True,
            location=path.name + "_data", size_threshold=1024,
        )
    else:
        onnx.save_model(model, str(path))


def copy_fp32(raw: Path, onnx_dir: Path) -> None:
    for name in KEEP:
        # Re-saving gathers the weights into one sidecar (Optimum writes one file per tensor for large models).
        model = onnx.load(str(raw / f"{name}.onnx"), load_external_data=True)
        save(model, onnx_dir / f"{name}.onnx")
        print(f"wrote {onnx_dir / name}.onnx")


def rename_input(graph: onnx.GraphProto, old: str, new: str) -> None:
    for node in graph.node:
        node.input[:] = [new if i == old else i for i in node.input]
        for attr in node.attribute:
            if attr.type == onnx.AttributeProto.GRAPH:
                rename_input(attr.g, old, new)


def pretranspose_constants(graph: onnx.GraphProto, initializers: dict[str, onnx.TensorProto]) -> int:
    """Replaces Transpose(initializer) with a transposed initializer, in graph and its subgraphs.

    The decoder's LM head shares the token embedding and multiplies by its transpose, computed at run
    time. The quantizer skips a MatMul whose weight comes out of a node, which would leave the LM head
    fp32 (80 MB for tiny), so the transpose is folded into a constant first.
    """
    folded = 0
    for node in list(graph.node):
        for attr in node.attribute:
            if attr.type == onnx.AttributeProto.GRAPH:
                folded += pretranspose_constants(attr.g, initializers)
        if node.op_type != "Transpose" or node.input[0] not in initializers:
            continue
        perm = next((list(a.ints) for a in node.attribute if a.name == "perm"), None)
        folded_name = node.input[0] + "_transposed"
        if folded_name not in initializers:
            weight = numpy_helper.to_array(initializers[node.input[0]])
            initializers[folded_name] = numpy_helper.from_array(np.ascontiguousarray(weight.transpose(perm)), folded_name)
        graph.node.remove(node)
        rename_input(graph, node.output[0], folded_name)
        folded += 1
    return folded


def quantize(onnx_dir: Path, tmp: Path) -> None:
    tmp.mkdir(exist_ok=True)
    for name in KEEP:
        model = onnx.load(str(onnx_dir / f"{name}.onnx"), load_external_data=True)
        initializers = {t.name: t for t in model.graph.initializer}
        known = set(initializers)
        if pretranspose_constants(model.graph, initializers):
            model.graph.initializer.extend(t for n, t in initializers.items() if n not in known)
        src = tmp / f"{name}.onnx"
        save(model, src)
        dst = onnx_dir / f"{name}_quantized.onnx"
        # MatMul and Gather (the token embedding), also inside the merged decoder's If branches.
        # The encoder's two Conv layers stay fp32, since ConvInteger has no kernel on
        # some mobile builds. Signed weights, no reduce_range: ARM's int8 dot product kernels need neither.
        quantize_dynamic(
            src, dst,
            op_types_to_quantize=["MatMul", "Gather"],
            weight_type=QuantType.QInt8,
            per_channel=False,
            reduce_range=False,
            use_external_data_format=src.with_name(src.name + "_data").exists(),
            extra_options={"MatMulConstBOnly": True, "EnableSubgraph": True},
        )
        model = onnx.load(str(dst), load_external_data=True)
        save(model, dst)  # int8 weights of every size fit under 2 GB, so this inlines them again
        dst.with_name(dst.name + ".data").unlink(missing_ok=True)  # quantize_dynamic's sidecar, now inlined
        print(f"wrote {dst}")
    shutil.rmtree(tmp)


def copy_tokenizer(raw: Path, repo: Path) -> None:
    for name in COPIED:
        if (raw / name).exists():
            shutil.copy2(raw / name, repo / name)


def main() -> None:
    if len(sys.argv) != 2 or sys.argv[1] not in SIZES:
        sys.exit(f"usage: {sys.argv[0]} {'|'.join(SIZES)}")
    size = sys.argv[1]
    work = HERE / "out" / size
    raw, repo = work / "raw", work / "repo"
    onnx_dir = repo / "onnx"
    onnx_dir.mkdir(parents=True, exist_ok=True)

    export(size, raw)
    copy_fp32(raw, onnx_dir)
    quantize(onnx_dir, work / "tmp")
    copy_tokenizer(raw, repo)


if __name__ == "__main__":
    main()