"""
focus_qc.py - input-image focus QC for hemocytometer blocks (v6, 1 Oct 2026).

Runs BEFORE segmentation. The 4-point block crop (the 1 mm x 1 mm counting
block) is warped to a 1024 x 1024 px square (0.977 um/px), so the six inner grid
lines always sit at 256, 512 and 768 px on both axes. For each inner line:

  * median intensity along the line, for every position across it (the median
    ignores cells lying on the line) -> a 1-D profile spanning +/-40 px;
  * background = median of the outer 12 px on each side of the profile;
  * contrast   = (peak - background) / background
  * FWHM       = sub-pixel full width at half of (peak - background), in um.

Per image: median over the six lines.

Decision: **FWHM > FWHM_LIMIT_UM (9 um) = out of focus, does not pass QC.**
Contrast is measured and reported but is NOT used to decide: it depends on
illumination and phase setting as well as focus (faint grey lines on sharp
photos), so it overlaps between in- and out-of-focus fields.

Validation (20260901_Asad, 48 fields, mutHEK / HT1080 / COS7; Kevin's labels):
in-focus FWHM 2.5-8.5 um (n=43), out-of-focus 9.8-10.6 um (n=5); AUC 1.00,
48/48 correct at 9 um. Contrast AUC 0.95. Write-up:
Z:\\Manuscripts\\cellposecellcounter\\manuscript\\out-of-focus qc control.md
Caveat: 5 positives from one COS7 session; re-check the limit on more data.

The measurement is identical to 20260901_Asad/results_v5/focus_qc_gridline.py.
"""
import numpy as np
import cv2

BLOCK_PX = 1024                 # 1 mm block warped to this many px
UM_PER_PX = 1000.0 / BLOCK_PX   # 0.977 um/px
LINES_PX = (256, 512, 768)      # the 3 inner rulings of a 4x4 block, per axis
HALF_WIN = 40                   # profile half-width, px (about 39 um)
BG_PX = 12                      # outer px per side used as background
FWHM_LIMIT_UM = 9.0             # out of focus above this


def _order_quad(pts):
    """(TL, TR, BR, BL) for any corner order, rotation-safe (as app.order_quad)."""
    pts = np.asarray(pts, dtype=np.float32)
    c = pts.mean(axis=0)
    ang = np.arctan2(pts[:, 1] - c[1], pts[:, 0] - c[0])
    o = pts[np.argsort(ang)]
    return np.roll(o, -int(np.argmin(o.sum(axis=1))), axis=0)


def warp_block(image_np, corners, size=BLOCK_PX):
    src = _order_quad(corners)
    dst = np.array([[0, 0], [size, 0], [size, size], [0, size]], np.float32)
    M = cv2.getPerspectiveTransform(src, dst)
    return cv2.warpPerspective(image_np, M, (size, size))


def _line_feats(profile, c, W=HALF_WIN, bg=BG_PX):
    seg = profile[c - W:c + W + 1].astype(float)
    base = float(np.median(np.r_[seg[:bg], seg[-bg:]]))
    if base <= 0:
        return None
    s = seg / base - 1.0
    k = int(np.argmax(s[bg:-bg])) + bg
    amp = float(s[k])
    if amp <= 0:
        return None
    half = amp / 2.0
    l = r = k
    while l > 0 and s[l - 1] > half:
        l -= 1
    while r < len(s) - 1 and s[r + 1] > half:
        r += 1
    fl = l - (s[l] - half) / (s[l] - s[l - 1] + 1e-9) if l > 0 else l
    fr = r + (s[r] - half) / (s[r] - s[r + 1] + 1e-9) if r < len(s) - 1 else r
    return amp, (fr - fl) * UM_PER_PX


def measure_block(block_rgb_or_gray):
    """Grid-line contrast and FWHM (um) of a 1024 px block. Returns dict."""
    g = block_rgb_or_gray
    if g.ndim == 3:
        g = cv2.cvtColor(g[..., :3], cv2.COLOR_RGB2GRAY)
    g = g.astype(np.float32)
    if g.shape != (BLOCK_PX, BLOCK_PX):
        g = cv2.resize(g, (BLOCK_PX, BLOCK_PX), interpolation=cv2.INTER_AREA)
    feats = []
    for ax in (0, 1):
        prof = np.median(g, axis=ax)
        for c in LINES_PX:
            f = _line_feats(prof, c)
            if f is not None:
                feats.append(f)
    if len(feats) < 3:
        return {"contrast": None, "fwhm_um": None, "n_lines": len(feats)}
    con, fw = np.median(np.array(feats), axis=0)
    return {"contrast": float(con), "fwhm_um": float(fw), "n_lines": len(feats)}


def check_focus(image_np, corners, limit_um=FWHM_LIMIT_UM):
    """QC an uploaded photo + its 4 block corners (full-resolution pixels).

    Returns {"ok", "measured", "fwhm_um", "contrast", "limit_um", "message"}.
    ok=False only when the grid lines were measured AND are too wide. If they
    cannot be measured (fewer than 3 of 6 lines found), ok=True with
    measured=False, so a measurement problem never blocks counting.
    """
    out = {"ok": True, "measured": False, "fwhm_um": None, "contrast": None,
           "limit_um": float(limit_um), "message": ""}
    if corners is None or len(corners) != 4:
        out["message"] = "Focus QC skipped: needs the 4-point block crop."
        return out
    img = np.asarray(image_np)
    if img.ndim == 3 and img.shape[2] == 4:
        img = cv2.cvtColor(img, cv2.COLOR_RGBA2RGB)
    m = measure_block(warp_block(img, corners))
    out.update(fwhm_um=m["fwhm_um"], contrast=m["contrast"])
    if m["fwhm_um"] is None:
        out["message"] = ("Focus QC could not measure the grid lines "
                          f"({m['n_lines']} of 6 found); not applied.")
        return out
    out["measured"] = True
    out["ok"] = m["fwhm_um"] <= limit_um
    if out["ok"]:
        out["message"] = (f"Focus QC passed: grid-line width {m['fwhm_um']:.1f} um "
                          f"(limit {limit_um:.0f} um), contrast {m['contrast']:.3f}.")
    else:
        out["message"] = (f"Out of focus, does not pass QC: grid-line width "
                          f"{m['fwhm_um']:.1f} um (limit {limit_um:.0f} um), "
                          f"contrast {m['contrast']:.3f}. "
                          "Refocus on the grid and retake this photo.")
    return out