"""
Pose validator: bbox mAP + keypoint mAP (OKS) + confusion matrix.
Uses Ultralytics PoseMetrics 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 OKS_SIGMA, ConfusionMatrix, PoseMetrics, box_iou, kpt_iou
from .utils_val import match_predictions
from .val_detection import DetectionValidator


def pose_gt_from_yolo_labels(labels: np.ndarray, orig_shape: tuple, num_kpts: int = 17) -> tuple:
    """
    Convert YOLO pose labels (N, 5 + num_kpts*3) to bboxes_xyxy, cls, keypoints (N, 17, 3).
    Returns (bboxes_xyxy, cls, keypoints) tensors.
    """
    if labels is None or (hasattr(labels, "shape") and labels.shape[0] == 0):
        return (
            torch.zeros((0, 4), dtype=torch.float32),
            torch.zeros(0, dtype=torch.long),
            torch.zeros((0, num_kpts, 3), dtype=torch.float32),
        )
    if hasattr(labels, "shape") and len(labels.shape) == 3:
        labels = labels[0]
    labels = np.asarray(labels, dtype=np.float64)
    valid = (labels[:, 2] > 0) & (labels[:, 3] > 0)
    labels = labels[valid]
    if labels.shape[0] == 0:
        return (
            torch.zeros((0, 4), dtype=torch.float32),
            torch.zeros(0, dtype=torch.long),
            torch.zeros((0, num_kpts, 3), dtype=torch.float32),
        )
    h, w = orig_shape[0], orig_shape[1]
    xc = labels[:, 1] * w
    yc = labels[:, 2] * h
    ww = labels[:, 3] * w
    hh = labels[:, 4] * h
    x1 = xc - ww / 2
    y1 = yc - hh / 2
    x2 = xc + ww / 2
    y2 = yc + hh / 2
    bboxes = np.stack([x1, y1, x2, y2], axis=1).astype(np.float32)
    cls = labels[:, 0].astype(np.int64)
    kpts = labels[:, 5:].reshape(-1, num_kpts, 3).astype(np.float32)
    kpts[:, :, 0] *= w
    kpts[:, :, 1] *= h
    return torch.from_numpy(bboxes), torch.from_numpy(cls), torch.from_numpy(kpts)


class PoseValidator(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,
        num_kpts: int = 17,
        sigma: Optional[List[float]] = None,
    ):
        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.num_kpts = num_kpts
        if sigma is not None:
            self.sigma = sigma
        else:
            # Default to OKS_SIGMA if it matches num_kpts, otherwise use a default value [0.026] * num_kpts
            if OKS_SIGMA is not None and len(OKS_SIGMA) == num_kpts:
                self.sigma = OKS_SIGMA.tolist()
            else:
                self.sigma = (np.ones(num_kpts) / num_kpts).tolist()
                # self.sigma = [0.026] * num_kpts

        self.metrics = PoseMetrics(names=self.names)

    def _get_loader(self) -> DatasetLoader:
        kpt_shape = (self.num_kpts, 3)
        return DatasetLoader(
            self.val_txt_path,
            task="pose",
            num_kpts=kpt_shape,
            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.keypoints is not None and result.keypoints.data is not None and len(result.keypoints) > 0:
            k = result.keypoints.data
            if isinstance(k, np.ndarray):
                k = torch.from_numpy(k)
            d["keypoints"] = k.float()
        else:
            d["keypoints"] = torch.zeros(0, self.num_kpts, 3, device=d["bboxes"].device)
        return d

    def _process_batch(self, pred: Dict[str, torch.Tensor], gt_bboxes: torch.Tensor, gt_cls: torch.Tensor, gt_kpts: Optional[torch.Tensor] = None) -> Dict[str, np.ndarray]:
        out = super()._process_batch(pred, gt_bboxes, gt_cls)
        if gt_kpts is not None and gt_kpts.shape[0] > 0 and pred.get("keypoints") is not None and pred["keypoints"].shape[0] > 0:
            area = (gt_bboxes[:, 2] - gt_bboxes[:, 0]) * (gt_bboxes[:, 3] - gt_bboxes[:, 1]) * 0.53
            iou_k = kpt_iou(gt_kpts, pred["keypoints"], area, self.sigma)
            tp_p = match_predictions(pred["cls"], gt_cls, iou_k, self.iouv)
            out["tp_p"] = tp_p.cpu().numpy()
        else:
            out["tp_p"] = np.zeros((pred["cls"].shape[0], self.niou), dtype=bool)
        return out

    def __call__(self) -> Dict[str, Any]:
        loader = self._get_loader()
        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)
                from nxva.toolbox.result import Results
                results = []
                for i, pred in enumerate(preds):
                    if isinstance(pred, dict) and "boxes" in pred and "keypoints" in pred:
                        boxes = torch.from_numpy(pred["boxes"])
                        kpts = torch.from_numpy(pred["keypoints"])
                        r = Results(batch_imgs[i], batch_paths[i], self.names, boxes=boxes)
                        r.update(keypoints=kpts)
                        results.append(r)
                    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]
                gt_bboxes, gt_cls, gt_kpts = pose_gt_from_yolo_labels(gt, orig_shape, self.num_kpts)
                predn = self._get_results_pred(result)
                stat = self._process_batch(predn, gt_bboxes, gt_cls, gt_kpts)
                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
