"""YOLO-seg postprocessor (YOLO26-seg instance segmentation).

Input: raw predictions dict with two outputs:
  - output0: [1, 300, 38] — 300 NMS-filtered detections, each [x1,y1,x2,y2,conf,class_id,coeff_0..coeff_31]
  - output1: [1, 32, 160, 160] — 32 prototype masks at 160x160

For raw (no NMS) models:
  - output0: [1, 4+numClasses+32, numBoxes] — standard YOLO format with mask coefficients appended

Output: {"kind": "detections", "detections": [{..., "mask": "<b64>", "maskWidth": N, "maskHeight": N}]}
"""
import base64

import numpy as np


COCO_80 = [
    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck",
    "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench",
    "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra",
    "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
    "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove",
    "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup",
    "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange",
    "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
    "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse",
    "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink",
    "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier",
    "toothbrush",
]

NUM_MASK_COEFFS = 32
PROTO_SIZE = 160


def _iou(a: tuple, b: tuple) -> float:
    """IoU between two boxes (x1, y1, x2, y2)."""
    ax1, ay1, ax2, ay2 = a
    bx1, by1, bx2, by2 = b
    ix1 = max(ax1, bx1)
    iy1 = max(ay1, by1)
    ix2 = min(ax2, bx2)
    iy2 = min(ay2, by2)
    inter = max(0, ix2 - ix1) * max(0, iy2 - iy1)
    if inter == 0:
        return 0.0
    area_a = (ax2 - ax1) * (ay2 - ay1)
    area_b = (bx2 - bx1) * (by2 - by1)
    union = area_a + area_b - inter
    return inter / union if union > 0 else 0.0


def _nms(boxes: list[dict], iou_threshold: float = 0.45) -> list[dict]:
    """Non-maximum suppression. Input: list of dicts with _xyxy, score, class."""
    if not boxes:
        return []
    sorted_boxes = sorted(boxes, key=lambda b: b["score"], reverse=True)
    kept = []
    suppressed = set()
    for i, box in enumerate(sorted_boxes):
        if i in suppressed:
            continue
        kept.append(box)
        for j in range(i + 1, len(sorted_boxes)):
            if j in suppressed:
                continue
            if _iou(box["_xyxy"], sorted_boxes[j]["_xyxy"]) > iou_threshold:
                suppressed.add(j)
    return kept


def _sigmoid(x: np.ndarray) -> np.ndarray:
    """Numerically stable sigmoid."""
    return np.where(x >= 0, 1 / (1 + np.exp(-x)), np.exp(x) / (1 + np.exp(x)))


def _find_tensors(predictions: dict, num_classes: int) -> tuple[np.ndarray, np.ndarray]:
    """Locate the detection tensor and prototype tensor from model outputs.

    Returns (detections, prototypes) as numpy arrays.
    detections: [num_dets, channels] row-major
    prototypes: [32, 160, 160]
    """
    det_tensor = None
    proto_tensor = None

    for key, val in predictions.items():
        arr = np.array(val, dtype=np.float32)

        # Prototype: shape containing 32 * 160 * 160 = 819200
        if arr.size == NUM_MASK_COEFFS * PROTO_SIZE * PROTO_SIZE:
            proto_tensor = arr.reshape(NUM_MASK_COEFFS, PROTO_SIZE, PROTO_SIZE)
            continue

        # Skip tiny tensors
        if arr.size < 10:
            continue

        # Detection tensor: the other large tensor
        if det_tensor is None:
            det_tensor = arr

    if det_tensor is None:
        raise ValueError("YOLO-seg postprocessor: could not find detection tensor")
    if proto_tensor is None:
        raise ValueError("YOLO-seg postprocessor: could not find prototype tensor")

    # Normalize detection tensor to [num_dets, channels]
    while det_tensor.ndim > 2:
        det_tensor = det_tensor[0]  # strip batch dim

    # NMS format: [300, 38] — rows are detections
    # Raw format: [4+nc+32, N] — cols are detections, needs transpose
    nms_cols = 6 + NUM_MASK_COEFFS  # 38
    raw_cols = 4 + num_classes + NUM_MASK_COEFFS

    if det_tensor.shape[1] == nms_cols:
        pass  # already [N, 38]
    elif det_tensor.shape[0] == nms_cols:
        det_tensor = det_tensor.T
    elif det_tensor.shape[1] == raw_cols:
        pass  # already [N, raw_cols]
    elif det_tensor.shape[0] == raw_cols:
        det_tensor = det_tensor.T
    elif det_tensor.shape[0] < det_tensor.shape[1]:
        det_tensor = det_tensor.T  # likely [C, N] -> [N, C]

    return det_tensor, proto_tensor


