"""GPU-accelerated template matching using FFT correlation + integral images.

Math is identical to OpenCV's TM_SQDIFF_NORMED + TM_CCOEFF_NORMED, but the
heavy operations are restructured to exploit GPU strengths:

  * Cross-correlation SI = sum(I[y+u, x+v] * T[u, v]) is computed in the
    frequency domain. We FFT the drawing ONCE, then per-variant pad+FFT the
    template and elementwise-multiply. On Blackwell (sm_120), cuDNN's direct
    conv kernels are unoptimized and take ~1 s for an 8000² drawing with a
    100² kernel — FFT correlation does the same work in ~12 ms.

  * Local sums S1 = sum(I in window) and S2 = sum(I² in window) are computed
    from integral images via four-corner subtraction — O(1) per output pixel,
    so per-variant cost is essentially free.

Putting these together: per-variant cost drops from ~3 s (direct conv) to
~15 ms, and the cost is dominated by the one-time drawing FFT (~230 ms on
an 8000×6000 page).

API and return shape match the CPU implementation in counting.py.
"""
from __future__ import annotations

import logging
from pathlib import Path
from typing import Iterable

import cv2
import numpy as np
import torch
import torch.fft as tfft

logger = logging.getLogger(__name__)

DEFAULT_THRESHOLD = 0.7
MAX_DRAWING_PX = 9000
EPS = 1e-6

# Binarization: pixels darker than this become "ink". Kept at 245 — measured
# the actual renders: drawing strokes are crisp (well below 245) so tightening
# barely changes the drawing ink (2.10% vs 2.24% at 210 vs 245), but faint
# templates (e.g. a low-contrast symbol crop, 241–255) lose ALL ink below ~245
# and silently match nothing. Not worth the fragility — the false-positive fix
# is the consensus scoring + verify gate below, not binarization.
INK_THRESHOLD = 245

# Propose-then-verify tuning.
#  - Consensus uses min(sim_sq, cc): far stricter than the old max(), so we
#    propose slightly below the user threshold to protect recall …
PROPOSE_RELAX = 0.08
#  - … then a geometric ink-overlap gate removes the candidates that only
#    matched by luck (right ink density, wrong shape).
VERIFY_INK_IOU = 0.25


def _prep(img_path: Path) -> np.ndarray:
    arr = cv2.imread(str(img_path), cv2.IMREAD_GRAYSCALE)
    if arr is None:
        raise RuntimeError(f"Could not load {img_path}")
    _, ink = cv2.threshold(arr, INK_THRESHOLD, 255, cv2.THRESH_BINARY_INV)
    return ink


def _crop_to_content(template: np.ndarray) -> np.ndarray:
    mask = template > 0
    if not mask.any():
        return template
    ys, xs = np.where(mask)
    y0, y1 = ys.min(), ys.max() + 1
    x0, x1 = xs.min(), xs.max() + 1
    pad = 2
    y0 = max(0, y0 - pad)
    x0 = max(0, x0 - pad)
    y1 = min(template.shape[0], y1 + pad)
    x1 = min(template.shape[1], x1 + pad)
    return template[y0:y1, x0:x1]


def _rotate(img: np.ndarray, angle_deg: int) -> np.ndarray:
    if angle_deg == 0:
        return img
    if angle_deg in (90, 180, 270):
        return np.rot90(img, k=angle_deg // 90)
    h, w = img.shape[:2]
    center = (w / 2, h / 2)
    M = cv2.getRotationMatrix2D(center, angle_deg, 1.0)
    cos = abs(M[0, 0]); sin = abs(M[0, 1])
    new_w = int(h * sin + w * cos)
    new_h = int(h * cos + w * sin)
    M[0, 2] += new_w / 2 - center[0]
    M[1, 2] += new_h / 2 - center[1]
    return cv2.warpAffine(img, M, (new_w, new_h),
                          flags=cv2.INTER_NEAREST, borderValue=0)


def _nms(boxes: list[tuple], scores: list[float],
         overlap_thresh: float = 0.3) -> list[int]:
    if not boxes:
        return []
    boxes_arr = np.array(boxes, dtype=np.float32)
    x1 = boxes_arr[:, 0]; y1 = boxes_arr[:, 1]
    x2 = boxes_arr[:, 2]; y2 = boxes_arr[:, 3]
    areas = (x2 - x1) * (y2 - y1)
    order = np.argsort(scores)[::-1]
    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(int(i))
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])
        w = np.maximum(0.0, xx2 - xx1)
        h = np.maximum(0.0, yy2 - yy1)
        inter = w * h
        union = areas[i] + areas[order[1:]] - inter
        iou = inter / np.maximum(union, EPS)
        order = order[1:][iou <= overlap_thresh]
    return keep


