from typing import List, Tuple, Union, Any
import time
import numpy as np
# import torch


def _sigmoid(x: np.ndarray) -> np.ndarray:
    """Apply sigmoid activation function with numerical stability."""
    return 1.0 / (1.0 + np.exp(-np.clip(x, -10, 10)))

# def _sigmoid(x: np.ndarray) -> np.ndarray:
#     """Apply sigmoid activation using torch (input/output are numpy)."""
#     xt = torch.from_numpy(x)
#     out = 1.0 / (1.0 + torch.exp(-torch.clamp(xt, min=-10, max=10)))
#     return out.numpy()

def _softmax(x: np.ndarray, axis: int) -> np.ndarray:
    """Apply softmax activation function along specified axis."""
    x = x - np.max(x, axis=axis, keepdims=True)
    ex = np.exp(np.clip(x, -10, 10))
    return ex / np.sum(ex, axis=axis, keepdims=True)

def _transform_feat(feat: List[np.ndarray], na: int, no: int) -> np.ndarray:
    """Transform feature maps to concatenated format."""
    x_feats = []
    B = feat[0].shape[0]
    for i, xi in enumerate(feat):
        *_, H, W = xi.shape
        # (B, na, no, H, W) => (B, na, H, W, no) => (B, -1, no)
        xi = xi.reshape(B, na, no, H, W).transpose(0, 1, 3, 4, 2).reshape(B, -1, no)  # if error occur, maybe is nc problem, check the nc in the config
        

        x_feats.append(xi)  # (B,na*H*W,no)
    return np.concatenate(x_feats, axis=1)  # (B, sum(na*H*W), no)


class DecodeHead:
    """
    Factory Class: Returns corresponding decoder instance based on version and task.
    Uses __new__ method to directly return concrete decoder instance, avoiding an extra layer of proxy object.
    """

    # Use string names to avoid reference issues caused by class definition order
    DECODE_CLASS_MAP = {
        ('yolov5', 'detect'): 'DecodeDetectV5',
        ('yolov5', 'face'): 'DecodeDetectFaceV5',
        ('yolov5', 'pose'): 'DecodePoseV5',
        ('yolov5', 'face'): 'DecodeFaceV5',
        ('yolov5', 'classify'): 'DecodeClassify',
        ('yolov8', 'detect'): 'DecodeDetectV11',
        ('yolov8', 'pose'): 'DecodePoseV11',
        ('yolo11', 'detect'): 'DecodeDetectV11',
        ('yolo11', 'pose'): 'DecodePoseV11',
        ('yolo11', 'classify'): 'DecodeClassify',
        ('yolo26', 'detect'): 'DecodeDetectV26',
        ('yolo26', 'pose'): 'DecodePoseV26',
    }

    def __new__(cls, version: str, task: str, nc: int = 80, **kwargs) -> Any:
        """
        Create and return corresponding decoder instance based on version and task.

        Args:
            version (str): Supported 'yolov5' | 'yolov8' | 'yolo11'
            task (str): Supported 'detect' | 'pose'
            nc (int): Number of classes
            **kwargs: Additional parameters passed to concrete decoder
                     Note: kpt_shape is only used for 'pose' task and will be removed for other tasks

        Returns:
            object: Corresponding decoder instance
        """

        # Create key
        key = (version, task)
        class_name = cls.DECODE_CLASS_MAP.get(key)

        if class_name is None:
            raise ValueError(f"Unsupported version-task combination: {version}-{task}")

        # Filter kwargs: only pose task should receive kpt_shape parameter
        # Remove kpt_shape from kwargs if task is not 'pose' to avoid passing wrong parameters
        filtered_kwargs = kwargs.copy()
        if task != 'pose' and 'kpt_shape' in filtered_kwargs:
            filtered_kwargs.pop('kpt_shape')

        # Get actual class from global namespace to avoid definition order issues
        decoder_class = globals().get(class_name)
        if decoder_class is None:
            raise RuntimeError(f"Decoder class '{class_name}' not found")

        # Return concrete decoder instance
        return decoder_class(nc=nc, **filtered_kwargs)



