import functools
import logging
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import override

import numpy as np

from ai_parade._toolkit.datasets.Dataset import Dataset
from ai_parade._toolkit.enums.formats import (
	ExactModelFormat,
	ModelFormat,
	Quantization,
)
from ai_parade._toolkit.logging import subprocess_run_log
from ai_parade._toolkit.run.ModelRunner import ModelRunner

from .converter import ConversionResult, Converter

logger = logging.getLogger(__name__)


class ncnnConverter(Converter):
	def __init__(self) -> None:
		super().__init__()

	@property
	@override
	def conversion(self):
		return (
			ExactModelFormat(ModelFormat.NCNN),
			ExactModelFormat(ModelFormat.NCNN, Quantization.int8),
		)

	@override
	def create(self, options: Converter.Options):
		assert options.calibration_data is not None
		return functools.partial(
			self.convert, calibration_data=options.calibration_data
		)

	def convert(
		self, runner: ModelRunner, calibration_data: Dataset
	) -> ConversionResult:
		with TemporaryDirectory() as d:
			shenanigans_dir = Path(d)
			calibration_files = []
			for i, (data, annotation, info) in enumerate(
				calibration_data.no_batch_iterator()
			):
				filename = shenanigans_dir / f"calibration_data{i}.npy"
				calibration_files.append(str(filename.resolve()))
				with open(filename, "wb") as f:
					np.save(f, data)
			with open(shenanigans_dir / "filelist.txt", "w") as f:
				f.write("\n".join(calibration_files))

			subprocess_run_log(
				[
					"ncnn2table",
					str(runner.weights / "model.param"),
					str(runner.weights / "model.bin"),
					str(shenanigans_dir / "filelist.txt"),
					str(shenanigans_dir / "calibration.table"),
					f"shape=[{','.join(map(str, reversed(data.shape)))}]",  # wtf is the shape reversed here (it does not work otherwise)
					"thread=8",
					"method=kl",
					"type=1",
				],
				logger=logger,
			)
			output_model_path = self.get_output_weights_path(runner, None)
			print(output_model_path)
			output_model_path.mkdir(parents=True, exist_ok=True)
			subprocess_run_log(
				[
					"ncnn2int8",
					str(runner.weights / "model.param"),
					str(runner.weights / "model.bin"),
					str(output_model_path / "model.param"),
					str(output_model_path / "model.bin"),
					str(shenanigans_dir / "calibration.table"),
				],
				logger=logger,
			)

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