from typing import TYPE_CHECKING, override

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

from ai_parade._toolkit.enums import ModelFormat

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


class NobucoConverter(FormatConverter):
	"""
	https://github.com/AlexanderLutsenko/nobuco
	"""

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

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

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

			import nobuco  # type: ignore  # noqa: F401
			from nobuco import (  # type: ignore  # noqa: F401
				ChannelOrder,
				ChannelOrderingStrategy,
			)
			from nobuco.layers.weight import WeightLayer  # type: ignore  # noqa: F401

			keras_model = nobuco.pytorch_to_keras(
				model,
				args=sample_inputs,
				kwargs=None,
				inputs_channel_order=ChannelOrder.TENSORFLOW,
				outputs_channel_order=ChannelOrder.TENSORFLOW,
			)
			keras_model.save(output_weights_path, save_format="tf")