class DecodeDetectV11:
    def __init__(self, nc: int = 80, reg_max: int = 16):
        self.reg_max = reg_max
        self.no = nc + reg_max*4

        self.strides = np.array([8., 16., 32.])


    def __call__(self, x: List[np.ndarray]) -> np.ndarray:
        """Process feature maps and return decoded bounding boxes and class predictions."""
        B = x[0].shape[0]

        #1. cat aligned feature maps        
        x_cat = np.concatenate([xi.reshape(B, self.no, -1) for xi in x], axis=2)

        #2. make anchors and strides
        self.anchors_grid, self.strides_grid = self._make_anchors(x, self.strides, grid_cell_offset=0.5)

        #3. split x_cat into box and cls
        box, cls = np.split(x_cat, [self.reg_max * 4], axis=1)

        #4. dfl
        box = self.dfl(box)
        cls = _sigmoid(cls)
        # 5. decode bbox
        dbox = self.decode_bboxes(box, self.anchors_grid) * self.strides_grid
        
        return np.concatenate([dbox, cls], axis=1)

    def _make_anchors(self, feat_shapes, strides, grid_cell_offset=0.5):
        """Generate anchors from features."""
        dtype = np.float32
        anchor_points, stride_tensor = [], []
        for i in range(len(feat_shapes)):  # use len(feats) to avoid TracerWarning from iterating over strides tensor
            stride = strides[i]
            h, w = feat_shapes[i].shape[2:]
            sx = np.arange(w, dtype=dtype) + grid_cell_offset  # shift x
            sy = np.arange(h, dtype=dtype) + grid_cell_offset  # shift y
            sy, sx = np.meshgrid(sy, sx, indexing='ij')
            anchor_points.append(np.stack((sx, sy), axis=-1).reshape(-1, 2))
            stride_tensor.append(np.full((h * w, 1), stride, dtype=dtype))
        return np.concatenate(anchor_points).T, np.concatenate(stride_tensor).T

    def dfl(self, feat: np.ndarray) -> np.ndarray:
        """Distribution Focal Loss decoding."""
        b, _, fn = feat.shape  # batch size, feature channel and feature number
        bins = np.arange(self.reg_max).reshape(1, 1, self.reg_max, 1)
        bins = np.broadcast_to(bins, (b, 4, self.reg_max, fn))
        feat = feat.reshape(b, 4, self.reg_max, fn)
        feat = _softmax(feat, axis=2)
        ltrb = (feat * bins).sum(axis=2) 
        return ltrb
    
    def decode_bboxes(self, box: np.ndarray, anchors: np.ndarray) -> np.ndarray:
        """Decode bounding boxes from predictions."""
        return self._dist2bbox(box, anchors, dim=1)

    def _dist2bbox(self, distance: np.ndarray, anchor_points: np.ndarray, dim: int = 1) -> np.ndarray:
        """Transform distance(ltrb) to box(xywh or xyxy)."""
        lt, rb = np.split(distance, 2, axis=dim)
        x1y1 = anchor_points - lt
        x2y2 = anchor_points + rb
        return np.concatenate((x1y1, x2y2), axis=dim)  # xywh bbox



class DecodePoseV11(DecodeDetectV11):
    def __init__(self, nc: int = 1, kpt_shape: Tuple[int, int] = (17, 3)):
        super().__init__(nc=nc)
        self.nc = nc
        self.kpt_shape = kpt_shape
        self.nk = kpt_shape[0] * kpt_shape[1]
    
    def __call__(self, x: Tuple[List[np.ndarray], List[np.ndarray]]) -> np.ndarray:
        """Process bounding box and keypoint feature maps and return decoded results."""
        bbox_heads, kpt_heads = x
        B = bbox_heads[0].shape[0]
        x_cat = np.concatenate([xi.reshape(B, self.nk, -1) for xi in kpt_heads], axis=-1)
        bbox = DecodeDetectV11.__call__(self, bbox_heads)
        pred_kpt = self.decode_kpts(x_cat)
        return np.concatenate((bbox, pred_kpt), axis=1)
    
    def decode_kpts(self, kpts: np.ndarray) -> np.ndarray:
        """Decode keypoints from predictions."""
        ndim = self.kpt_shape[1]
        if ndim == 3:
            kpts[:, 2::ndim] = _sigmoid(kpts[:, 2::ndim])
        kpts[:, 0::ndim] = (kpts[:, 0::ndim] * 2.0 + (self.anchors_grid[0] - 0.5)) * self.strides_grid
        kpts[:, 1::ndim] = (kpts[:, 1::ndim] * 2.0 + (self.anchors_grid[1] - 0.5)) * self.strides_grid
        return kpts



