"""Generalized (no per-image constants) map-screenshot grid calibration.

find_grid_labels() -> OCR sweep for visible large-grid labels.
estimate_pitch()   -> derive cell pitch from label-to-label spacing.
find_label_corner()-> per-label corner search via a real corner detector
                       (cv2.goodFeaturesToTrack), window scaled to pitch.
fit_grid_homography() -> RANSAC homography from however many corners
                          were found, refusing to fit degenerate/thin data.
"""
import re
import cv2
import numpy as np
import pytesseract

LARGE_X = "ABCDEFGHIJKLMNOPQRST"
LABEL_RE = re.compile(r'^([A-T])(10|[1-9])$')
ANCHOR_FRAC_X = 0.10
ANCHOR_FRAC_Y = 0.15


def desaturate(rgb):
    """Replace saturated (colored) pixels with a neutral gray of the same
    brightness. The hatching overlay and unit markers are drawn in
    deliberately saturated colors while the grid + its labels are neutral
    white/cream on a grayscale photo, hatch lines crossing straight
    through a label glyph otherwise corrupt its shape enough to break
    OCR (seen: 'L7' -> 'AF', 'N8' -> 'NWS')."""
    hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV)
    sat = hsv[:, :, 1]
    gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
    out = gray.copy()
    return out, sat


