"""
Detection validator: run NXVA YOLO on val set and compute mAP + confusion matrix.
Uses Ultralytics DetMetrics and ConfusionMatrix.
"""

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

import numpy as np
import torch

from .dataloader import DatasetLoader
from .metrics import ConfusionMatrix, DetMetrics, box_iou
from .utils_val import det_gt_from_yolo_labels, match_predictions


class 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,
    ):
        self.model = model
        self.val_txt_path = val_txt_path
        self.batch_size = batch_size
        self.save_dir = Path(save_dir)
        self.save_dir.mkdir(parents=True, exist_ok=True)
        self.plots = plots
        self.conf = conf
        self.iou = iou
        self.names = names or {}
        self.nc = nc
        self.img_size = img_size
        self.verbose = verbose

        self.iouv = torch.linspace(0.5, 0.95, 10)
        self.niou = self.iouv.numel()
        self.seen = 0
        self.metrics = DetMetrics(names=self.names)
        self.confusion_matrix = ConfusionMatrix(names=self.names, task="detect", save_matches=False)
        self.speed = {"preprocess": 0.0, "inference": 0.0, "postprocess": 0.0}

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

    def _results_to_pred_dict(self, result: Any) -> Dict[str, torch.Tensor]:
        """Convert one Results instance to dict with bboxes (xyxy), conf, cls (tensors)."""
        if result.boxes is None or len(result.boxes) == 0:
            return {
                "bboxes": torch.zeros((0, 4), dtype=torch.float32),
                "conf": torch.zeros(0, dtype=torch.float32),
                "cls": torch.zeros(0, dtype=torch.long),
            }
        data = result.boxes.data
        if isinstance(data, np.ndarray):
            data = torch.from_numpy(data)
        bboxes = data[:, :4].float()
        conf = data[:, 4].float()
        cls = data[:, 5].long()
        return {"bboxes": bboxes, "conf": conf, "cls": cls}

    def _process_batch(self, pred: Dict[str, torch.Tensor], gt_bboxes: torch.Tensor, gt_cls: torch.Tensor) -> Dict[str, np.ndarray]:
        if gt_cls.shape[0] == 0 or pred["cls"].shape[0] == 0:
            return {"tp": np.zeros((pred["cls"].shape[0], self.niou), dtype=bool)}
        iou = box_iou(gt_bboxes, pred["bboxes"])
        tp = match_predictions(pred["cls"], gt_cls, iou, self.iouv)
        return {"tp": tp.cpu().numpy()}

    def __call__(self) -> Dict[str, Any]:
        loader = self._get_loader()
        # Use YOLOToolBox if available (needs img_paths for Results)
        use_toolbox = hasattr(self.model, "__construct_results")
        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)
                # Build minimal Results-like list: each pred can be raw array (detect: (N,6))
                results = []
                for i, pred in enumerate(preds):
                    from nxva.toolbox.result import Results
                    if isinstance(pred, np.ndarray) and pred.ndim == 2 and pred.shape[1] >= 6:
                        boxes = torch.from_numpy(pred[:, :6])
                        results.append(Results(batch_imgs[i], batch_paths[i], self.names, boxes=boxes))
                    else:
                        results.append(Results(batch_imgs[i], batch_paths[i], self.names, boxes=torch.zeros((0, 6))))

            n_batch = len(results)
            for si in range(n_batch):
                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]
                gt_bboxes, gt_cls = det_gt_from_yolo_labels(gt, orig_shape)
                predn = self._get_results_pred(result)
                stat = self._process_batch(predn, gt_bboxes, gt_cls)
                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

    def _get_results_pred(self, result: Any) -> Dict[str, torch.Tensor]:
        return self._results_to_pred_dict(result)

    def _print_results(self) -> None:
        keys = getattr(self.metrics, "keys", ["metrics/precision(B)", "metrics/recall(B)", "metrics/mAP50(B)", "metrics/mAP50-95(B)"])
        mr = self.metrics.mean_results()
        nt = getattr(self.metrics, "nt_per_class", None)
        total_inst = nt.sum() if nt is not None else 0
        pf = "%22s" + "%11i" * 2 + "%11.3g" * len(keys)
        print(pf % ("all", self.seen, int(total_inst), *mr))
        if self.nc > 1 and hasattr(self.metrics, "ap_class_index") and len(self.metrics.ap_class_index):
            for i, c in enumerate(self.metrics.ap_class_index):
                ni = getattr(self.metrics, "nt_per_image", None)
                nc_per = self.metrics.nt_per_class[c] if nt is not None else 0
                ni_c = ni[c] if ni is not None else 0
                print(pf % (self.names.get(c, str(c)), ni_c, nc_per, *self.metrics.class_result(i)))