class DecodeSegmentV11(DecodeDetectV11):
    def __init__(self, nc: int = 80) -> None:
        super().__init__(nc=nc)
    
    def __call__(self, x: Tuple[List[np.ndarray], List[np.ndarray], np.ndarray]) -> Tuple[np.ndarray, np.ndarray]:
        """Process segmentation, detection features and mask coefficients, return decoded results."""
        segment_feat, detect_feat, mask_coeffs = x
        x = DecodeDetectV11.__call__(self, detect_feat)
        return (np.concatenate((x[0], mask_coeffs), axis=1), segment_feat)



class DecodeDetectV5:
    anchors = np.array([
        [[10, 13],  [16, 30],  [33, 23]],
        [[30, 61],  [62, 45],  [59, 119]],
        [[116, 90], [156, 198], [373, 326]]
    ], dtype=np.float32)
    strides = np.array([8., 16., 32.])
    
    def __init__(self, nc: int = 80):
        self.na = len(self.anchors[0])
        self.no = nc + 5
        self.anchor_grid, self.strides_grid = None, None

    def __call__(self, x: List[np.ndarray]) -> np.ndarray:
        """Process feature maps and return decoded bounding boxes."""
        #1. cat aligned feature maps
        x_cat = _transform_feat(x, self.na, self.no)
        
        #2. make anchors and strides
        if self.anchor_grid is None:
            self._make_anchorsV5(x, self.na, self.anchors, self.strides)

        # 4. decode bbox
        return self.decode_bboxes(x_cat)

    def decode_bboxes(self, boxes: np.ndarray) -> np.ndarray:
        """Decode bounding boxes from predictions."""
        return self._dist2bbox(boxes)
    
    def _dist2bbox(self, boxes: np.ndarray) -> np.ndarray:
        """Transform xy(ltrb) to box(xywh or xyxy)."""
        xy, wh, conf = np.split(_sigmoid(boxes), [2, 4], axis=2)
        xy = (xy * 2 + self.grid) * self.strides_grid
        wh = (wh * 2) ** 2 * self.anchor_grid
        x1y1 = xy - wh / 2
        x2y2 = xy + wh / 2
        return np.concatenate((x1y1, x2y2, conf), axis=-1)

    def _make_anchorsV5(
        self, 
        feats: List[np.ndarray], 
        na: int, 
        anchors: np.ndarray, 
        strides: Union[List[float], np.ndarray], 
        grid_cell_offset: float = -0.5
    ) -> None:
        """Use NumPy to generate anchors."""
        grids, anchor_grids, strides_grids = [], [], []
        for i, stride in enumerate(strides):
            *_, h, w = feats[i].shape
            # create grid coordinates
            sx = np.arange(w, dtype=np.float32) + grid_cell_offset
            sy = np.arange(h, dtype=np.float32) + grid_cell_offset
            sy, sx = np.meshgrid(sy, sx, indexing="ij")  # shape: (h, w)

            # stack coordinate points
            points = np.stack((sx, sy), axis=-1).reshape(1, 1, -1, 2)  # shape: (h * w, 2)
            points = np.repeat(points, na, axis=1).reshape(1, -1, 2)
            grids.append(points)

            # make anchor grids
            anchor = anchors[i]
            anchor = anchor.reshape(1, na, 1, 1, 2) + np.zeros((1,1,h,w,1), dtype=anchors.dtype)
            anchor = anchor.reshape(1, -1, 2)
            anchor_grids.append(anchor)

            # make stride grids
            strides = np.full((1, na * h * w, 1), stride, dtype=np.float32)
            strides_grids.append(strides)
        self.grid = np.concatenate(grids, axis=1)
        self.anchor_grid = np.concatenate(anchor_grids, axis=1)
        self.strides_grid = np.concatenate(strides_grids, axis=1)
        
        