def find_grid_labels(gray, band_height=160, band_step=60, scale=3.0, min_conf=30):
    H, W = gray.shape
    config = "--psm 11 -c tessedit_char_whitelist=ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
    found = {}
    for y0 in range(0, H, band_step):
        y1 = min(y0 + band_height, H)
        band = gray[y0:y1, :]
        big = cv2.resize(band, None, fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
        data = pytesseract.image_to_data(big, output_type=pytesseract.Output.DICT, config=config)
        for i, txt in enumerate(data['text']):
            t = txt.strip()
            m = LABEL_RE.match(t)
            if not m or int(data['conf'][i]) < min_conf:
                continue
            conf = int(data['conf'][i])
            key = t
            if key not in found or conf > found[key]['conf']:
                found[key] = dict(
                    text=t, letter=m.group(1), number=int(m.group(2)),
                    x=data['left'][i]/scale, y=y0 + data['top'][i]/scale,
                    w=data['width'][i]/scale, h=data['height'][i]/scale, conf=conf,
                )
    return list(found.values())


def estimate_pitch(labels):
    xs_by_row, ys_by_col = {}, {}
    for lb in labels:
        xs_by_row.setdefault(lb['number'], []).append((LARGE_X.index(lb['letter']), lb['x']+lb['w']/2))
        ys_by_col.setdefault(lb['letter'], []).append((lb['number'], lb['y']+lb['h']/2))
    px = []
    for pts in xs_by_row.values():
        pts.sort()
        for (ci,xi),(cj,xj) in zip(pts, pts[1:]):
            if cj > ci:
                px.append((xj-xi)/(cj-ci))
    py = []
    for pts in ys_by_col.values():
        pts.sort()
        for (ri,yi),(rj,yj) in zip(pts, pts[1:]):
            if rj != ri:
                py.append(abs(yj-yi)/abs(rj-ri))
    pitch_x = float(np.median(px)) if px else None
    pitch_y = float(np.median(py)) if py else None
    if pitch_x and not pitch_y:
        pitch_y = pitch_x
    if pitch_y and not pitch_x:
        pitch_x = pitch_y
    return pitch_x, pitch_y


def cross_kernel(size, arm_width):
    """Matched filter for a bright axis-aligned line CROSSING in both
    directions through the center, not just any corner-like feature. A
    generic corner detector (Shi-Tomasi/Harris) fires just as happily on
    a diamond marker's vertex or a single line's endpoint kink as on a
    real grid intersection, this is specific to the one shape we
    actually want: positive along both the horizontal and vertical arm,
    negative in the four quadrant gaps between them (a lone single-
    direction line only lights up one arm and loses on the other three
    quadrants plus the missing arm, scoring far below a true crossing).
    size/arm_width are in PIXELS, computed by the caller from the
    current image's own pitch, never a fixed constant, a 25px kernel
    tuned for a 750px cell is meaningless on a 150px cell."""
    size = max(int(size) | 1, 9)      # odd, sane minimum
    arm_width = max(int(round(arm_width)), 1)
    k = np.full((size, size), -1.0, dtype=np.float32)
    c = size // 2
    half = arm_width // 2
    k[c-half:c+half+1, :] = 1.0
    k[:, c-half:c+half+1] = 1.0
    k -= k.mean()
    k /= np.abs(k).sum()
    return k


def find_crossing(gray, hsv_sat, guess_x, guess_y, pitch_x, pitch_y, window_frac=0.12,
                   label_bbox=None, label_margin=6, sat_thresh=60, min_score=0.12):
    """Locate the grid-line intersection nearest this label's anchor
    guess, by convolving a cross/intersection matched filter (see
    cross_kernel()) against a local, saturation-masked, label-masked
    window, then taking the response peak, weighted by a Gaussian
    falloff in distance from the guess so a stronger-but-farther false
    cross elsewhere in the window doesn't win over the real, closer one.
    window_frac is deliberately small: the anchor guess (from label
    position + pitch, itself derived from real label spacing) should
    already be close, a small margin covers its own slop without
    covering enough area to catch an unrelated intersection."""
    H, W = gray.shape
    # search window and match kernel both scale off THIS image's pitch,
    # not fixed pixel constants, so this works whether a cell is 150px or
    # 750px across. Kernel arm needs to reach far enough to genuinely
    # distinguish 'line extends in this direction' from noise, but must
    # stay smaller than the window it slides within.
    wx, wy = max(window_frac*pitch_x, 12), max(window_frac*pitch_y, 12)
    x0, x1 = max(0,int(guess_x-wx)), min(W,int(guess_x+wx))
    y0, y1 = max(0,int(guess_y-wy)), min(H,int(guess_y+wy))
    kernel_size = max(int(0.9 * min(wx, wy)), 7)
    kernel_arm = max(pitch_x, pitch_y) * 0.01
    cross = cross_kernel(kernel_size, kernel_arm)
    blur_kernel = max(int(kernel_arm * 4) | 1, 5)
    if x1-x0 < cross.shape[1] or y1-y0 < cross.shape[0]:
        return None, (guess_x, guess_y), (x0,y0,x1,y1)
    roi = gray[y0:y1, x0:x1].astype(np.float32)
    baseline = cv2.medianBlur(gray[y0:y1, x0:x1], blur_kernel).astype(np.float32)
    excess = np.clip(roi - baseline, 0, 60)
    sat_roi = hsv_sat[y0:y1, x0:x1]
    excess[sat_roi > sat_thresh] = 0
    if label_bbox is not None:
        lx0 = int(label_bbox['x']) - label_margin - x0
        ly0 = int(label_bbox['y']) - label_margin - y0
        lx1 = int(label_bbox['x'] + label_bbox['w']) + label_margin - x0
        ly1 = int(label_bbox['y'] + label_bbox['h']) + label_margin - y0
        lx0, ly0 = max(0, lx0), max(0, ly0)
        lx1, ly1 = min(excess.shape[1], lx1), min(excess.shape[0], ly1)
        if lx1 > lx0 and ly1 > ly0:
            excess[ly0:ly1, lx0:lx1] = 0

    response = cv2.filter2D(excess, -1, cross)
    half = cross.shape[0] // 2
    response[:half, :] = -1e9; response[-half:, :] = -1e9
    response[:, :half] = -1e9; response[:, -half:] = -1e9

    # let the matched-filter response do the actual selecting (find the
    # strongest genuine cross in the window), the Gaussian prior only
    # nudges among near-tied candidates and sanity-checks the winner
    # isn't implausibly far from the guess, it was previously multiplied
    # straight into the per-pixel score, which let a weak-but-central
    # false response beat a real, stronger crossing nearby.
    yy, xx = np.mgrid[0:response.shape[0], 0:response.shape[1]].astype(np.float32)
    gx, gy = guess_x - x0, guess_y - y0
    sigma = 1.5 * min(wx, wy)
    prior = np.exp(-((xx-gx)**2 + (yy-gy)**2) / (2*sigma**2))
    weighted = response * (0.85 + 0.15*prior)
    py, px = np.unravel_index(np.argmax(weighted), weighted.shape)
    peak_response = response[py, px]
    if peak_response < min_score * np.abs(cross).sum() * 60:
        return None, (guess_x, guess_y), (x0,y0,x1,y1)
    return (float(px+x0), float(py+y0)), (guess_x, guess_y), (x0,y0,x1,y1)


def find_label_corner(gray, hsv_sat, label, pitch_x, pitch_y, **kw):
    guess_x = label['x'] - ANCHOR_FRAC_X*pitch_x
    guess_y = label['y'] - ANCHOR_FRAC_Y*pitch_y
    return find_crossing(gray, hsv_sat, guess_x, guess_y, pitch_x, pitch_y, label_bbox=label, **kw)


def span_grid_crossings(gray, hsv_sat, labels, pitch_x, pitch_y, margin_cells=1):
    """Once we have a rough pitch/origin from however many labels OCR'd
    (even just 2-3), the whole grid is regular, so predict and directly
    test EVERY crossing across the visible frame, not just the ones next
    to a label that happened to be readable. Returns (ideal_pts,
    img_pts, debug) for every crossing that matched; a bad/outlier label
    just contributes points that RANSAC discards downstream rather than
    limiting how much of the grid we ever attempt.

    dc = column offset in cells from the reference label. dr = offset in
    cells DOWN the image (increasing y) from the reference, so it moves
    opposite to row number (row 8 sits above row 7 on screen): the ideal
    row coordinate is -(ref_number - dr) = dr - ref_number, matching the
    (col, -row_number) convention used everywhere else in this module."""
    H, W = gray.shape
    if not labels:
        return [], [], []
    ref = max(labels, key=lambda lb: lb['conf'])
    ref_col = LARGE_X.index(ref['letter'])
    ref_corner_x = ref['x'] - ANCHOR_FRAC_X*pitch_x
    ref_corner_y = ref['y'] - ANCHOR_FRAC_Y*pitch_y

    dc_lo = int(np.floor((0 - ref_corner_x) / pitch_x)) - margin_cells
    dc_hi = int(np.ceil((W - ref_corner_x) / pitch_x)) + margin_cells
    dr_lo = int(np.floor((0 - ref_corner_y) / pitch_y)) - margin_cells
    dr_hi = int(np.ceil((H - ref_corner_y) / pitch_y)) + margin_cells

    ideal_pts, img_pts, debug = [], [], []
    for dc in range(dc_lo, dc_hi + 1):
        for dr in range(dr_lo, dr_hi + 1):
            gx = ref_corner_x + dc * pitch_x
            gy = ref_corner_y + dr * pitch_y
            if not (0 <= gx < W and 0 <= gy < H):
                continue
            corner, guess, window = find_crossing(gray, hsv_sat, gx, gy, pitch_x, pitch_y)
            debug.append((guess, corner, window))
            if corner is not None:
                ideal_pts.append((ref_col + dc, dr - ref['number']))
                img_pts.append(corner)
    return ideal_pts, img_pts, debug


def has_diversity(ideal_pts):
    cols = set(p[0] for p in ideal_pts)
    rows = set(p[1] for p in ideal_pts)
    return len(cols) >= 2 and len(rows) >= 2


def fit_grid_homography(ideal_pts, img_pts):
    if len(ideal_pts) < 4 or not has_diversity(ideal_pts):
        return None
    ideal = np.array(ideal_pts, dtype=np.float32)
    img = np.array(img_pts, dtype=np.float32)
    Hmat, mask = cv2.findHomography(ideal, img, cv2.RANSAC, 15.0)
    return Hmat, mask