from dataclasses import replace
from typing import TYPE_CHECKING, Literal, override

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

from ai_parade._toolkit.data.ImageInput import ImageInput
from ai_parade._toolkit.enums import (
	ExactModelFormat,
	ModelFormat,
	Quantization,
)

from ._pytorch_base import PyTorchConvertBase
from .converter import Converter


class AIEdgeConverterEntry(Converter):
	"""
	https://github.com/google-ai-edge/ai-edge-torch
	and
	https://github.com/google-ai-edge/ai-edge-torch/blob/main/docs/pytorch_converter/README.md#quantization
	"""

	def __init__(
		self,
		quantization_type: Literal[Quantization.none]
		| Literal[Quantization.float16]
		| Literal[Quantization.int8],
	):
		self.quantization_type = quantization_type

	@property
	@override
	def conversion(self):
		return (
			ExactModelFormat(ModelFormat.PyTorch),
			ExactModelFormat(ModelFormat.LiteRT, self.quantization_type),
		)

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

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

			import ai_edge_torch  # type: ignore
			import tensorflow as tf  # type: ignore

			# set it to channels last mode, so that the model does not contain transpose layers (the backend does work in NHWC)
			model = ai_edge_torch.to_channel_last_io(model, args=[0])
			model.eval()  # apparently the mode is reset by the to_channel_last_io
			sample_inputs = (sample_inputs[0].permute(0, 2, 3, 1),)

			converter_flags = {}
			match self.conversion[1].quantization:
				case Quantization.float32:
					pass

				case Quantization.int8:
					assert self.options.calibration_data is not None
					# with fallback to other datatypes
					converter_flags["optimizations"] = [tf.lite.Optimize.DEFAULT]
					converter_flags["representative_dataset"] = list(
						map(
							lambda x: [x[0].transpose(0, 2, 3, 1)],
							self.options.calibration_data,
						)
					).__iter__
				case Quantization.float16:
					converter_flags = {
						"optimizations": [tf.lite.Optimize.DEFAULT],
						"target_spec": {"supported_types": [tf.float16]},
					}

				case _:
					raise NotImplementedError(
						f"Quantization {self.conversion[1].quantization} not supported"
					)

			edge_model = ai_edge_torch.convert(
				model, sample_inputs, _ai_edge_converter_flags=converter_flags
			)
			edge_model.export(str(output_weights_path))

			default_results.image_input = replace(
				default_results.image_input, dataOrder=ImageInput.DataOrder.NHWC
			)
			return default_results


capabilities: list[Converter] = [
	AIEdgeConverterEntry(Quantization.none),
	AIEdgeConverterEntry(Quantization.int8),
	AIEdgeConverterEntry(Quantization.float16),
]
