# -*- coding: utf-8 -*-
from __future__ import annotations
import math
from typing import Dict, List, Tuple, Any

import numpy as np

def sigmoid(x: np.ndarray) -> np.ndarray:
    return 1.0 / (1.0 + np.exp(-x))

def temp_scale(p: np.ndarray, temperature: float) -> np.ndarray:
    p = np.clip(p, 1e-8, 1 - 1e-8)
    z = np.log(p) - np.log(1 - p)
    return 1.0 / (1.0 + np.exp(-z / max(temperature, 1e-6)))

def xywh2xyxy(xywh: np.ndarray) -> np.ndarray:
    x, y, w, h = xywh.T
    return np.stack([x - w/2, y - h/2, x + w/2, y + h/2], axis=1)

def nms(boxes: np.ndarray, scores: np.ndarray, iou_thr: float, topk: int) -> List[int]:
    # simple NMS
    idxs = scores.argsort()[::-1]
    keep = []
    while idxs.size > 0 and len(keep) < topk:
        i = idxs[0]
        keep.append(i)
        if idxs.size == 1:
            break
        ious = iou(boxes[i], boxes[idxs[1:]])
        idxs = idxs[1:][ious < iou_thr]
    return keep

def iou(a: np.ndarray, bs: np.ndarray) -> np.ndarray:
    ax1, ay1, ax2, ay2 = a
    bx1, by1, bx2, by2 = bs.T
    ix1, iy1 = np.maximum(ax1, bx1), np.maximum(ay1, by1)
    ix2, iy2 = np.minimum(ax2, bx2), np.minimum(ay2, by2)
    iw, ih = np.maximum(0.0, ix2 - ix1), np.maximum(0.0, iy2 - iy1)
    inter = iw * ih
    area_a = (ax2 - ax1) * (ay2 - ay1)
    area_b = (bx2 - bx1) * (by2 - by1)
    union = area_a + area_b - inter + 1e-9
    return inter / union

def scale_coords(xyxy: np.ndarray, img_meta: Tuple[int, int], imgsz: int) -> np.ndarray:
    # img_meta should be (w0, h0). Falls vertauscht geliefert, korrigieren wir unten adaptiv.
    def _scale(xyxy_in: np.ndarray, w0: int, h0: int) -> np.ndarray:
        r = min(imgsz / h0, imgsz / w0)
        nw, nh = int(round(w0 * r)), int(round(h0 * r))
        pad_w, pad_h = (imgsz - nw) // 2, (imgsz - nh) // 2
        x1, y1, x2, y2 = xyxy_in.T
        x1 = (x1 - pad_w) / r
        y1 = (y1 - pad_h) / r
        x2 = (x2 - pad_w) / r
        y2 = (y2 - pad_h) / r
        out = np.stack([x1, y1, x2, y2], axis=1)
        out[:, 0] = np.clip(out[:, 0], 0, w0 - 1)
        out[:, 1] = np.clip(out[:, 1], 0, h0 - 1)
        out[:, 2] = np.clip(out[:, 2], 0, w0 - 1)
        out[:, 3] = np.clip(out[:, 3], 0, h0 - 1)
        return out

    w0, h0 = img_meta
    out = _scale(xyxy, int(w0), int(h0))

    # Heuristik: wenn Breite ~0 aber Höhe deutlich > 0 → probiere vertauschte Reihenfolge (h0,w0)
    w_pix = out[:, 2] - out[:, 0]
    h_pix = out[:, 3] - out[:, 1]
    if out.size and np.median(w_pix) < 2.0 and np.median(h_pix) > 5.0:
        out_alt = _scale(xyxy, int(h0), int(w0))
        w_pix_alt = out_alt[:, 2] - out_alt[:, 0]
        if np.median(w_pix_alt) > np.median(w_pix):
            out = out_alt

    return out

def postprocess_yolov8_seg(outputs: Dict[str, np.ndarray],
                           imgsz: int,
                           img_meta: Tuple[int, int],
                           names: List[str],
                           conf_thr: float,
                           iou_thr: float,
                           max_det: int,
                           calibration: Dict[str, Any] | None) -> Dict[str, Any]:
    """
    Minimal post-processor for YOLOv8-Seg ONNX export.
    Expects Ultralytics-like outputs:
        out0: [1, N, 84+?] (boxes xywh, objectness, class confs, mask coeffs)
        out1: proto: [1, C, H/4, W/4] (for masks)
    Layout can differ across versions/exports; we handle the common path:
    - boxes xywh in pixel coords on imgsz, then scaled back to original
    - score = obj * max(class_prob)
    - optional temperature calibration
    """
    # Heuristic: find 'proto' (largest 4D tensor) and 'pred' (2D/3D)
    outs = list(outputs.values())
    outs.sort(key=lambda a: len(a.shape))
    pred = outs[-2]  # [1,N,K]
    proto = outs[-1] # [1,C,H,W] – masks ignored in v1 (optional later)

    pred = np.squeeze(pred, axis=0)  # [N,K]
    num_classes = len(names)

    # boxes
    b = pred[:, :4]  # expected xywh
    # handle normalized boxes in [0..1]
    if np.nanmax(b) <= 1.5:
        b = b * float(imgsz)

    # scores
    obj = pred[:, 4:5]
    cls = pred[:, 5:5 + num_classes]
    cls = np.clip(cls, 0.0, 1.0)
    obj = np.clip(obj, 0.0, 1.0)

    cls_idx = cls.argmax(axis=1)
    cls_prob = cls.max(axis=1, keepdims=True)
    score = (obj * cls_prob).flatten()

    # calibration
    if calibration and "temperature" in calibration:
        score = temp_scale(score, float(calibration["temperature"]))

    # threshold
    m = score >= conf_thr
    if not np.any(m):
        return {"boxes": [], "scores": [], "labels": [], "masks": None}

    b = b[m]; score = score[m]; cls_idx = cls_idx[m]

    # decode xywh -> xyxy, scale back to original
    xyxy = xywh2xyxy(b)
    xyxy_scaled = scale_coords(xyxy, img_meta, imgsz)

    # heuristic: if widths collapse (~0) but heights not, try yxwh swap
    w_pix = (xyxy_scaled[:, 2] - xyxy_scaled[:, 0])
    h_pix = (xyxy_scaled[:, 3] - xyxy_scaled[:, 1])

    if len(w_pix) > 0 and (np.median(w_pix) < 2.0 and np.median(h_pix) > 5.0):
        b_swapped = b.copy()
        b_swapped[:, [0, 1]] = b_swapped[:, [1, 0]]  # swap x<->y
        b_swapped[:, [2, 3]] = b_swapped[:, [3, 2]]  # swap w<->h
        xyxy = xywh2xyxy(b_swapped)
        xyxy_scaled = scale_coords(xyxy, img_meta, imgsz)

    xyxy = xyxy_scaled

    # NMS
    keep = nms(xyxy, score, iou_thr=iou_thr, topk=max_det)
    xyxy = xyxy[keep]; score = score[keep]; cls_idx = cls_idx[keep]

    # nicer output
    xyxy = np.round(xyxy, 2)

    boxes = xyxy.tolist()
    scores = score.tolist()
    labels = [int(c) for c in cls_idx]
    return {"boxes": boxes, "scores": scores, "labels": labels, "masks": None}