"""Detector pre/post-processing shared by evaluation and inference."""

import cv2
import numpy as np

from ..geometry import expand_quad, quad_size
from ..models.detector import KERNEL_END_TRIM, KERNEL_HEIGHT_RATIO

DET_MAX_SIDE = 1920
DET_MIN_SIDE = 960


def detector_input(img_rgb, max_side=DET_MAX_SIDE, min_side=DET_MIN_SIDE):
    """Resize so the long side is within [min_side, max_side], pad to a multiple of 32.

    Returns (float32 CHW in [0,1], scale) where model coords = original coords * scale.
    """
    h, w = img_rgb.shape[:2]
    long = max(h, w)
    scale = min(max_side / long, max(1.0, min_side / long))
    nh, nw = max(32, round(h * scale)), max(32, round(w * scale))
    img = cv2.resize(img_rgb, (nw, nh), interpolation=cv2.INTER_AREA if scale < 1 else cv2.INTER_LINEAR)
    ph, pw = (nh + 31) // 32 * 32, (nw + 31) // 32 * 32
    pad = np.empty((ph, pw, 3), np.uint8)
    pad[:] = np.median(img[::8, ::8].reshape(-1, 3), 0).astype(np.uint8)   # page colour from a pixel sample
    pad[:nh, :nw] = img
    return pad.transpose(2, 0, 1).astype(np.float32) / 255.0, scale


def _text_quad(rect):
    """minAreaRect -> quad ordered TL, TR, BR, BL with the top edge as close to horizontal as possible."""
    pts = cv2.boxPoints(rect)
    c = pts.mean(0)
    pts = pts[np.argsort(np.arctan2(pts[:, 1] - c[1], pts[:, 0] - c[0]))]  # clockwise in image coords
    best, best_score = 0, -1e9
    for s in range(4):
        v = pts[(s + 1) % 4] - pts[s]
        score = v[0] / (np.linalg.norm(v) + 1e-6)   # cos of angle to +x
        if score > best_score:
            best, best_score = s, score
    return np.roll(pts, -best, axis=0).astype(np.float32)


def kernels_to_quads(prob, bin_thresh=0.3, box_thresh=0.55, min_area=4):
    """prob: (H, W) kernel probability map. Returns list of (quad, score) in map coordinates."""
    bitmap = (prob > bin_thresh).astype(np.uint8)
    n, labels, stats, _ = cv2.connectedComponentsWithStats(bitmap, connectivity=4)
    out = []
    for i in range(1, n):
        x, y, w, h, area = stats[i]
        if area < min_area:
            continue
        comp = labels[y:y + h, x:x + w] == i
        score = float(prob[y:y + h, x:x + w][comp].mean())
        if score < box_thresh:
            continue
        ys, xs = np.nonzero(comp)
        pts = np.stack([xs + x, ys + y], 1).astype(np.float32)
        (cx, cy), (rw, rh), ang = cv2.minAreaRect(pts)
        q = _text_quad(((cx, cy), (rw, rh), ang))  # fillPoly includes boundary pixels, so centres already span the kernel
        kw, kh = quad_size(q)
        full_h = kh / KERNEL_HEIGHT_RATIO
        dy = (full_h - kh) / 2
        dx = min(KERNEL_END_TRIM * full_h, kw / 2)
        out.append((expand_quad(q, dx, dy, dx, dy), score))
    return out


def quad_iou(a, b):
    a, b = np.asarray(a, np.float32), np.asarray(b, np.float32)
    inter, _ = cv2.intersectConvexConvex(a, b)
    union = abs(cv2.contourArea(a)) + abs(cv2.contourArea(b)) - inter
    return float(inter / union) if union > 0 else 0.0


def match_quads(pred, gt, thresh=0.5):
    """Greedy one-to-one matching by IoU. Returns (matches [(pi, gi, iou)], unmatched_pred, unmatched_gt)."""
    pairs = []
    for i, p in enumerate(pred):
        pb = cv2.boundingRect(np.asarray(p, np.float32))
        for j, g in enumerate(gt):
            gb = cv2.boundingRect(np.asarray(g, np.float32))
            if pb[0] > gb[0] + gb[2] or gb[0] > pb[0] + pb[2] or pb[1] > gb[1] + gb[3] or gb[1] > pb[1] + pb[3]:
                continue
            iou = quad_iou(p, g)
            if iou >= thresh:
                pairs.append((iou, i, j))
    pairs.sort(reverse=True)
    used_p, used_g, matches = set(), set(), []
    for iou, i, j in pairs:
        if i not in used_p and j not in used_g:
            used_p.add(i)
            used_g.add(j)
            matches.append((i, j, iou))
    return matches, [i for i in range(len(pred)) if i not in used_p], [j for j in range(len(gt)) if j not in used_g]


