"""Supervised classification from fixed-size YOLO tiles and human box labels."""
from __future__ import annotations

import math
from pathlib import Path

import yaml
from PIL import Image

from datasets.classification import FolderClassificationDataset
from datasets.transforms import build_transforms


def tile_label(text: str, num_classes: int, background: int) -> int:
    boxes = []
    for line in text.splitlines():
        if not line.strip():
            continue
        parts = line.split()
        if len(parts) != 5:
            raise ValueError(f"Malformed YOLO label: {line}")
        cls, cx, cy, w, h = map(float, parts)
        if (not all(math.isfinite(v) for v in (cls, cx, cy, w, h))
                or not cls.is_integer() or not 0 <= cls < num_classes
                or not 0 <= cx <= 1 or not 0 <= cy <= 1 or not 0 < w <= 1 or not 0 < h <= 1):
            raise ValueError(f"Invalid YOLO label: {line}")
        boxes.append((int(cls), w * h))
    return max(boxes, key=lambda b: b[1])[0] if boxes else background


class YoloTileClassificationDataset(FolderClassificationDataset):
    """Keep the existing YOLO split. Empty labels are normal; largest box wins.

    No dataset transform may resample the image. The size is checked on every read.
    """

    def __init__(self, yaml_path, split, image_size=(128, 128), transforms=(), classes=None):
        path = Path(yaml_path).resolve()
        cfg = yaml.safe_load(path.read_text())
        names = cfg["names"]
        self.classes = ([names.get(i, names.get(str(i))) for i in range(len(names))]
                        if isinstance(names, dict) else list(names))
        if not self.classes or any(not isinstance(n, str) for n in self.classes):
            raise ValueError("Invalid YOLO names mapping")
        if classes is not None and list(classes) != self.classes:
            raise ValueError("Classifier class order must match YOLO names")
        background = self.classes.index("normal")
        self.class_to_idx = dict(zip(self.classes, range(len(self.classes)), strict=True))
        self.root = Path(cfg.get("path", path.parent))
        if not self.root.is_absolute():
            self.root = path.parent / self.root
        folder = self.root / cfg[split]
        if not folder.is_dir():
            raise FileNotFoundError(folder)
        relative = folder.relative_to(self.root)
        if relative.parts[0] != "images":
            raise ValueError("YOLO split must be beneath images/")
        labels = self.root / "labels" / Path(*relative.parts[1:])
        allowed = {"ToTensor", "RandomHorizontalFlip", "RandomVerticalFlip", "PerImageStandardize"}
        if any(t["name"] not in allowed for t in transforms):
            raise ValueError("Fixed tiles forbid resampling transforms")
        self.transforms = build_transforms(transforms)
        self.image_size = tuple(image_size)
        self.samples = []
        for image in sorted(folder.rglob("*.png")):
            label = labels / image.relative_to(folder).with_suffix(".txt")
            if not label.is_file():
                raise FileNotFoundError(f"Missing label (not an implicit background): {label}")
            self.samples.append((image, tile_label(label.read_text(), len(self.classes), background)))
        if not self.samples:
            raise ValueError(f"No PNG tiles in {folder}")

    def __getitem__(self, index):
        path, label = self.samples[index]
        with Image.open(path) as raw:
            if (raw.height, raw.width) != self.image_size:
                raise ValueError(f"{path}: size {raw.size}, expected {self.image_size}; resize is forbidden")
            image = raw.convert("RGB")
        sample = self.transforms({"image": image})
        return {"image": sample["image"], "label": label,
                "meta": {"path": str(path), "index": index, "class": self.classes[label]}}