import functools
from pathlib import Path
from typing import Any, cast

from ai_parade._toolkit.data.output import ModelOutput
from ai_parade._toolkit.structural_mapping import MappingDefinition, structural_map

from ._PascalVOC import PascalVOC
from .Dataset import DatasetSplit

id_to_info: dict[str, tuple[int, str]] | None = None
index_to_info: dict[int, tuple[str, str]] | None = None


def _map_ImageNet_to_ModelOutput(
	raw: Any, mapping: MappingDefinition, mapping_path: Path
) -> ModelOutput:
	if raw == {}:
		return {}
	if not id_to_info:
		_load_imageNet_mapping(mapping_path)
		assert id_to_info

	# make sure the detections are a list
	if raw.get("annotation", {}).get("object", None) is not None:
		if not isinstance(raw["annotation"]["object"], list):
			raw["annotation"]["object"] = [raw["annotation"]["object"]]
	else:
		raw["annotation"]["object"] = []

	mapped = cast(ModelOutput, structural_map(raw, mapping, skip_missing=True))
	if (
		"detections" in mapped
		and len(mapped["detections"]) == 0
		and "class_ImageNet1k_id" not in mapped
	):
		return mapped

	# use class id from detection if not present
	if (
		"class_ImageNet1k_id" not in mapped
		and "detections" in mapped
		and "class_ImageNet1k_id" in mapped["detections"][0]
	):
		mapped["class_ImageNet1k_id"] = mapped["detections"][0]["class_ImageNet1k_id"]

	# map class id to index: in both classification and detection
	if "class_ImageNet1k_id" in mapped:
		if mapped["class_ImageNet1k_id"] in id_to_info:
			mapped["class_ImageNet1k"] = id_to_info[mapped["class_ImageNet1k_id"]][0]
	if "detections" in mapped:
		for detection in mapped["detections"]:
			if "class_ImageNet1k_id" in detection:
				if detection["class_ImageNet1k_id"] in id_to_info:
					detection["class_ImageNet1k"] = id_to_info[
						detection["class_ImageNet1k_id"]
					][0]
	return mapped


def _load_imageNet_mapping(path: Path):
	global id_to_info, index_to_info
	id_to_info = {}
	index_to_info = {}

	with open(path, "r") as f:
		for index, line in enumerate(f):
			line = line.strip()
			if not line:
				continue

			wnid, labels = line.split(" ", 1)

			id_to_info[wnid] = (index, labels)
			index_to_info[index] = (wnid, labels)


def ImageNet_VID(
	path: Path = Path("datasets/ImageNet/ILSVRC2015_VID"),
	split: DatasetSplit = DatasetSplit.train,
):
	if (path / "ImageNet").exists():
		path = path / "ImageNet"
	if (path / "ILSVRC2015_VID").exists():
		path = path / "ILSVRC2015_VID"
	return PascalVOC(
		path,
		split,
		functools.partial(
			_map_ImageNet_to_ModelOutput,
			mapping={
				"annotation": {
					"object": [
						{
							"name": "class_ImageNet1k_id",
							"bndbox": {
								"xmin": "bbox_xmin_absolute",
								"xmax": "bbox_xmax_absolute",
								"ymin": "bbox_ymin_absolute",
								"ymax": "bbox_ymax_absolute",
							},
						},
						"...repeat as detections",
					]
				}
			},
			mapping_path=path.parent / "LOC_synset_mapping.txt",
		),
	)


def ImageNet(
	path: Path = Path("datasets/ImageNet/ILSVRC"),
	split: DatasetSplit = DatasetSplit.train,
):
	if (path / "ImageNet").exists():
		path = path / "ImageNet"
	if (path / "ILSVRC").exists():
		path = path / "ILSVRC"

	return PascalVOC(
		path,
		split,
		functools.partial(
			_map_ImageNet_to_ModelOutput,
			mapping={
				"annotation": {
					"object": [
						{
							"name": "class_ImageNet1k_id",
							"bndbox": {
								"xmin": "bbox_xmin_absolute",
								"xmax": "bbox_xmax_absolute",
								"ymin": "bbox_ymin_absolute",
								"ymax": "bbox_ymax_absolute",
							},
						},
						"...repeat as detections",
					]
				}
			},
			mapping_path=path.parent / "LOC_synset_mapping.txt",
		),
	)