def _crop_mask_to_bbox(
    mask_full: np.ndarray,
    x1: float,
    y1: float,
    x2: float,
    y2: float,
    input_size: int,
    mask_threshold: float = 0.5,
) -> tuple[np.ndarray, int, int]:
    """Crop sigmoid mask to bbox region, threshold to binary, return (binary, w, h)."""
    proto_scale = PROTO_SIZE / input_size
    px1 = max(0, int(np.floor(x1 * proto_scale)))
    py1 = max(0, int(np.floor(y1 * proto_scale)))
    px2 = min(PROTO_SIZE, int(np.ceil(x2 * proto_scale)))
    py2 = min(PROTO_SIZE, int(np.ceil(y2 * proto_scale)))

    crop_w = max(1, px2 - px1)
    crop_h = max(1, py2 - py1)

    cropped = mask_full[py1:py1 + crop_h, px1:px1 + crop_w]
    binary = (cropped > mask_threshold).astype(np.uint8) * 255
    return binary, crop_w, crop_h


def postprocess_yolo_seg(
    predictions: dict,
    config: dict,
    orig_w: int,
    orig_h: int,
    scale: float,
    pad: tuple[int, int],
) -> dict:
    """Postprocess YOLO-seg output to detections with per-instance masks.

    Handles:
    1. NMS-enabled export: [N, 6+32] with [x1,y1,x2,y2,conf,class_id,...coeffs]
    2. Raw export: [4+numClasses+32, N] standard YOLO transposed format
    """
    conf_threshold = float(config.get("confidence", 0))
    labels = config.get("labels", COCO_80)
    num_classes = config.get("numClasses", len(labels))
    input_size = config.get("inputSize", 640)
    mask_threshold = config.get("maskThreshold", 0.5)

    det_tensor, proto_tensor = _find_tensors(predictions, num_classes)

    num_dets = det_tensor.shape[0]
    num_cols = det_tensor.shape[1]
    is_nms_format = num_cols == 6 + NUM_MASK_COEFFS

    candidates = []

    for i in range(num_dets):
        row = det_tensor[i]

        if is_nms_format:
            # [x1, y1, x2, y2, confidence, class_id, coeff_0..coeff_31]
            x1, y1, x2, y2 = float(row[0]), float(row[1]), float(row[2]), float(row[3])
            best_score = float(row[4])
            best_class = int(row[5])
            coeffs = row[6:6 + NUM_MASK_COEFFS]

            if best_score < conf_threshold:
                continue
            if best_class < 0:
                continue  # padding row
        else:
            # [cx, cy, w, h, class_scores..., coeff_0..coeff_31]
            cx, cy, w, h = float(row[0]), float(row[1]), float(row[2]), float(row[3])
            class_scores = row[4:4 + num_classes]
            best_class = int(np.argmax(class_scores))
            best_score = float(class_scores[best_class])
            coeffs = row[4 + num_classes:4 + num_classes + NUM_MASK_COEFFS]

            if best_score < conf_threshold:
                continue

            x1 = cx - w / 2
            y1 = cy - h / 2
            x2 = cx + w / 2
            y2 = cy + h / 2

        # Compute instance mask: sigmoid(coefficients @ prototypes)
        # coeffs: [32], proto_tensor: [32, 160, 160]
        mask_raw = _sigmoid(coeffs @ proto_tensor.reshape(NUM_MASK_COEFFS, -1)).reshape(
            PROTO_SIZE, PROTO_SIZE
        )

        # Crop mask to bbox, threshold to binary
        binary, crop_w, crop_h = _crop_mask_to_bbox(
            mask_raw, x1, y1, x2, y2, input_size, mask_threshold
        )
        mask_b64 = base64.b64encode(binary.tobytes()).decode("ascii")

        # Transform bbox from model coords to original image coords
        ox1 = max(0, min(orig_w, (x1 - pad[0]) / scale))
        oy1 = max(0, min(orig_h, (y1 - pad[1]) / scale))
        ox2 = max(0, min(orig_w, (x2 - pad[0]) / scale))
        oy2 = max(0, min(orig_h, (y2 - pad[1]) / scale))

        label = labels[best_class] if best_class < len(labels) else str(best_class)
        candidates.append({
            "class": label,
            "score": round(best_score, 4),
            "bbox": [round(ox1, 1), round(oy1, 1), round(ox2, 1), round(oy2, 1)],
            "mask": mask_b64,
            "maskWidth": crop_w,
            "maskHeight": crop_h,
            "_xyxy": (ox1, oy1, ox2, oy2),
        })

    kept = _nms(candidates)
    # Remove internal _xyxy field
    for d in kept:
        del d["_xyxy"]

    return {"kind": "detections", "detections": kept}