from typing import TYPE_CHECKING, Any, override

from ai_parade._toolkit.run.ONNX import ONNXRunner

if TYPE_CHECKING:
	import torch  # type: ignore  # noqa: F401

from ai_parade._toolkit.enums import ModelFormat

from ._onnx_simplify import ONNXSimplify
from ._pytorch_base import PyTorchConvertBase
from .converter import Converter, FormatConverter


class PyTorch_onnx_export(FormatConverter):
	"""
	Built-in converter from PyTorch to ONNX.
	"""

	@property
	@override
	def format_conversion(self):
		return (ModelFormat.PyTorch, ModelFormat.ONNX)

	@override
	def create(self, options: Converter.Options):
		super().create(options)
		return self.PyTorchConverter(self.conversion, options)

	class PyTorchConverter(PyTorchConvertBase):
		@override
		def _export_call(self, params: PyTorchConvertBase.Params):
			model, sample_inputs, output_weights_path, *_ = params

			import torch  # type: ignore
			from onnxscript import opset18 as op  # type: ignore

			def sym_not(x: Any):
				return op.Not(x)

			# This notation is kind of cryptic but:
			# - tuple define model input argument (could be dict instead to act as kwargs)
			# - key in the dictionary is the index of its dimensions
			# - value is the dimension
			dynamic_shapes = ({0: torch.export.Dim("batch_dim", min=1)},)

			onnx_model = torch.onnx.export(
				model,
				sample_inputs,
				external_data=False,  # ?todo: at least for now this large models are useless for us
				dynamo=True,
				report=True,
				artifacts_dir=output_weights_path.parent,
				verify=True,
				dynamic_shapes=dynamic_shapes,
				input_names=[ONNXRunner.INPUT_NAME],
				# todo after https://github.com/pytorch/pytorch/issues/136572 : remove
				custom_translation_table={torch.sym_not: sym_not},
			)
			if not onnx_model:
				raise RuntimeError("Failed to export ONNX model")
			onnx_model.save(output_weights_path)

			onnx_model.release()
			onnx_model = None

			if not ONNXSimplify().simplify(output_weights_path):
				raise RuntimeError("Failed to simplify ONNX model")