def _integral(I: torch.Tensor) -> torch.Tensor:
    """Build integral image. Output shape: (H+1, W+1), with zero first row/col.

    II[i, j] = sum(I[:i, :j]). Then for any rectangle [y, y+h) × [x, x+w):
        local_sum = II[y+h, x+w] - II[y, x+w] - II[y+h, x] + II[y, x]
    """
    H, W = I.shape
    II = torch.zeros((H + 1, W + 1), device=I.device, dtype=I.dtype)
    II[1:, 1:] = I.cumsum(0).cumsum(1)
    return II


def _local_sum(II: torch.Tensor, h: int, w: int) -> torch.Tensor:
    """Sum of every h×w window using integral image. Output: (H-h+1, W-w+1)."""
    return (II[h:, w:] - II[:-h, w:] - II[h:, :-w] + II[:-h, :-w])


def _fft_correlate(
    Ifft: torch.Tensor, T: torch.Tensor, H: int, W: int,
) -> torch.Tensor:
    """Cross-correlation of drawing (whose rfft2 is Ifft) with template T.

    Returns the 'valid' region [0, H-h+1) × [0, W-w+1) — matching the
    output region of cv2.matchTemplate / F.conv2d.

    Math: cross-correlation in time domain = FFT(I) * conj(FFT(T_padded))
    in frequency domain. T_padded has T in the top-left of a zero canvas
    of size (H, W).
    """
    h, w = T.shape
    Tpad = torch.zeros((H, W), device=T.device, dtype=T.dtype)
    Tpad[:h, :w] = T
    Tfft = tfft.rfft2(Tpad)
    full = tfft.irfft2(Ifft * Tfft.conj(), s=(H, W))
    return full[: H - h + 1, : W - w + 1]