def find_rules(img, text_height):
    """Printed horizontal rules (thin, long dark lines) as [x0, y0, x1, y1] boxes.

    Newspapers and magazines separate stories with such rules; reading order uses them as separators.
    Text survives the long horizontal opening only as underlines, which are harmless separators.
    """
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) if img.ndim == 3 else img
    h, w = gray.shape
    ink = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY_INV, 31, 12)
    length = max(30, int(max(8 * text_height, w / 25)))
    lines = _long_runs(ink, length)
    lines = cv2.dilate(lines, cv2.getStructuringElement(cv2.MORPH_RECT, (max(3, int(text_height)), 1)))  # bridge dashes
    n, _, stats, _ = cv2.connectedComponentsWithStats(lines, connectivity=8)
    out = []
    for x, y, bw, bh, _ in stats[1:n]:
        if bw >= length and bh <= max(4, 0.5 * text_height):
            out.append([float(x), float(y), float(x + bw), float(y + bh)])
    return out


def _long_runs(ink, length):
    """Binary opening with a 1 x length horizontal line, computed from run lengths: keeps exactly the horizontal
    ink runs at least `length` long (several times faster than the sliding-window morphology on full pages)."""
    b = np.zeros((ink.shape[0], ink.shape[1] + 2), np.int8)
    b[:, 1:-1] = ink > 0
    d = np.diff(b, axis=1)
    ys, x0 = np.nonzero(d == 1)
    _, x1 = np.nonzero(d == -1)
    keep = np.nonzero(x1 - x0 >= length)[0]
    out = np.zeros_like(ink)
    for k in keep:
        out[ys[k], x0[k]:x1[k]] = 255
    return out


def split_gutter_crossers(gray, dets):
    """Split detected lines that run across a column gutter (two columns' lines merged by the detector).

    A line is split only where (a) prose lines just above and below it leave a clear vertical gap - a
    column gutter - on both sides, (b) the line itself has whitespace there, and (c) it is body-text sized
    (headlines legitimately span columns). dets: list of (quad, score); returns the same.
    """
    if len(dets) < 6:
        return dets
    boxes = np.array([[q[:, 0].min(), q[:, 1].min(), q[:, 0].max(), q[:, 1].max()] for q, _ in dets], np.float64)
    hs = boxes[:, 3] - boxes[:, 1]
    hmed = float(np.median(hs))
    cy = (boxes[:, 1] + boxes[:, 3]) / 2
    out = []
    for k, (q, score) in enumerate(dets):
        x0, y0, x1, y1 = boxes[k]
        if not (0.7 * hmed < hs[k] < 1.4 * hmed and x1 - x0 > 6 * hmed):
            out.append((q, score))
            continue
        near = (np.abs(cy - cy[k]) < 6 * hmed) & (boxes[:, 2] > x0) & (boxes[:, 0] < x1)
        near[k] = False
        if near.sum() < 4:
            out.append((q, score))
            continue
        lo = int(np.floor(x0))
        cov = np.zeros(int(np.ceil(x1)) - lo + 1, np.int32)
        for bx0, _, bx1, _ in boxes[near]:
            cov[max(0, int(bx0) - lo):max(0, int(bx1) - lo)] += 1
        split_at = None
        free = np.flatnonzero(np.diff(np.r_[1, (cov == 0).astype(np.int8), 1]))   # boundaries of zero runs
        for a, b in zip(free[0::2], free[1::2]):
            if b - a < 0.3 * hmed or a < hmed or b > len(cov) - hmed:
                continue                                    # too narrow, or at the line's ends
            gx0, gx1 = lo + a, lo + b
            widths = boxes[:, 2] - boxes[:, 0]
            ok = True
            for side in (near & (boxes[:, 2] <= gx0 + 1), near & (boxes[:, 0] >= gx1 - 1)):
                # a real column gutter has prose lines on both sides, above and below this line
                # (a table's column gaps under a caption are not one)
                above, below = side & (cy < cy[k]), side & (cy > cy[k])
                if above.sum() < 1 or below.sum() < 1 or np.median(widths[side]) < 6 * hmed:
                    ok = False
            if not ok:
                continue
            crop = gray[max(0, int(y0)):int(np.ceil(y1)), max(0, int(gx0 - 0.3 * hmed)):int(np.ceil(gx1 + 0.3 * hmed))]
            if crop.size == 0:
                continue
            thr = min(200, int(np.median(crop)) - 40)
            empty = (crop < thr).sum(0) == 0
            runs = np.flatnonzero(np.diff(np.r_[0, empty.astype(np.int8), 0]))
            best = max(((e - s_, s_) for s_, e in zip(runs[0::2], runs[1::2])), default=(0, 0))
            if best[0] >= 0.25 * hmed:
                split_at = int(gx0 - 0.3 * hmed) + best[1] + best[0] / 2
                break
        if split_at is None:
            out.append((q, score))
            continue
        for a, b in ((x0, split_at - 1), (split_at + 1, x1)):
            out.append((np.array([[a, y0], [b, y0], [b, y1], [a, y1]], np.float32), score))
    return out