import copy
import argparse
from optimum.exporters.onnx import onnx_export_from_model
from collections import OrderedDict
from typing import Dict
from optimum.exporters.onnx.model_configs import XLMRobertaOnnxConfig
from bgem3_model import BGEM3InferenceModel


class BGEM3OnnxConfig(XLMRobertaOnnxConfig):
    @property
    def outputs(self) -> Dict[str, Dict[int, str]]:
        """
        Dict containing the axis definition of the output tensors to provide to the model.

        Returns:
            `Dict[str, Dict[int, str]]`: A mapping of each output name to a mapping of axis position to the axes symbolic name.
        """
        return copy.deepcopy(
            OrderedDict(
                {
                    "dense_vecs": {0: "batch_size", 1: "embedding"},
                    "sparse_vecs": {0: "batch_size", 1: "token", 2: "weight"},
                    "colbert_vecs": {0: "batch_size", 1: "token", 2: "embedding"},
                }
            )
        )


def main(output: str, opset: int, device: str, optimize: str, atol: str):
    model = BGEM3InferenceModel()
    bgem3_onnx_config = BGEM3OnnxConfig(model.config)
    onnx_export_from_model(
        model,
        output=output,
        task="feature-extraction",
        custom_onnx_configs={"model": bgem3_onnx_config},
        opset=opset,
        optimize=optimize,
        atol=atol,
        device=device,
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--output",
        type=str,
        help="Path indicating the directory where to store the generated ONNX model.",
    )
    parser.add_argument(
        "--opset",
        type=int,
        default=None,
        help="If specified, ONNX opset version to export the model with. Otherwise, the default opset for the given model architecture will be used.",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cpu",
        help='The device to use to do the export. Defaults to "cpu".',
    )
    parser.add_argument(
        "--optimize",
        type=str,
        default=None,
        choices=["O1", "O2", "O3", "O4"],
        help=(
            "Allows to run ONNX Runtime optimizations directly during the export. Some of these optimizations are specific to ONNX Runtime, and the resulting ONNX will not be usable with other runtime as OpenVINO or TensorRT. Possible options:\n"
            "    - O1: Basic general optimizations\n"
            "    - O2: Basic and extended general optimizations, transformers-specific fusions\n"
            "    - O3: Same as O2 with GELU approximation\n"
            "    - O4: Same as O3 with mixed precision (fp16, GPU-only, requires `--device cuda`)"
        ),
    )
    parser.add_argument(
        "--atol",
        type=float,
        default=None,
        help="If specified, the absolute difference tolerance when validating the model. Otherwise, the default atol for the model will be used.",
    )
    args = parser.parse_args()

    main(args.output, args.opset, args.device, args.optimize, args.atol)