class DecodePoseV5(DecodeDetectV5):
    def __init__(self, nc: int = 1, kpt_shape: Tuple[int, int] = (17, 3)) -> None:
        self.strides = np.array([8., 16., 32., 64.])
        self.anchors = np.array([
               [[ 19.,  27.],
                [ 44.,  40.],
                [ 38.,  94.]],

               [[ 96.,  68.],
                [ 86., 152.],
                [180., 137.]],

               [[140., 301.],
                [303., 264.],
                [238., 542.]],

               [[436., 615.],
                [739., 380.],
                [925., 792.]]])
        self.nc = 1
        self.no = nc + 5
        self.na = len(self.anchors[0])
        self.kpt_shape = kpt_shape
        self.nk = kpt_shape[0] * kpt_shape[1]    
        self.anchor_grid, self.strides_grid = None, None    

    def __call__(self, x: List[np.ndarray]) -> np.ndarray:
        """Process feature maps and return decoded bounding boxes and keypoints."""
        #feature transform
        feats = _transform_feat(x, self.na, self.nk+self.no)

        #split feats to boxes, kpts
        boxes, kpts = np.split(feats, [6], axis=-1)

        # make grid and anchor
        if self.anchor_grid is None:
            self._make_anchorsV5(x, self.na, self.anchors, self.strides)
        
        #process boxes
        pred_boxes = self.decode_bboxes(boxes)

        #process kpts
        pred_kpt = self.decode_kpts(kpts)
        
        #concat outputs
        return np.concatenate([pred_boxes, pred_kpt], axis=-1)
    
    def decode_kpts(self, kpts: np.ndarray) -> np.ndarray:
        """Decode keypoints from predictions."""
        return self._dist2kpts(kpts)

    def _dist2kpts(self, kpts: np.ndarray) -> np.ndarray:
        """Transform keypoint distances to coordinates."""
        x = (kpts[..., 0::3] * 2.0 + self.grid[..., 0:1]) * self.strides_grid
        y = (kpts[..., 1::3] * 2.0 + self.grid[..., 1:2]) * self.strides_grid
        conf = _sigmoid(kpts[..., 2::3])
        return np.stack((x, y, conf), axis=-1).reshape(kpts.shape) 



# class DecodeFaceV5:
#     def __init__(self, nc=80):
#         self.nc = nc  # number of classes
#         self.no = nc + 5 + 10  # number of outputs per anchor
#         self.nl = 3  # number of detection layers
#         self.na = 3  # number of anchors
#         self.anchors = np.array([
#             [4, 5, 8, 10, 13, 16],      # P3/8
#             [23, 29, 43, 55, 73, 105],  # P4/16
#             [146, 217, 231, 300, 335, 433] # P5/32
#         ]).reshape(3, 3, 2)
#         self.strides = np.array([8., 16., 32.])
        
#         # --- Initialize grid ---
#         self.grid = [np.zeros(1)] * self.nl

#     def __call__(self, x):
#         # x is a list of 9 numpy arrays: [head1_1, head1_2, head1_3, ...]
#         z = []
#         for i in range(self.nl):
#             # _1~_3 are p1 (sigmoid xywhc), p2 (raw landm), p3 (sigmoid cls)
#             p1 = x[i*3]     
#             p2 = x[i*3 + 1] 
#             p3 = x[i*3 + 2] 
#             bs, _, ny, nx = p1.shape  
#             p1 = p1.reshape(bs, self.na, -1, ny, nx).transpose(0, 1, 3, 4, 2)
#             p2 = p2.reshape(bs, self.na, -1, ny, nx).transpose(0, 1, 3, 4, 2)
#             p3 = p3.reshape(bs, self.na, -1, ny, nx).transpose(0, 1, 3, 4, 2)
#             # p1 = p1.reshape(bs, -1, self.na, ny, nx).transpose(0, 1, 3, 4, 2)
#             # p2 = p2.reshape(bs, -1, self.na, ny, nx).transpose(0, 1, 3, 4, 2)
#             # p3 = p3.reshape(bs, -1, self.na, ny, nx).transpose(0, 1, 3, 4, 2)



#             y = np.concatenate([p1, p2, p3], axis=-1)
#             # print(y.shape, 'y.shape')
#             bs, na, ny, nx, _ = y.shape
#             # y = y.reshape(bs, self.na, self.no, ny, nx).transpose(0, 1, 3, 4, 2)

