"""
Segment validator: bbox mAP + mask mAP + confusion matrix.
Uses Ultralytics SegmentMetrics and ConfusionMatrix.
"""

from pathlib import Path
from typing import Any, Dict, List, Optional, Union

import numpy as np
import torch

from .dataloader import (
    DatasetLoader,
    segmentation_labels_to_bboxes,
    segmentation_labels_to_masks,
)
from .metrics import ConfusionMatrix, SegmentMetrics, box_iou, mask_iou
from .utils_val import match_predictions
from .val_detection import DetectionValidator


class SegmentValidator(DetectionValidator):
    def __init__(
        self,
        model: Any,
        val_txt_path: str,
        batch_size: int = 1,
        save_dir: Union[str, Path] = ".",
        plots: bool = False,
        conf: float = 0.25,
        iou: float = 0.45,
        names: Optional[Dict[int, str]] = None,
        nc: int = 80,
        img_size: tuple = (640, 640),
        verbose: bool = True,
    ):
        super().__init__(
            model=model,
            val_txt_path=val_txt_path,
            batch_size=batch_size,
            save_dir=save_dir,
            plots=plots,
            conf=conf,
            iou=iou,
            names=names,
            nc=nc,
            img_size=img_size,
            verbose=verbose,
        )
        self.metrics = SegmentMetrics(names=self.names)

    def _get_loader(self) -> DatasetLoader:
        return DatasetLoader(
            self.val_txt_path,
            task="segment",
            num_kpts=None,
            batch_size=self.batch_size,
            img_size=self.img_size,
        )

    def _get_results_pred(self, result: Any) -> Dict[str, torch.Tensor]:
        d = super()._get_results_pred(result)
        if result.masks is not None and result.masks.data is not None and len(result.masks) > 0:
            m = result.masks.data
            if isinstance(m, np.ndarray):
                m = torch.from_numpy(m)
            # mask_iou expects (N, H*W); masks are (N, H, W)
            d["masks"] = m.flatten(1).float()
        else:
            d["masks"] = torch.zeros(0, 1, device=d["bboxes"].device)
        return d

    def _process_batch(self, pred: Dict[str, torch.Tensor], gt_bboxes: torch.Tensor, gt_cls: torch.Tensor, gt_masks: Optional[torch.Tensor] = None) -> Dict[str, np.ndarray]:
        out = super()._process_batch(pred, gt_bboxes, gt_cls)
        if gt_masks is not None and gt_masks.shape[0] > 0 and pred.get("masks") is not None and pred["masks"].shape[0] > 0:
            iou_m = mask_iou(gt_masks, pred["masks"])
            tp_m = match_predictions(pred["cls"], gt_cls, iou_m, self.iouv)
            out["tp_m"] = tp_m.cpu().numpy()
        else:
            n_pred = pred["cls"].shape[0]
            out["tp_m"] = np.zeros((n_pred, self.niou), dtype=bool)
        return out

    def __call__(self) -> Dict[str, Any]:
        loader = self._get_loader()
        use_toolbox = hasattr(self.model, "__construct_results")
        input_shape = (self.img_size[0], self.img_size[1]) if isinstance(self.img_size[0], int) else self.img_size

        for batch_imgs, batch_gt, batch_paths in loader:
            if use_toolbox:
                results = self.model(batch_imgs, batch_paths)
            else:
                preds = self.model(batch_imgs)
                from nxva.toolbox.result import Results
                results = []
                for i, pred in enumerate(preds):
                    if isinstance(pred, dict) and "boxes" in pred and "mask" in pred:
                        boxes = torch.from_numpy(pred["boxes"])
                        masks = torch.from_numpy(pred["mask"]) if pred["mask"] is not None else torch.zeros(0, *input_shape)
                        results.append(Results(batch_imgs[i], batch_paths[i], self.names, boxes=boxes, masks=masks))
                    else:
                        results.append(Results(batch_imgs[i], batch_paths[i], self.names, boxes=torch.zeros((0, 6))))

            for si in range(len(results)):
                result = results[si]
                gt = batch_gt[si]
                self.seen += 1
                orig_shape = result.orig_shape if hasattr(result, "orig_shape") else batch_imgs[si].shape[:2]
                seg_labels = gt if isinstance(gt, list) else (gt.tolist() if hasattr(gt, "tolist") else [])
                gt_bboxes = torch.from_numpy(segmentation_labels_to_bboxes(seg_labels, orig_shape))
                gt_masks_np = segmentation_labels_to_masks(seg_labels, orig_shape, input_shape)
                gt_masks = torch.from_numpy(gt_masks_np).float().flatten(1)
                gt_cls = torch.tensor([int(x["class_id"]) for x in seg_labels], dtype=torch.long) if seg_labels else torch.zeros(0, dtype=torch.long)
                if gt_bboxes.shape[0] == 0:
                    gt_cls = torch.zeros(0, dtype=torch.long)

                predn = self._get_results_pred(result)
                stat = self._process_batch(predn, gt_bboxes, gt_cls, gt_masks)
                cls_np = gt_cls.cpu().numpy()
                stat["target_cls"] = cls_np
                stat["target_img"] = np.unique(cls_np) if cls_np.size else np.array([], dtype=np.int64)
                stat["conf"] = predn["conf"].cpu().numpy() if predn["conf"].numel() else np.zeros(0)
                stat["pred_cls"] = predn["cls"].cpu().numpy() if predn["cls"].numel() else np.zeros(0, dtype=np.int64)
                self.metrics.update_stats(stat)
                if self.plots and gt_bboxes.numel() > 0 and predn["cls"].numel() > 0:
                    pbatch = {"cls": gt_cls, "bboxes": gt_bboxes}
                    self.confusion_matrix.process_batch(predn, pbatch, conf=self.conf, iou_thres=self.iou)

        self.metrics.process(save_dir=self.save_dir, plot=self.plots)
        if self.plots:
            for normalize in (True, False):
                self.confusion_matrix.plot(save_dir=str(self.save_dir), normalize=normalize)
        self.metrics.confusion_matrix = self.confusion_matrix
        self.metrics.save_dir = self.save_dir
        if self.verbose:
            self._print_results()
        return self.metrics.results_dict
