"""Heatmap decoding, box fusion and the metric-aware post-processing.

Two dataset facts are exploited here:
  * no ground-truth box is a pure deletion (0 of 1404), so a candidate whose
    template side has ink and whose photo side does not is dropped;
  * the ink-blob box sits within (0, 0, +1, +1) of the ground-truth box, so a
    constant margin is added after snapping.
"""
from __future__ import annotations

import cv2
import numpy as np
import torch
import torch.nn.functional as F

from .config import BOX_MARGIN, MAX_BOXES_PER_IMAGE, OUT_STRIDE, TOPK_PER_TILE, WBF_IOU
from .metric import iou_matrix


def decode_heatmap(hm_logits: torch.Tensor, wh: torch.Tensor, off: torch.Tensor,
                   score_thr: float = 0.05, topk: int = TOPK_PER_TILE,
                   stride: int = OUT_STRIDE):
    """Single-image decode. Returns (boxes (N,4) in tile pixels, scores (N,))."""
    hm = torch.sigmoid(hm_logits)
    keep = (F.max_pool2d(hm, 3, stride=1, padding=1) == hm).float()
    hm = hm * keep
    scores, flat = hm.reshape(-1).topk(min(topk, hm.numel()))
    size = hm.shape[-1]
    ys = (flat // size).float()
    xs = (flat % size).float()

    off_flat = off.reshape(2, -1)[:, flat]
    wh_flat = wh.reshape(2, -1)[:, flat]
    cx = (xs + off_flat[0]) * stride
    cy = (ys + off_flat[1]) * stride
    w, h = wh_flat[0].clamp(min=1.0), wh_flat[1].clamp(min=1.0)
    boxes = torch.stack([cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2], 1)

    m = scores >= score_thr
    return boxes[m].cpu().numpy(), scores[m].cpu().numpy()


def wbf(boxes: np.ndarray, scores: np.ndarray, iou_thr: float = WBF_IOU):
    """Weighted Boxes Fusion. Averages coordinates instead of discarding them,
    which matters at IoU 0.5 on 8x8 boxes."""
    if len(boxes) == 0:
        return boxes, scores
    order = np.argsort(-scores)
    boxes, scores = boxes[order], scores[order]
    clusters: list[list[int]] = []
    fused: list[np.ndarray] = []
    fused_scores: list[float] = []

    for i in range(len(boxes)):
        placed = False
        if fused:
            ious = iou_matrix(boxes[i:i + 1], np.asarray(fused, np.float32))[0]
            j = int(np.argmax(ious))
            if ious[j] >= iou_thr:
                clusters[j].append(i)
                members = clusters[j]
                w = scores[members]
                fused[j] = (boxes[members] * w[:, None]).sum(0) / w.sum()
                fused_scores[j] = float(w.sum() / min(len(members) + 1, 3))
                placed = True
        if not placed:
            clusters.append([i])
            fused.append(boxes[i].copy())
            fused_scores.append(float(scores[i]))

    f = np.asarray(fused, np.float32)
    s = np.clip(np.asarray(fused_scores, np.float32), 0, 1)
    order = np.argsort(-s)
    return f[order], s[order]


def polarity_filter(boxes: np.ndarray, scores: np.ndarray, tn: np.ndarray, pn: np.ndarray,
                    ink_thr: float = 3.0):
    """Drop candidates that look like a deletion (template inked, photo blank)."""
    if len(boxes) == 0:
        return boxes, scores
    keep = np.ones(len(boxes), bool)
    for i, (x1, y1, x2, y2) in enumerate(boxes.astype(int)):
        x1, y1 = max(0, x1), max(0, y1)
        x2, y2 = min(tn.shape[1], x2), min(tn.shape[0], y2)
        if x2 <= x1 or y2 <= y1:
            keep[i] = False
            continue
        t_ink = float(tn[y1:y2, x1:x2].mean())
        p_ink = float(pn[y1:y2, x1:x2].mean())
        if t_ink > ink_thr and p_ink < ink_thr:
            keep[i] = False
    return boxes[keep], scores[keep]


def snap_to_ink(boxes: np.ndarray, tn: np.ndarray, pn: np.ndarray, pad: int = 6,
                diff_thr: float = 25.0, margin: tuple[int, int, int, int] = BOX_MARGIN,
                max_shift: int = 6) -> np.ndarray:
    """Refit each box to the local ink-difference blob, then apply the measured margin."""
    if len(boxes) == 0:
        return boxes
    diff = np.abs(pn.astype(np.float32) - tn.astype(np.float32))
    out = boxes.copy()
    h, w = diff.shape
    for i, (x1, y1, x2, y2) in enumerate(boxes):
        a, b = int(max(0, y1 - pad)), int(min(h, y2 + pad))
        c, d = int(max(0, x1 - pad)), int(min(w, x2 + pad))
        if b <= a or d <= c:
            continue
        m = diff[a:b, c:d] > diff_thr
        if m.sum() < 3:
            continue
        ys, xs = np.nonzero(m)
        nx1, ny1 = xs.min() + c, ys.min() + a
        nx2, ny2 = xs.max() + 1 + c, ys.max() + 1 + a
        cand = np.array([nx1 - margin[0], ny1 - margin[1], nx2 + margin[2], ny2 + margin[3]],
                        np.float32)
        if np.abs(cand - boxes[i]).max() <= max_shift:
            out[i] = cand
    return out


def clip_boxes(boxes: np.ndarray, h: int, w: int) -> np.ndarray:
    if len(boxes) == 0:
        return boxes
    b = boxes.copy()
    b[:, [0, 2]] = np.clip(b[:, [0, 2]], 0, w)
    b[:, [1, 3]] = np.clip(b[:, [1, 3]], 0, h)
    return b


def postprocess(boxes: np.ndarray, scores: np.ndarray, tn: np.ndarray, pn: np.ndarray,
                use_snap: bool = True, use_polarity: bool = True):
    if len(boxes) == 0:
        return boxes, scores
    h, w = tn.shape[:2]
    boxes = clip_boxes(boxes, h, w)
    if use_polarity:
        boxes, scores = polarity_filter(boxes, scores, tn, pn)
    if use_snap and len(boxes):
        boxes = snap_to_ink(boxes, tn, pn)
    keep = (boxes[:, 2] - boxes[:, 0] > 2) & (boxes[:, 3] - boxes[:, 1] > 2)
    boxes, scores = boxes[keep], scores[keep]
    if len(boxes) > MAX_BOXES_PER_IMAGE:
        order = np.argsort(-scores)[:MAX_BOXES_PER_IMAGE]
        boxes, scores = boxes[order], scores[order]
    return boxes, scores