def count_template(
    template_path: Path,
    drawing_path: Path,
    threshold: float = DEFAULT_THRESHOLD,
    rotations: Iterable[int] = (0, 45, 90, 135, 180, 225, 270, 315),
    scales: Iterable[float] = (0.6, 0.75, 0.9, 1.0, 1.15, 1.35, 1.6),
    exclude_box: tuple | None = None,
    area_tolerance: float = 200.0,  # ignored — kept for API compat
    mirror: bool = True,
) -> dict:
    """GPU template matcher. Same return shape as the CPU implementation."""
    device = torch.device("cuda")

    template = _prep(template_path)
    template = _crop_to_content(template)
    drawing = _prep(drawing_path)

    coord_scale = 1.0
    if max(drawing.shape) > MAX_DRAWING_PX:
        coord_scale = max(drawing.shape) / MAX_DRAWING_PX
        new_w = int(drawing.shape[1] / coord_scale)
        new_h = int(drawing.shape[0] / coord_scale)
        drawing = cv2.resize(drawing, (new_w, new_h),
                             interpolation=cv2.INTER_AREA)

    if exclude_box is not None:
        ex_x, ex_y, ex_w, ex_h = exclude_box
        ex_x = int(ex_x / coord_scale); ex_y = int(ex_y / coord_scale)
        ex_w = int(ex_w / coord_scale); ex_h = int(ex_h / coord_scale)
        H_, W_ = drawing.shape
        ex_x = max(0, min(W_, ex_x)); ex_y = max(0, min(H_, ex_y))
        ex_x2 = max(0, min(W_, ex_x + ex_w))
        ex_y2 = max(0, min(H_, ex_y + ex_h))
        drawing[ex_y:ex_y2, ex_x:ex_x2] = 0

    if min(template.shape) < 8 or int((template > 0).sum()) < 5:
        return {"count": 0, "matches": [], "threshold_used": threshold}

    H, W = drawing.shape

    # One-time GPU setup. Use float32 normalized to [0,1] to keep magnitudes
    # tame in the FFT (large summations on uint8 → 0..255 can overflow into
    # noisier float regions of the spectrum).
    drawing_t = torch.from_numpy(drawing.astype(np.float32) / 255.0).to(device)
    drawing_sq_t = drawing_t * drawing_t

    # FFT of drawing (used by every variant for SI)
    Ifft = tfft.rfft2(drawing_t)

    # Integral images for S1 and S2 (used by every variant)
    II = _integral(drawing_t)
    II2 = _integral(drawing_sq_t)

    # Build all variants
    base_variants = [template]
    if mirror:
        base_variants.append(cv2.flip(template, 1))

    variants_np: list[np.ndarray] = []
    for variant in base_variants:
        for rot in rotations:
            rotated = _rotate(variant, rot)
            for scale in scales:
                nw = max(8, int(rotated.shape[1] * scale))
                nh = max(8, int(rotated.shape[0] * scale))
                if nh > H or nw > W:
                    continue
                tmpl = cv2.resize(rotated, (nw, nh),
                                  interpolation=cv2.INTER_AREA)
                variants_np.append(tmpl)

    if not variants_np:
        return {"count": 0, "matches": [], "threshold_used": threshold}

    all_boxes: list[tuple] = []
    all_scores: list[float] = []
    all_vidx: list[int] = []  # which variant produced each candidate (verify)

    propose_thr = max(0.0, threshold - PROPOSE_RELAX)

    for vi, v in enumerate(variants_np):
        v_norm = (v.astype(np.float32) / 255.0)
        h_t, w_t = v_norm.shape
        T = torch.from_numpy(v_norm).to(device)

        # SI = sum_{u,v} I[y+u, x+v] * T[u, v]
        SI = _fft_correlate(Ifft, T, H, W)

        # S1 = sum_{u,v} I[y+u, x+v]
        # S2 = sum_{u,v} I[y+u, x+v]^2
        S1 = _local_sum(II, h_t, w_t)
        S2 = _local_sum(II2, h_t, w_t)

        n = float(h_t * w_t)
        sumT = float(T.sum())
        sumT2 = float((T * T).sum())
        meanT = sumT / n
        var_T_sum = sumT2 - n * meanT * meanT

        # TM_SQDIFF_NORMED  →  similarity = 1 - normalized squared difference
        sq = (sumT2 - 2.0 * SI + S2) / torch.sqrt(
            torch.clamp(sumT2 * S2, min=EPS))
        sim_sq = 1.0 - sq

        if var_T_sum <= EPS:
            sim = sim_sq
        else:
            # TM_CCOEFF_NORMED
            cc_num = SI - meanT * S1
            var_I_sum = torch.clamp(S2 - S1 * S1 / n, min=0.0)
            cc = cc_num / torch.sqrt(
                torch.clamp(var_T_sum * var_I_sum, min=EPS))
            # Score = CCOEFF_NORMED only. The old code used max(sim_sq, cc),
            # which let sim_sq's spurious highs over big white/low-ink regions
            # through → thousands of false positives. cc is the shape-
            # discriminative metric and keeps recall; the ink-IoU verify gate
            # below removes cc's remaining thin-line spurious hits. (min() was
            # tried and over-pruned: sim_sq is structurally low for sparse
            # linework even on true matches, so it killed real symbols too.)
            sim = cc

        hits = sim >= propose_thr
        if hits.any():
            ys, xs = torch.where(hits)
            scores_t = sim[ys, xs]
            ys = ys.cpu().numpy(); xs = xs.cpu().numpy()
            scores_np = scores_t.cpu().numpy()
            for y, x, s in zip(ys, xs, scores_np):
                all_boxes.append(
                    (float(x), float(y), float(x + w_t), float(y + h_t)))
                all_scores.append(float(s))
                all_vidx.append(vi)

        del SI, S1, S2, sq, sim_sq, T

    del drawing_t, drawing_sq_t, Ifft, II, II2
    torch.cuda.empty_cache()

    if not all_boxes:
        return {"count": 0, "matches": [], "threshold_used": threshold}

    keep = _nms(all_boxes, all_scores, overlap_thresh=0.3)

    # ── Verify pass ──────────────────────────────────────────────
    # For each surviving box: (a) enforce the real user threshold (propose was
    # relaxed for recall), and (b) require the candidate window's ink to
    # geometrically overlap the matched template's ink. (b) is what removes the
    # leftover "right ink density, wrong shape" false positives.
    matches = []
    for i in keep:
        score = all_scores[i]
        if score < threshold:
            continue
        x0, y0, x1, y1 = all_boxes[i]
        tmpl = variants_np[all_vidx[i]]
        th, tw = tmpl.shape
        yi, xi = int(round(y0)), int(round(x0))
        if yi < 0 or xi < 0 or yi + th > H or xi + tw > W:
            continue
        win = drawing[yi:yi + th, xi:xi + tw]
        if win.shape != tmpl.shape:
            continue
        t_ink = tmpl > 0
        w_ink = win > 0
        union = int(np.logical_or(t_ink, w_ink).sum())
        if union == 0:
            continue
        iou = int(np.logical_and(t_ink, w_ink).sum()) / union
        if iou < VERIFY_INK_IOU:
            continue
        matches.append({
            "x": int(x0 * coord_scale),
            "y": int(y0 * coord_scale),
            "w": int((x1 - x0) * coord_scale),
            "h": int((y1 - y0) * coord_scale),
            "score": round(score, 3),
        })
    return {"count": len(matches), "matches": matches,
            "threshold_used": threshold}