"""Shared locked-protocol dataset, model, and COCO evaluator utilities."""
from __future__ import annotations

import contextlib
import io
import json
import xml.etree.ElementTree as ET
from pathlib import Path

import torch
from PIL import Image
from torch.utils.data import Dataset

CLASS_NAMES = ("car", "motorbike", "bicycle", "chair", "diningtable", "bottle", "tvmonitor", "bus")
EXPECTED = {"train": 1647, "val": 183, "test": 400}


def resolve_image(dataset_root: Path, domain: str, basename: str) -> Path:
    folder = "RGB_normal" if domain == "bright" else "RGB_Dark"
    hits = list(dataset_root.rglob(f"{folder}/{basename}"))
    if len(hits) != 1:
        hits = [p for p in dataset_root.rglob(basename) if folder.lower() in str(p).lower()]
    if len(hits) != 1:
        raise RuntimeError(f"expected exactly one {domain} image for {basename}, found {hits}")
    return hits[0]


def read_yolo(path: Path, width: int, height: int):
    boxes, labels = [], []
    for line in path.read_text().splitlines():
        if not line.strip():
            continue
        cls, cx, cy, bw, bh = map(float, line.split())
        if not 0 <= int(cls) < len(CLASS_NAMES):
            raise RuntimeError(f"invalid class id {cls} in {path}")
        boxes.append([(cx - bw / 2) * width, (cy - bh / 2) * height,
                      (cx + bw / 2) * width, (cy + bh / 2) * height])
        labels.append(int(cls))
    return (torch.tensor(boxes, dtype=torch.float32).reshape(-1, 4),
            torch.tensor(labels, dtype=torch.int64))


def read_voc(path: Path):
    boxes, labels = [], []
    for obj in ET.parse(path).getroot().findall("object"):
        name = obj.findtext("name")
        if name not in CLASS_NAMES:
            raise RuntimeError(f"unknown VOC class {name!r} in {path}")
        box = obj.find("bndbox")
        boxes.append([float(box.findtext("xmin")), float(box.findtext("ymin")),
                      float(box.findtext("xmax")), float(box.findtext("ymax"))])
        labels.append(CLASS_NAMES.index(name))
    return (torch.tensor(boxes, dtype=torch.float32).reshape(-1, 4),
            torch.tensor(labels, dtype=torch.int64))


class LockedLODDataset(Dataset):
    def __init__(self, manifest: Path, dataset_root: Path, labels_root: Path, limit: int | None = None):
        payload = json.loads(Path(manifest).read_text())
        self.domain = payload["domain"]
        self.split = payload["split"]
        self.items = payload["items"][:limit]
        if len(payload["items"]) != EXPECTED[self.split]:
            raise RuntimeError(f"{manifest}: expected {EXPECTED[self.split]} locked rows, got {len(payload['items'])}")
        self.dataset_root, self.labels_root = Path(dataset_root), Path(labels_root)

    def __len__(self):
        return len(self.items)

    def __getitem__(self, index):
        item = self.items[index]
        image = Image.open(resolve_image(self.dataset_root, self.domain, item["image_basename"])).convert("RGB")
        # The public LOD Kaggle dataset contains the source VOC annotations.
        # Prefer them, retaining copied YOLO labels only as a local fallback.
        group = "RGB-normal-Annotations" if self.domain == "bright" else "RGB-dark-Annotations"
        voc = self.dataset_root / group / group / (Path(item["image_basename"]).stem + ".xml")
        if voc.is_file():
            boxes, labels = read_voc(voc)
        else:
            label_domain = "normal" if self.domain == "bright" else "dark"
            label = self.labels_root / label_domain / self.split / "labels" / item["label_basename"]
            if not label.is_file():
                raise FileNotFoundError(f"missing VOC and fallback label for {item['pair_id']}")
            boxes, labels = read_yolo(label, *image.size)
        return image, {"boxes": boxes, "labels": labels,
                       "image_id": torch.tensor(index),
                       "orig_size": torch.tensor([image.height, image.width])}


def collate(batch):
    return tuple(zip(*batch))


