from dataclasses import replace
from pathlib import Path
from typing import TYPE_CHECKING, override

from ai_parade._toolkit.converters._pytorch_base import PyTorchConvertBase
from ai_parade._toolkit.data.ImageInput import ImageInput
from ai_parade._toolkit.enums import (
	ExactModelFormat,
	HardwareAcceleration,
	ModelFormat,
	Quantization,
)

from .converter import ConversionResult, Converter

if TYPE_CHECKING:
	import torch  # type: ignore


class Executorch(Converter):
	"""
	https://docs.pytorch.org/executorch/main/using-executorch-export.html
	"""

	def __init__(
		self,
		quantization: Quantization,
		hw_acceleration: HardwareAcceleration,
		soc: tuple[str, ...] | None,
	):
		self._quantization = quantization
		self._hw_acceleration = hw_acceleration
		self._soc = soc

	chip_compatibility_lists: list[tuple[str, tuple[str, ...]]] = [
		("SM8450", ("SM8450",)),
		("SM8550", ("SM8550", "SM8550-AC")),
		("SM8650", ("SM8650", "SM8650-AC")),
		("SM8750", ("SM8750", "SM8750-AC")),
	]

	@property
	@override
	def conversion(self):
		return (
			ExactModelFormat(ModelFormat.PyTorch),
			ExactModelFormat(
				ModelFormat.ExecutorTorch,
				self._quantization,
				self._hw_acceleration,
				self._soc,
			),
		)

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

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

			import torch  # type: ignore
			from executorch.exir import to_edge_transform_and_lower  # type: ignore
			from torchao.quantization.pt2e.quantize_pt2e import (  # type: ignore
				convert_pt2e,
				prepare_pt2e,
			)

			image_input = default_results.image_input
			if self.conversion[1].quantization == Quantization.float16:
				image_input = replace(image_input, dataType=ImageInput.DataType.float16)

			# todo
			# model = model.to(memory_format=torch.channels_last)  # pyright: ignore[reportCallIssue]
			# sample_inputs = (sample_inputs[0].to(memory_format=torch.channels_last),)

			# ?TODO: add support for dynamic shapes,
			# problem is the executorch doesn't support unbound dims (i.e. [1, inf])
			# and for some reason it complains with sample output batch dim set to 1:
			# "You marked batch_dim as dynamic but your code specialized it to be a constant (1)"
			dynamic_shapes = ()  # ({0: torch.export.Dim("batch_dim", min=1, max=8)},)
			match self.conversion[1].hw_acceleration_restriction:
				case HardwareAcceleration.XNNPACK:
					partitioner, quantizer = self.xnnpack_setup()
				case HardwareAcceleration.Vulkan:
					partitioner, quantizer = self.vulkan_setup()
				case HardwareAcceleration.QNN:
					partitioner, quantizer = self.qnn_setup()
				case _:
					raise NotImplementedError(
						f"{self.conversion[1].hw_acceleration_restriction} is not supported"
					)
			match self.conversion[1].quantization:
				case Quantization.none:
					pass
				case Quantization.float16:
					model = model.half()
				case Quantization.int8:
					assert self.options.calibration_data is not None

					training_ep = torch.export.export(
						model, sample_inputs, dynamic_shapes=dynamic_shapes
					).module()
					prepared_model = prepare_pt2e(training_ep, quantizer)

					# calibration
					for cal_sample, _, _ in self.options.calibration_data:
						prepared_model(torch.from_numpy(cal_sample))

					model = convert_pt2e(prepared_model)

			if (
				self.conversion[1].hw_acceleration_restriction
				== HardwareAcceleration.QNN
			):
				models = self.convert_qnn(model, sample_inputs, output_weights_path)
				return [
					ConversionResult(
						format=replace(self.conversion[1], chip_compatibility=[soc]),
						image_input=image_input,
						path=model,
						output=model_metadata.output,
					)
					for model, soc in models
				]
			else:
				assert partitioner is not None
				executorch_program = to_edge_transform_and_lower(
					torch.export.export(
						model,
						sample_inputs
						if self.conversion[1].quantization != Quantization.float16
						else (sample_inputs[0].half(),),
						dynamic_shapes=dynamic_shapes,
					),
					partitioner=[partitioner],
				).to_executorch()

				with open(output_weights_path, "wb") as file:
					executorch_program.write_to_file(file)

				return ConversionResult(
					format=self.conversion[1],
					image_input=image_input,
					path=output_weights_path,
					output=model_metadata.output,
				)

		def vulkan_setup(self):
			from executorch.backends.vulkan.partitioner.vulkan_partitioner import (  # type: ignore
				VulkanPartitioner,
			)
			from executorch.backends.vulkan.quantizer.vulkan_quantizer import (  # type: ignore
				VulkanQuantizer,
				get_symmetric_quantization_config,
			)

			qparams = get_symmetric_quantization_config(is_dynamic=False, weight_bits=8)
			quantizer = VulkanQuantizer()
			quantizer.set_global(qparams)

			return VulkanPartitioner(), quantizer

		def xnnpack_setup(self):
			from executorch.backends.xnnpack.partition.xnnpack_partitioner import (  # type: ignore
				XnnpackPartitioner,
			)
			from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import (  # type: ignore
				XNNPACKQuantizer,
				get_symmetric_quantization_config,
			)

			qparams = get_symmetric_quantization_config(is_per_channel=True)
			quantizer = XNNPACKQuantizer()
			quantizer.set_global(qparams)

			return XnnpackPartitioner(), quantizer

		def qnn_setup(self):
			from executorch.backends.qualcomm.quantizer.quantizer import (  # type: ignore
				QnnQuantizer,
			)

			return None, QnnQuantizer()

		def convert_qnn(
			self,
			model: "torch.nn.Module",
			sample_inputs: tuple["torch.Tensor"],
			output_weights_path: Path,
		):
			from executorch.backends.qualcomm.utils.utils import (  # type: ignore
				generate_htp_compiler_spec,
				generate_qnn_executorch_compiler_spec,
				get_soc_to_chipset_map,
				to_edge_transform_and_lower_to_qnn,
			)

			output_model_paths: list[tuple[Path, str]] = []

			soc_str = next(
				x[0]
				for x in Executorch.chip_compatibility_lists
				if x[1] == self.conversion[1].chip_compatibility
			)
			soc = get_soc_to_chipset_map()[soc_str]
			try:
				# QNN Compiler Spec
				compile_spec = generate_qnn_executorch_compiler_spec(
					soc_model=soc,  # Your target SoC
					backend_options=generate_htp_compiler_spec(  # HTP Compiler Configuration
						use_fp16=self.conversion[1].quantization
						== Quantization.float16,
					),
				)
				# Lower to QNN backend
				executorch_program = to_edge_transform_and_lower_to_qnn(
					model,
					sample_inputs,
					compile_spec,
				).to_executorch()

				output_model = output_weights_path / soc.name
				with open(output_model, "wb") as file:
					file.write(executorch_program.buffer)

				output_model_paths.append((output_model, soc.name))
			except Exception as e:
				print(e)

			return output_model_paths


capabilities = [
	Executorch(quantization, hw_acceleration, soc)
	for quantization in Quantization
	for hw_acceleration in HardwareAcceleration
	for soc in (
		(x[1] for x in Executorch.chip_compatibility_lists)
		if hw_acceleration == HardwareAcceleration.QNN
		else [None]
	)
]