#             # --- Grid cache logic ---
#             if self.grid[i].shape[2:4] != (ny, nx):
#                 yv, xv = np.meshgrid(np.arange(ny), np.arange(nx), indexing='ij')
#                 self.grid[i] = np.stack((xv, yv), axis=-1).reshape(1, 1, ny, nx, 2)

#             # Anchor grid (already in pixel space, no need to multiply by stride)
#             anchor_grid = self.anchors[i].reshape(1, na, 1, 1, 2)

#             # Post-processing
#             # box_xy uses self.grid[i]
#             box_xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.strides[i]
#             box_wh = (y[..., 2:4] * 2) ** 2 * anchor_grid
#             conf = y[..., 4:5]

#             # Landmarks use self.grid[i]
#             # Reshape y[..., 5:15] from (..., 10) to (..., 5, 2) for broadcasting with anchor_grid of shape (..., 2)
#             lms = y[..., 5:15].reshape(bs, na, ny, nx, 5, 2) * anchor_grid[:, :, :, :, np.newaxis, :] + \
#                   self.grid[i][:, :, :, :, np.newaxis, :] * self.strides[i]
#             landmarks = lms.reshape(bs, na, ny, nx, 10)

#             # Classes
#             cls = y[..., 15:15+self.nc]

#             # Concatenate results for this scale
#             # Reorder to match the keypoints function signature: [box(4), conf(1), cls(1), landmarks(10)]
#             print(box_xy.shape, box_wh.shape, conf.shape, cls.shape, landmarks.shape, '==')
#             y_out = np.concatenate([box_xy, box_wh, conf, cls, landmarks], axis=-1)

#             z.append(y_out.reshape(bs, -1, self.no))
#         # np.save('decode_output.npy', np.concatenate(z, axis=1))
#         return np.load('./decode_output_from_torch.npy')
        
#         # return np.concatenate(z, axis=1)



class DecodeSegmentV5(DecodeDetectV5):
    def __init__(self, nc: int = 80) -> None:
        super().__init__(nc=nc)

    def __call__(self, x: Tuple[List[np.ndarray], List[np.ndarray]]) -> Tuple[np.ndarray, np.ndarray]:
        """Process segmentation and detection features and return decoded results."""
        segment_feat, detect_feat = x
        x = DecodeDetectV5.__call__(self, detect_feat)
        return (x[0], segment_feat)
        


class DecodeClassify:
    def __init__(self, nc: int = 80) -> None:
        pass
    
    def __call__(self, x: np.ndarray) -> np.ndarray:
        """Process classification predictions and return probabilities."""
        return x[0]