def build_model(cfg):
    detector = cfg["detector"]
    if detector == "rtdetr":
        from transformers import RTDetrV2ForObjectDetection, RTDetrImageProcessor
        processor = RTDetrImageProcessor.from_pretrained(cfg["pretrained"])
        model = RTDetrV2ForObjectDetection.from_pretrained(
            cfg["pretrained"], num_labels=len(CLASS_NAMES),
            id2label=dict(enumerate(CLASS_NAMES)), label2id={n: i for i, n in enumerate(CLASS_NAMES)},
            ignore_mismatched_sizes=True,
        )
        return model, processor, "hf"
    if detector == "fasterrcnn_r50_fpn":
        from torchvision.models.detection import fasterrcnn_resnet50_fpn_v2, FasterRCNN_ResNet50_FPN_V2_Weights
        from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
        model = fasterrcnn_resnet50_fpn_v2(weights=FasterRCNN_ResNet50_FPN_V2_Weights.COCO_V1)
        model.roi_heads.box_predictor = FastRCNNPredictor(model.roi_heads.box_predictor.cls_score.in_features, len(CLASS_NAMES) + 1)
        return model, None, "torchvision"
    if detector == "ssdlite320_mobilenet_v3_large":
        from torchvision.models.detection import ssdlite320_mobilenet_v3_large, SSDLite320_MobileNet_V3_Large_Weights
        from torchvision.models.detection.ssdlite import SSDLiteClassificationHead
        model = ssdlite320_mobilenet_v3_large(weights=SSDLite320_MobileNet_V3_Large_Weights.COCO_V1)
        # TorchVision exposes each SSDLite predictor as a Sequential block;
        # its depthwise Conv2d is nested under block[0][0].  Derive the
        # feature widths from the loaded COCO head rather than hard-coding a
        # version-specific list.
        in_channels = [module[0][0].in_channels for module in model.head.classification_head.module_list]
        num_anchors = model.anchor_generator.num_anchors_per_location()
        model.head.classification_head = SSDLiteClassificationHead(
            in_channels, num_anchors, len(CLASS_NAMES) + 1, torch.nn.BatchNorm2d
        )
        return model, None, "torchvision"
    raise ValueError(f"unsupported detector: {detector}")


def image_tensors(images, device):
    from torchvision.transforms.functional import pil_to_tensor
    return [pil_to_tensor(image).float().div(255).to(device) for image in images]


def hf_batch(processor, images, targets, device):
    annotations = []
    for index, target in enumerate(targets):
        boxes = target["boxes"]
        xywh = torch.stack((boxes[:, 0], boxes[:, 1], boxes[:, 2] - boxes[:, 0], boxes[:, 3] - boxes[:, 1]), 1)
        annotations.append({"image_id": index, "annotations": [
            {"bbox": box.tolist(), "category_id": int(label), "area": float(box[2] * box[3]), "iscrowd": 0}
            for box, label in zip(xywh, target["labels"])
        ]})
    encoded = processor(images=list(images), annotations=annotations, return_tensors="pt")
    return {key: (value.to(device) if hasattr(value, "to") else [{k: v.to(device) for k, v in x.items()} for x in value])
            for key, value in encoded.items()}


class CanonicalMAP:
    def __init__(self):
        self.images, self.annotations, self.predictions, self.annotation_id = {}, [], [], 1

    @staticmethod
    def _xywh(box):
        x1, y1, x2, y2 = map(float, box.tolist())
        return [x1, y1, max(0.0, x2 - x1), max(0.0, y2 - y1)]

    def update(self, predictions, targets):
        for pred, target in zip(predictions, targets):
            image_id = int(target["image_id"])
            h, w = map(int, target["orig_size"].tolist())
            self.images[image_id] = {"id": image_id, "height": h, "width": w}
            for box, label in zip(target["boxes"], target["labels"]):
                bbox = self._xywh(box)
                self.annotations.append({"id": self.annotation_id, "image_id": image_id, "category_id": int(label) + 1,
                                         "bbox": bbox, "area": bbox[2] * bbox[3], "iscrowd": 0})
                self.annotation_id += 1
            for box, score, label in zip(pred["boxes"], pred["scores"], pred["labels"]):
                self.predictions.append({"image_id": image_id, "category_id": int(label) + 1,
                                         "bbox": self._xywh(box), "score": float(score)})

    def compute(self):
        from pycocotools.coco import COCO
        from pycocotools.cocoeval import COCOeval
        gt = COCO(); gt.dataset = {"images": list(self.images.values()), "annotations": self.annotations,
                                   "categories": [{"id": i + 1, "name": n} for i, n in enumerate(CLASS_NAMES)]}; gt.createIndex()
        dt = gt.loadRes(self.predictions) if self.predictions else COCO()
        if not self.predictions:
            dt.dataset = {"images": list(self.images.values()), "annotations": [], "categories": gt.dataset["categories"]}; dt.createIndex()
        evaluator = COCOeval(gt, dt, "bbox"); evaluator.params.imgIds = sorted(self.images); evaluator.params.catIds = list(range(1, 9)); evaluator.params.maxDets = [1, 10, 100]
        with contextlib.redirect_stdout(io.StringIO()):
            evaluator.evaluate(); evaluator.accumulate(); evaluator.summarize()
        return {"map50_95": float(evaluator.stats[0]), "map50": float(evaluator.stats[1])}