"""
Core postprocessing functions.
"""
from typing import List, Dict, Any, Tuple, Union
import numpy as np

from .ops import non_max_suppression, scale_boxes, scale_coords, process_mask


# =============================================================================
# Allowed steps per task
# =============================================================================
TASK_ALLOWED_STEPS = {
    'detect'  : ['unpack', 'nms', 'scale'],
    'segment' : ['engine_swap', 'proto', 'nms', 'mask'],
    'pose'    : ['nms', 'keypoints'],
    'classify': ['classify']
}

# Default steps per task
TASK_DEFAULT_STEPS = {
    'detect'  : ['nms', 'scale'],
    'segment' : ['engine_swap', 'proto', 'nms', 'mask'],
    'pose'    : ['nms', 'keypoints'],
    'classify': ['classify']
}


# =============================================================================
# Core processing functions
# =============================================================================

def nms(
    pred: np.ndarray,
    conf_threshold: float = 0.25,
    iou_threshold: float = 0.45,
    nc: int = 80,
    classes: List[int] = None,
    agnostic: bool = False,
    version: str = 'yolo11',
    task: str = None
) -> List[np.ndarray]:
    """Run NMS on raw model predictions."""
    return non_max_suppression(
        pred,
        conf_threshold=conf_threshold,
        iou_threshold=iou_threshold,
        nc=nc,
        classes=classes,
        agnostic=agnostic,
        version=version,
        task=task
    )


def unpack(
    preds: Union[List[np.ndarray], Tuple],
    input_shape: Tuple[int, int]
) -> List[np.ndarray]:
    """
    Unpack the output format of .hef models with built-in decode+nms.

    Applies when hef_has_nms=True; the output is in packed (B, 1000) format.

    Packed buffer layout (per image, up to 1000 floats total):
        For each class (repeated until buffer is exhausted):
            [count]  1 float  — number of detections for this class (cast to int)
            [objs]   count × 5 floats — one object per row: [y1, x1, y2, x2, conf]
        Classes with count=0 occupy exactly 1 float (the count byte itself).

    Args:
        preds: Model predictions where preds[0] is packed data of shape (B, 1000).
        input_shape: Input image size as (height, width).

    Returns:
        List[np.ndarray]: Per-batch detection results of shape (N, 6) [x1, y1, x2, y2, conf, cls].
    """
    # Get prediction data, typically of shape (B, 1000)
    data = preds[0]
    if data.ndim == 1:
        data = data[None, :]

    # Only process the first 1000 bytes
    data = data[:, :1000]
    
    batch_size = data.shape[0]
    img_width = input_shape[1]
    results = []
    
    for b in range(batch_size):
        batch_data = data[b]
        curr_idx = 0
        batch_bboxes = []
        class_id = 0
        
        # Iterate through the buffer sequentially by class.
        # Coordinate values (float32) may be non-zero, so np.where(data != 0) cannot
        # be used to locate count bytes directly.
        while curr_idx < 1000:
            # Read the detection count for the current class
            count = int(batch_data[curr_idx])

            if count > 0:
                # Check that the remaining buffer is large enough to read coordinates
                # (5 values per object: y1, x1, y2, x2, conf)
                if curr_idx + 1 + count * 5 > 1000:
                    break

                # Extract coordinates and confidence [y1, x1, y2, x2, conf]
                objs = batch_data[curr_idx + 1 : curr_idx + 1 + count * 5].reshape(count, 5)

                # Build result array [x1, y1, x2, y2, conf, cls]
                res = np.empty((count, 6), dtype=np.float32)
                res[:, :5] = objs
                res[:, 5] = class_id

                # Swap and scale coordinates (y1, x1, y2, x2) -> (x1, y1, x2, y2)
                # Equivalent to res[:, [1, 0, 3, 2]] in the original code
                res[:, :4] = res[:, [1, 0, 3, 2]] * img_width

                batch_bboxes.append(res)
                curr_idx += 1 + count * 5
            else:
                # Count is 0; skip the count byte for this class
                curr_idx += 1
            
            class_id += 1
            
        if batch_bboxes:
            results.append(np.vstack(batch_bboxes))
        else:
            results.append(np.empty((0, 6), dtype=np.float32))
            
    return results


def scale(
    detections: List[np.ndarray],
    input_shape: Tuple[int, int],
    orig_imgs_shapes: List[Tuple[int, int]]
) -> List[np.ndarray]:
    """Scale bboxes from input_shape to the original image size."""
    results = []
    for i, det in enumerate(detections):
        if len(det):
            det[:, :4] = scale_boxes(input_shape, det[:, :4], orig_imgs_shapes[i]).round()
            det = det[np.argsort(det[:, 0])]
        else:
            det = np.zeros((0, 6), dtype=np.float32)
        results.append(det)
    return results


def mask(
    detections: List[np.ndarray],
    proto: np.ndarray,
    input_shape: Tuple[int, int],
    orig_imgs_shapes: List[Tuple[int, int]],
    version: str = 'yolo11'
) -> List[Dict[str, Any]]:
    """Process segmentation masks."""
    results = []
    for i, pred in enumerate(detections):
        if len(pred) == 0:
            results.append({'boxes': np.empty((0, 6)), 'mask': None})
            continue
        masks = process_mask(proto[i], pred[:, 6:], pred[:, :4], input_shape, upsample=True)
        pred[:, :4] = scale_boxes(input_shape, pred[:, :4], orig_imgs_shapes[i])
        results.append({'boxes': pred[:, :6], 'mask': masks})
    return results


def keypoints(
    detections: List[np.ndarray],
    input_shape: Tuple[int, int],
    orig_imgs_shapes: List[Tuple[int, int]],
    kpt_shape: Tuple[int, int] = (17, 3)
) -> List[Dict[str, Any]]:
    """Process pose keypoints."""
    results = []
    for i, pred in enumerate(detections):
        if len(pred) == 0:
            results.append({'boxes': np.empty((0, 6)), 'keypoints': np.empty((0, *kpt_shape))})
            continue
        pred[:, :4] = scale_boxes(input_shape, pred[:, :4], orig_imgs_shapes[i]).round()
        pred_kpts = pred[:, 6:].reshape((len(pred), *kpt_shape))
        pred_kpts = scale_coords(input_shape, pred_kpts, orig_imgs_shapes[i], kpt_shape=kpt_shape)
        results.append({'boxes': pred[:, :6], 'keypoints': pred_kpts})
    return results
    

def engine_swap(preds, weights: str):
    """Reorder outputs for .engine format.

    Args:
        preds: Model predictions.
        weights: Model weights path, used to detect .engine format.

    Returns:
        preds: Output list with first two elements swapped if .engine, otherwise unchanged.
    """
    if weights.endswith('.engine'):
        preds = list(preds)
        preds[0], preds[1] = preds[1], preds[0]
    return preds


def proto(preds):
    """Extract the segmentation proto tensor from model predictions.

    Args:
        preds: Model predictions.

    Returns:
        proto_data: Proto tensor with a batch dimension ensured.
    """
    return preds[1][None, ...] if preds[1].ndim == 3 else preds[1]


def classify(preds):
    """Postprocess classification predictions.

    Args:
        preds: Model predictions.

    Returns:
        Tuple of raw predictions and (top5_scores, top5_indices).
    """
    preds = preds[0] if isinstance(preds, (list, tuple)) else preds
    top5_idx = np.argsort(preds, axis=1)[:, -5:][:, ::-1]
    top5 = np.take_along_axis(preds, top5_idx, axis=1)
    return preds, (top5, top5_idx)