class DecodeDetectV26:
    def __init__(self, nc: int = 80, max_det: int = 300) -> None:
        self.nc = nc
        self.max_det = max_det
        self.strides = np.array([8., 16., 32.])
        self.anchors_grid, self.strides_grid = None, None

    def __call__(self, x: List[np.ndarray]) -> np.ndarray:
        """
        Process feature maps and return decoded bounding boxes.
        Args:
            x (List[np.ndarray]): List of feature maps.
            bbox: x[:3]
            cls : x[3:]
        Returns:
            np.ndarray: Decoded bounding boxes.
        """
        s1 = time.time()
        if self.anchors_grid is None:
            self._make_anchors(x[:3], self.strides, grid_cell_offset=0.5)
        feats_dict = self.feats_process(x)
        x_cat = self._inference(feats_dict)
        results = self._postprocess(np.transpose(x_cat, (0, 2, 1)), self.nc, self.max_det)
        return results

    def _inference(self, x: np.ndarray) -> np.ndarray:
        dbox = self._dist2bbox(x["boxes"], self.anchors_grid[np.newaxis, :]) * self.strides_grid
        sigmoid_scores = _sigmoid(x["scores"])
        return np.concatenate([dbox, sigmoid_scores], axis=1)

    def _postprocess(self, preds: np.ndarray, nc: int, max_det: int) -> np.ndarray:
        """Post-processes YOLO model predictions.

        Args:
            preds (np.ndarray): Raw predictions with shape (batch_size, num_anchors, 4 + nc) with last dimension
                format [x1, y1, x2, y2, class_probs] when end2end=True, or [x, y, w, h, class_probs] otherwise.

        Returns:
            (np.ndarray): Processed predictions with shape (batch_size, min(max_det, num_anchors), 6) and last
                dimension format [x1, y1, x2, y2, max_class_prob, class_index] when end2end=True,
                or [x, y, w, h, max_class_prob, class_index] otherwise.
        """
        boxes = preds[..., :4]  # (batch_size, num_anchors, 4)
        scores = preds[..., 4:]  # (batch_size, num_anchors, nc)
        scores, conf, idx = self.get_topk_index(scores, max_det)
        
        # Gather boxes: for each batch, select boxes at indices in idx
        # batch_size = boxes.shape[0]
        # gathered_boxes = np.array([boxes[i, idx[i]] for i in range(batch_size)])  # (batch_size, k, 4)
        batch_idx = np.arange(boxes.shape[0])[:, None]
        gathered_boxes = boxes[batch_idx, idx]   # (batch_size, k, 4)        
        return np.concatenate([gathered_boxes, scores, conf], axis=-1)

    def feats_process(self, x: List[np.ndarray]):
        """Process feature maps."""
        B = x[0].shape[0]
        boxes_feat = np.concatenate([arr.reshape(B, 4, -1) for arr in x[:3]], axis=-1)
        cls_feat = np.concatenate([arr.reshape(B, self.nc, -1) for arr in x[3:6]], axis=-1)
        return dict(boxes=boxes_feat, scores=cls_feat)

    def get_topk_index(self, scores: np.ndarray, max_det: int):
        """Get top-k indices from scores.

        Args:
            scores (np.ndarray): Scores array with shape (batch_size, num_anchors, num_classes).
            max_det (int): Maximum detections per image.

        Returns:
            (np.ndarray, np.ndarray, np.ndarray): Top scores, class indices, and filtered indices.
        """
        batch_size, anchors, nc = scores.shape  # i.e. shape(16,8400,84)
        # Use max_det directly during export for TensorRT compatibility (requires k to be constant),
        # otherwise use min(max_det, anchors) for safety with small inputs during Python inference
        k = min(max_det, anchors)
        
        # Get max scores per anchor and find top-k indices
        max_scores = np.max(scores, axis=-1)  # (batch_size, num_anchors)
        ori_index = np.argsort(max_scores, axis=-1)[:, -k:][:, ::-1]  # top-k indices, descending (batch_size, k)
        ori_index_expanded = np.expand_dims(ori_index, axis=-1)  # (batch_size, k, 1)
        
        # Gather scores using advanced indexing
        batch_indices = np.arange(batch_size)[:, None, None]  # (batch_size, 1, 1)
        gathered_scores = scores[batch_indices, ori_index_expanded, :]  # (batch_size, k, nc)
        
        # Flatten and get top-k
        gathered_scores_flat = gathered_scores.reshape(batch_size, -1)  # (batch_size, k*nc)
        topk_indices_flat = np.argsort(gathered_scores_flat, axis=-1)[:, -k:][:, ::-1]  # top-k, descending
        topk_scores_flat = np.take_along_axis(gathered_scores_flat, topk_indices_flat, axis=-1)
        
        # Calculate class indices and original anchor indices
        class_indices = (topk_indices_flat % nc)[:, :, None].astype(np.float32)  # (batch_size, k, 1)
        anchor_indices = topk_indices_flat // nc  # (batch_size, k)
        
        # Get original anchor indices - use advanced indexing
        batch_idx = np.arange(batch_size)[:, None]  # (batch_size, 1)
        idx = ori_index[batch_idx, anchor_indices]  # (batch_size, k)
        
        return topk_scores_flat[:, :, None], class_indices, idx

    def _make_anchors(self, feat_shapes, strides, grid_cell_offset=0.5):
        """Generate anchors from features."""
        dtype = np.float32
        anchor_points, stride_tensor = [], []
        for i in range(len(feat_shapes)):  # use len(feats) to avoid TracerWarning from iterating over strides tensor
            stride = strides[i]
            h, w = feat_shapes[i].shape[2:]
            sx = np.arange(w, dtype=dtype) + grid_cell_offset  # shift x
            sy = np.arange(h, dtype=dtype) + grid_cell_offset  # shift y
            sy, sx = np.meshgrid(sy, sx, indexing='ij')
            anchor_points.append(np.stack((sx, sy), axis=-1).reshape(-1, 2))
            stride_tensor.append(np.full((h * w, 1), stride, dtype=dtype))
        self.anchors_grid = np.concatenate(anchor_points).T
        self.strides_grid = np.concatenate(stride_tensor).T

    def _dist2bbox(self, distance: np.ndarray, anchor_points: np.ndarray, dim: int = 1) -> np.ndarray:
        """Transform distance(ltrb) to box(xywh or xyxy)."""
        lt, rb = np.split(distance, 2, axis=dim)
        x1y1 = anchor_points - lt
        x2y2 = anchor_points + rb
        return np.concatenate((x1y1, x2y2), axis=dim)  # xywh bbox



