"""Generate the reduced-precision variants of the SHARP ONNX model.

Each variant is written to its own folder next to `fp32/`, as `model.onnx` plus
`model.onnx.data` (the external weight file).

    python quantize.py fp16
    python quantize.py int8
    python quantize.py int4
    python quantize.py fp8

Requires: onnx, onnxruntime, onnxconverter-common (fp16), ml_dtypes (fp8).
"""

import argparse
import os

import numpy as np
import onnx
from onnx import TensorProto, helper, numpy_helper

SRC_MODEL = os.path.join("fp32", "model.onnx")
DATA_FILE = "model.onnx.data"


def save_model(model: onnx.ModelProto, folder: str) -> str:
    os.makedirs(folder, exist_ok=True)
    path = os.path.join(folder, "model.onnx")
    # onnx.save appends to an existing external-data file instead of replacing it,
    # so any leftovers from an earlier run have to go first
    for stale in (path, os.path.join(folder, DATA_FILE)):
        if os.path.exists(stale):
            os.remove(stale)
    onnx.save(
        model,
        path,
        save_as_external_data=True,
        all_tensors_to_one_file=True,
        location=DATA_FILE,
        size_threshold=1024,
        convert_attribute=False,
    )
    return path


def _cast_resize_fp16_inputs(model: onnx.ModelProto) -> int:
    """Insert Cast nodes so Resize gets float32 roi/scales.

    Operator block lists do not apply to `Constant` nodes folded into the graph, so the
    fp16 scales have to be cast back explicitly.
    """
    initializers = {t.name: t for t in model.graph.initializer}
    producers = {out: node for node in model.graph.node for out in node.output}
    patched = 0
    nodes = []
    for node in model.graph.node:
        if node.op_type != "Resize":
            nodes.append(node)
            continue
        for index in range(1, len(node.input)):
            name = node.input[index]
            if not name:
                continue
            producer = producers.get(name)
            is_fp16 = initializers.get(name) is not None and initializers[name].data_type == TensorProto.FLOAT16
            if producer is not None and producer.op_type == "Constant":
                tensor = [a.t for a in producer.attribute if a.name == "value"][0]
                is_fp16 = tensor.data_type == TensorProto.FLOAT16
            if not is_fp16:
                continue
            cast_out = f"{name}_float32"
            nodes.append(helper.make_node("Cast", [name], [cast_out], name=f"{cast_out}_cast", to=TensorProto.FLOAT))
            node.input[index] = cast_out
            patched += 1
        nodes.append(node)
    if patched:
        model.graph.ClearField("node")
        model.graph.node.extend(nodes)
    return patched


def convert_fp16() -> str:
    # onnxruntime's converter inserts the casts needed to keep every node's operand
    # dtypes consistent, which onnxconverter-common does not do reliably.
    from onnxruntime.transformers.float16 import convert_float_to_float16

    model = onnx.load(SRC_MODEL)  # resolves the external .data file
    # use_symbolic_shape_infer=False: the in-memory proto is >2 GB and cannot be serialized
    model = convert_float_to_float16(model, keep_io_types=True, disable_shape_infer=True)
    print(f"fp16: cast {_cast_resize_fp16_inputs(model)} Resize inputs back to float32")
    return save_model(model, "fp16")


def convert_int8() -> str:
    from onnxruntime.quantization import QuantType, quantize_dynamic

    out = os.path.join("int8", "model.onnx")
    os.makedirs("int8", exist_ok=True)
    quantize_dynamic(
        SRC_MODEL,
        out,
        weight_type=QuantType.QInt8,
        per_channel=True,
        reduce_range=False,
        use_external_data_format=True,
        extra_options={"MatMulConstBOnly": False},
    )
    return out


def convert_int4() -> str:
    from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer

    model = onnx.load(SRC_MODEL)
    quantizer = MatMulNBitsQuantizer(model, block_size=32, is_symmetric=True, bits=4)
    quantizer.process()
    # newer onnxruntime wraps the ModelProto in onnx_ir.Model
    quantized = quantizer.model.model if hasattr(quantizer.model, "model") else quantizer.model
    return save_model(quantized, "int4")


def convert_fp8() -> str:
    """Store weights as float8_e4m3fn with a per-output-channel scale.

    Each weight becomes `<name>`_fp8` plus a float32 scale, and a DequantizeLinear node
    restores `<name>`. float8 needs opset 19 and ONNX Runtime >= 1.25 for the CPU kernel.
    """
    import ml_dtypes

    model = onnx.load(SRC_MODEL)
    for opset in model.opset_import:
        if opset.domain in ("", "ai.onnx"):
            opset.version = max(opset.version, 19)

    dequantize = []
    new_initializers = []
    converted = 0
    for init in model.graph.initializer:
        array = numpy_helper.to_array(init)
        if array.dtype != np.float32 or array.ndim < 2 or array.size < 1024:
            continue
        name = init.name
        amax = np.abs(array).reshape(array.shape[0], -1).max(axis=1)
        scale = np.maximum(amax, 1e-9).astype(np.float32) / 448.0  # 448 = E4M3 max
        shape = (array.shape[0],) + (1,) * (array.ndim - 1)
        quantized = (array / scale.reshape(shape)).astype(ml_dtypes.float8_e4m3fn)
        # the fp8 tensor takes a new name; DequantizeLinear rebuilds the original one
        init.CopyFrom(numpy_helper.from_array(quantized, name + "_fp8"))
        new_initializers.append(numpy_helper.from_array(scale, name + "_scale"))
        dequantize.append(
            helper.make_node("DequantizeLinear", [name + "_fp8", name + "_scale"], [name],
                             name=f"{name}_dequantize", axis=0)
        )
        converted += 1

    # the DequantizeLinear nodes read only initializers, so they can run first
    nodes = list(dequantize) + list(model.graph.node)
    model.graph.ClearField("node")
    model.graph.node.extend(nodes)
    model.graph.initializer.extend(new_initializers)
    print(f"fp8: converted {converted} weight tensors")
    return save_model(model, "fp8")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("variant", choices=["fp16", "int8", "int4", "fp8", "all"])
    args = parser.parse_args()

    steps = (
        [convert_fp16, convert_int8, convert_int4, convert_fp8]
        if args.variant == "all"
        else [globals()[f"convert_{args.variant}"]]
    )
    for step in steps:
        print(f"-> {step.__name__}")
        print("   ", step())


if __name__ == "__main__":
    main()