class DecodePoseV26(DecodeDetectV26):
    def __init__(self, nc: int = 1, kpt_shape: Tuple[int, int] = (17, 3), max_det: int = 300) -> None:
        super().__init__(nc=nc, max_det=max_det)
        self.kpt_shape = kpt_shape
        self.nk = kpt_shape[0] * kpt_shape[1]
        self.nc = nc
        self.max_det = max_det
        # self.anchors_grid, self.strides_grid = None, None

    def __call__(self, x: List[np.ndarray]) -> np.ndarray:
        """Process feature maps and return decoded bounding boxes and keypoints."""
        if self.anchors_grid is None:
            self._make_anchors(x[:3], self.strides, grid_cell_offset=0.5)
        feats_dict = self._feats_process(x)
        boxes_results = super()._inference(feats_dict)
        kpts_results = self.decode_kpts(feats_dict["kpts"])
        cat_results = np.concatenate((boxes_results, kpts_results), axis=1)
        return self._postprocess(cat_results.transpose(0, 2, 1), self.nc, nk=self.nk, max_det=self.max_det)

    def _postprocess(self, preds: np.ndarray, nc: int, nk: int, max_det: int) -> np.ndarray:
        """Post-process YOLO model predictions.

        Args:
            preds (np.ndarray): Raw predictions with shape (batch_size, num_anchors, 4 + nc + nk) with last dimension
                format [x, y, w, h, class_probs, keypoints].
            nc (int): Number of classes.
            nk (int): Number of keypoint values.
            max_det (int): Maximum detections per image.

        Returns:
            (np.ndarray): Processed predictions with shape (batch_size, min(max_det, num_anchors), 6 + nk) and
                last dimension format [x, y, w, h, max_class_prob, class_index, keypoints].
        """
        # Split predictions into boxes, scores, and keypoints
        boxes, scores, kpts = np.split(preds, [4, 4 + self.nc], axis=-1)
        
        # Get top-k indices
        scores, conf, idx = self.get_topk_index(scores, self.max_det)
        
        # Gather boxes and keypoints using advanced indexing
        batch_size = boxes.shape[0]
        gathered_boxes = np.array([boxes[i, idx[i]] for i in range(batch_size)])  # (batch_size, k, 4)
        gathered_kpts = np.array([kpts[i, idx[i]] for i in range(batch_size)])  # (batch_size, k, nk)
        
        # Concatenate results
        return np.concatenate([gathered_boxes, scores, conf, gathered_kpts], axis=-1)

    def decode_kpts(self, kpts: np.ndarray) -> np.ndarray:
        ndim = self.kpt_shape[1]
        if ndim == 3:
            kpts[:, 2::ndim] = _sigmoid(kpts[:, 2::ndim])
        kpts[:, 0::ndim] = (kpts[:, 0::ndim] + self.anchors_grid[0]) * self.strides_grid
        kpts[:, 1::ndim] = (kpts[:, 1::ndim] + self.anchors_grid[1]) * self.strides_grid
        return kpts

    def _feats_process(self, x: List[np.ndarray]):
        B = x[0].shape[0]
        boxes_feat = np.concatenate([arr.reshape(B, 4, -1) for arr in x[:3]], axis=-1)
        cls_feat = np.concatenate([arr.reshape(B, self.nc, -1) for arr in x[3:6]], axis=-1)
        kpt_feat = np.concatenate([arr.reshape(B, self.nk, -1) for arr in x[6:]], axis=-1)
        return dict(boxes=boxes_feat, scores=cls_feat, kpts=kpt_feat)