"""Self-test for the antigen puncta quantification app.

Builds synthetic fluorescence images with known ground truth and verifies
each antigen-specific behavior, image loading, calibration, and that the
Gradio app starts and serves. Run: python selftest.py
"""

import os
import tempfile
import urllib.request

import numpy as np
import tifffile
from PIL import Image

import app

RNG = np.random.default_rng(42)
FAILURES = []


def check(name, cond, detail=""):
    status = "PASS" if cond else "FAIL"
    print(f"[{status}] {name}" + (f" - {detail}" if detail else ""))
    if not cond:
        FAILURES.append(name)


def synth_field(size=512, bg=100.0, noise=5.0, gradient=40.0):
    """Background with a smooth gradient plus Gaussian noise."""
    yy, xx = np.mgrid[0:size, 0:size].astype(np.float64)
    img = bg + gradient * (xx / size)
    img += RNG.normal(0, noise, (size, size))
    return img


def add_puncta(img, n, sigma, amp, amp_jitter=0.1, margin=20, min_sep=14,
               coords=None):
    """Add n Gaussian puncta with similar peak amplitude; returns centers."""
    size = img.shape[0]
    centers = []
    existing = list(coords) if coords else []
    tries = 0
    while len(centers) < n and tries < 20000:
        tries += 1
        y = RNG.uniform(margin, size - margin)
        x = RNG.uniform(margin, size - margin)
        if any((y - cy) ** 2 + (x - cx) ** 2 < min_sep ** 2
               for cy, cx in existing + centers):
            continue
        centers.append((y, x))
    for (y, x) in centers:
        a = amp * (1 + RNG.uniform(-amp_jitter, amp_jitter))
        r = int(np.ceil(4 * sigma))
        y0, y1 = int(y) - r, int(y) + r + 1
        x0, x1 = int(x) - r, int(x) + r + 1
        yy, xx = np.mgrid[y0:y1, x0:x1].astype(np.float64)
        bump = a * np.exp(-((yy - y) ** 2 + (xx - x) ** 2) / (2 * sigma ** 2))
        img[max(0, y0):y1, max(0, x0):x1] += bump[max(0, -y0):, max(0, -x0):]
    return centers


def match_stats(truth, detected, tol=5.0):
    """Greedy nearest-neighbor matching; returns (recall, precision)."""
    truth = list(truth)
    det = [(row["y_px"], row["x_px"]) for _, row in detected.iterrows()]
    unmatched_det = list(range(len(det)))
    tp = 0
    for (ty, tx) in truth:
        best_j, best_d = None, tol
        for j in unmatched_det:
            d = np.hypot(det[j][0] - ty, det[j][1] - tx)
            if d <= best_d:
                best_j, best_d = j, d
        if best_j is not None:
            unmatched_det.remove(best_j)
            tp += 1
    recall = tp / len(truth) if truth else 1.0
    precision = tp / len(det) if det else 1.0
    return recall, precision


def test_cda_ptau():
    """CDA and PTAU: similar-intensity puncta detected, background and
    bright artifacts rejected."""
    img = synth_field()
    truth = add_puncta(img, 60, sigma=2.0, amp=300.0)
    # Bright artifacts (dust) that the similarity filter should reject.
    dust = add_puncta(img, 5, sigma=2.5, amp=4000.0, coords=truth)

    for antigen in ("CDA", "PTAU"):
        df, info = app.detect_puncta(img, app.DEFAULT_PARAMS[antigen])
        recall, precision = match_stats(truth, df)
        check(f"{antigen} recall >= 0.90", recall >= 0.90,
              f"recall={recall:.3f}, n={len(df)}")
        check(f"{antigen} precision >= 0.90", precision >= 0.90,
              f"precision={precision:.3f}")
        dust_recall, _ = match_stats(dust, df)
        check(f"{antigen} similarity filter rejects bright artifacts",
              dust_recall <= 0.2,
              f"artifacts kept={dust_recall * 5:.0f}/5, "
              f"rejected_sim={info['n_rejected_sim']}")


def test_syn_small_puncta():
    """SYN: small puncta with similar intensity that default-size settings
    tend to miss are recovered."""
    img = synth_field()
    truth = add_puncta(img, 40, sigma=1.1, amp=200.0)

    df_syn, _ = app.detect_puncta(img, app.DEFAULT_PARAMS["SYN"])
    recall_syn, precision_syn = match_stats(truth, df_syn)
    df_cda, _ = app.detect_puncta(img, app.DEFAULT_PARAMS["CDA"])
    recall_cda, _ = match_stats(truth, df_cda)

    check("SYN recall on small puncta >= 0.85", recall_syn >= 0.85,
          f"recall={recall_syn:.3f}")
    check("SYN precision >= 0.85", precision_syn >= 0.85,
          f"precision={precision_syn:.3f}")
    print(f"       (info: CDA defaults recover {recall_cda:.2f} of the same "
          f"small puncta; SYN recovers {recall_syn:.2f})")
    check("SYN recovers at least as many small puncta as CDA defaults",
          recall_syn >= recall_cda,
          f"SYN={recall_syn:.2f} CDA={recall_cda:.2f}")


def test_mbp():
    """MBP: fluorescence measured inside cells only, background excluded."""
    size = 512
    img = synth_field(size=size, gradient=10.0)
    yy, xx = np.mgrid[0:size, 0:size].astype(np.float64)
    cell = (yy - 256) ** 2 + (xx - 256) ** 2 <= 150 ** 2
    img[cell] += 60.0
    pos_truth = (yy - 256) ** 2 + (xx - 256) ** 2 <= 80 ** 2
    img[pos_truth] += 250.0

    metrics, pos, cells = app.analyze_mbp(img, app.DEFAULT_PARAMS["MBP"])
    true_frac = pos_truth.sum() / cell.sum()
    got_frac = metrics["mbp_area_fraction_of_cells"]
    check("MBP area fraction within 30 percent of truth",
          abs(got_frac - true_frac) / true_frac <= 0.30,
          f"true={true_frac:.3f} got={got_frac:.3f}")
    outside = pos & ~cell
    check("MBP positive pixels stay inside cells",
          outside.sum() <= 0.01 * max(pos.sum(), 1),
          f"outside={int(outside.sum())} of {int(pos.sum())}")
    check("MBP mean intensity reflects positive signal",
          metrics["mbp_mean_intensity_raw"] > metrics["background_median_raw"]
          + 100,
          f"mean={metrics['mbp_mean_intensity_raw']:.1f} "
          f"bg={metrics['background_median_raw']:.1f}")


def test_io_roundtrip():
    """16-bit TIFF and RGB PNG loading give the same detections."""
    img = synth_field()
    truth = add_puncta(img, 50, sigma=2.0, amp=300.0)
    tmp = tempfile.mkdtemp(prefix="puncta_selftest_")

    tif_path = os.path.join(tmp, "test16.tif")
    tifffile.imwrite(tif_path, np.clip(img, 0, 65535).astype(np.uint16))
    loaded = app.to_2d(app.load_image(tif_path))
    df, _ = app.detect_puncta(loaded, app.DEFAULT_PARAMS["CDA"])
    recall, precision = match_stats(truth, df)
    check("16-bit TIFF load + detect", recall >= 0.9 and precision >= 0.9,
          f"recall={recall:.2f} precision={precision:.2f}")

    png_path = os.path.join(tmp, "test_rgb.png")
    rgb = np.zeros((*img.shape, 3), dtype=np.uint8)
    rgb[..., 1] = np.clip(img / img.max() * 255, 0, 255).astype(np.uint8)
    Image.fromarray(rgb).save(png_path)
    loaded = app.to_2d(app.load_image(png_path), "Auto (strongest channel)")
    df, _ = app.detect_puncta(loaded, app.DEFAULT_PARAMS["CDA"])
    recall, precision = match_stats(truth, df)
    check("RGB PNG auto-channel load + detect",
          recall >= 0.85 and precision >= 0.85,
          f"recall={recall:.2f} precision={precision:.2f}")


def test_calibration():
    """Calibration matches gold-standard counts at least as well as
    defaults."""
    tmp = tempfile.mkdtemp(prefix="puncta_cal_")
    paths, counts = [], []
    for i in range(3):
        img = synth_field(size=384)
        truth = add_puncta(img, 30 + 10 * i, sigma=1.2, amp=180.0)
        p = os.path.join(tmp, f"roi_{i}.tif")
        tifffile.imwrite(p, np.clip(img, 0, 65535).astype(np.uint16))
        paths.append(p)
        counts.append(len(truth))

    state, report, msg = app.cb_calibrate(
        paths, ", ".join(str(c) for c in counts), "CDA",
        app.CHANNEL_CHOICES[0], {})
    check("Calibration stores CDA parameters", "CDA" in state)
    cal_err = np.mean([abs(r - c) / c for r, c in
                       zip(report["calibrated_count"], counts)])
    def_err = np.mean([abs(r - c) / c for r, c in
                       zip(report["default_params_count"], counts)])
    check("Calibrated count error <= default error", cal_err <= def_err + 1e-9,
          f"calibrated={cal_err:.3f} default={def_err:.3f}")
    check("Calibrated count error <= 15 percent", cal_err <= 0.15,
          f"calibrated={cal_err:.3f}")
    print(f"       (info: {msg.splitlines()[0] if msg else ''})")


def test_callbacks_and_ui():
    """End-to-end callback runs and the Gradio app serves HTTP."""
    tmp = tempfile.mkdtemp(prefix="puncta_ui_")
    img = synth_field()
    add_puncta(img, 40, sigma=2.0, amp=300.0)
    path = os.path.join(tmp, "field.tif")
    tifffile.imwrite(path, np.clip(img, 0, 65535).astype(np.uint16))

    slider_vals = [app.DEFAULT_PARAMS["CDA"][k] for k in app.PARAM_KEYS]
    overlay, summary, table, csv_path, status = app.cb_analyze(
        path, "CDA", app.CHANNEL_CHOICES[0],
        "Auto (calibrated when available, else antigen defaults)", 0.11,
        {}, *slider_vals)
    check("cb_analyze returns overlay image",
          overlay.ndim == 3 and overlay.dtype == np.uint8)
    check("cb_analyze writes CSV", os.path.isfile(csv_path))
    check("cb_analyze summary non-empty", len(summary) > 0)

    mbp_img = synth_field(gradient=10.0)
    yy, xx = np.mgrid[0:512, 0:512].astype(np.float64)
    mbp_img[(yy - 256) ** 2 + (xx - 256) ** 2 <= 150 ** 2] += 60.0
    mbp_img[(yy - 256) ** 2 + (xx - 256) ** 2 <= 80 ** 2] += 250.0
    mbp_path = os.path.join(tmp, "mbp.tif")
    tifffile.imwrite(mbp_path, np.clip(mbp_img, 0, 65535).astype(np.uint16))
    overlay, summary, table, csv_path, status = app.cb_analyze(
        mbp_path, "MBP", app.CHANNEL_CHOICES[0],
        "Auto (calibrated when available, else antigen defaults)", 0,
        {}, *slider_vals)
    check("cb_analyze MBP mode works", os.path.isfile(csv_path))

    df, csv_path, status = app.cb_batch(
        [path, mbp_path], "CDA", app.CHANNEL_CHOICES[0],
        "Auto (calibrated when available, else antigen defaults)", 0,
        {}, *slider_vals)
    check("cb_batch processes multiple files",
          len(df) == 2 and os.path.isfile(csv_path))

    ppath = app.cb_export_params({})
    check("Parameter export writes JSON", os.path.isfile(ppath))
    state, msg = app.cb_import_params(ppath, {})
    check("Parameter import restores all antigens",
          all(a in state for a in app.ALL_ANTIGENS))

    app.demo.launch(prevent_thread_lock=True, server_name="127.0.0.1",
                    quiet=True, show_error=True)
    try:
        url = app.demo.local_url
        with urllib.request.urlopen(url, timeout=15) as resp:
            ok = resp.status == 200
        check("Gradio app serves HTTP 200", ok, url)
    finally:
        app.demo.close()


def test_multichannel_layouts():
    """Channel handling for CZYX 4D stacks and 2-channel channel-last
    images (regressions found in adversarial review)."""
    img = synth_field(size=256)
    truth = add_puncta(img, 20, sigma=2.0, amp=300.0, margin=15)

    czyx = np.zeros((3, 5, 256, 256), dtype=np.float64)
    czyx[0] = 4000.0  # bright irrelevant channel (e.g. nuclear stain)
    for z in range(5):
        czyx[1, z] = img  # green carries the puncta
    czyx[2] = 50.0
    reduced = app.to_2d(czyx, "Green")
    df, _ = app.detect_puncta(reduced, app.DEFAULT_PARAMS["CDA"])
    recall, precision = match_stats(truth, df)
    check("CZYX 4D stack: selected channel survives projection",
          recall >= 0.9 and precision >= 0.9,
          f"recall={recall:.2f} precision={precision:.2f}")

    two = np.stack([img, np.full_like(img, 60.0)], axis=-1)
    reduced = app.to_2d(two, "Auto (strongest channel)")
    check("(H, W, 2) image keeps spatial shape",
          reduced.shape == (256, 256), str(reduced.shape))
    df, _ = app.detect_puncta(reduced, app.DEFAULT_PARAMS["CDA"])
    recall, _ = match_stats(truth, df)
    check("(H, W, 2) image detects puncta", recall >= 0.9,
          f"recall={recall:.2f}")


def test_nonfinite_and_lzw():
    """NaN sanitization and LZW-compressed TIFF support."""
    tmp = tempfile.mkdtemp(prefix="puncta_nan_")
    img = synth_field()
    truth = add_puncta(img, 30, sigma=2.0, amp=300.0)

    img_nan = img.astype(np.float32).copy()
    img_nan[0:4, 0:4] = np.nan
    p = os.path.join(tmp, "nan.tif")
    tifffile.imwrite(p, img_nan)
    loaded = app.to_2d(app.load_image(p))
    check("NaN pixels replaced with finite values",
          np.isfinite(loaded).all())
    df, _ = app.detect_puncta(loaded, app.DEFAULT_PARAMS["CDA"])
    recall, _ = match_stats(truth, df)
    check("Image with NaN padding still detects puncta", recall >= 0.9,
          f"recall={recall:.2f}")

    p_lzw = os.path.join(tmp, "lzw.tif")
    tifffile.imwrite(p_lzw, np.clip(img, 0, 65535).astype(np.uint16),
                     compression="lzw")
    loaded = app.to_2d(app.load_image(p_lzw))
    df, _ = app.detect_puncta(loaded, app.DEFAULT_PARAMS["CDA"])
    recall, _ = match_stats(truth, df)
    check("LZW-compressed TIFF loads and detects", recall >= 0.9,
          f"recall={recall:.2f}")

    small = np.full((10, 10), 100, dtype=np.uint16)
    p_small = os.path.join(tmp, "small.tif")
    tifffile.imwrite(p_small, small)
    slider_vals = [app.DEFAULT_PARAMS["CDA"][k] for k in app.PARAM_KEYS]
    dfb, _, _ = app.cb_batch(
        [p_small], "CDA", app.CHANNEL_CHOICES[0],
        "Auto (calibrated when available, else antigen defaults)", 0,
        {}, *slider_vals)
    check("Batch flags degenerate images in the error column",
          str(dfb.iloc[0]["error"]) != "", str(dfb.iloc[0]["error"])[:60])


def test_confusers_and_artifacts():
    """The similarity filter must quantify the signal population, not the
    majority population, and survive large saturated artifacts."""
    img = synth_field()
    real = add_puncta(img, 20, sigma=2.0, amp=300.0)
    confusers = add_puncta(img, 60, sigma=2.0, amp=60.0, coords=real)
    df, info = app.detect_puncta(img, app.DEFAULT_PARAMS["CDA"])
    r_real, _ = match_stats(real, df)
    r_conf, _ = match_stats(confusers, df)
    check("Bright true puncta kept despite dim-confuser majority",
          r_real >= 0.9, f"recall={r_real:.2f}, kept={len(df)}")
    check("Dim background confusers rejected", r_conf <= 0.2,
          f"confuser recall={r_conf:.2f}")
    check("Bimodal population reported to the user",
          info["sim_population"] == "brighter mode", info["sim_population"])

    img = synth_field()
    truth = add_puncta(img, 50, sigma=2.0, amp=120.0)
    yy, xx = np.mgrid[0:512, 0:512]
    img[(yy - 450) ** 2 + (xx - 450) ** 2 <= 40 ** 2] = 60000.0
    far_truth = [(y, x) for (y, x) in truth
                 if (y - 450) ** 2 + (x - 450) ** 2 > 70 ** 2]
    df, info = app.detect_puncta(img, app.DEFAULT_PARAMS["CDA"])
    recall, precision = match_stats(far_truth, df)
    check("Dim puncta survive one large saturated artifact",
          recall >= 0.9, f"recall={recall:.2f} of {len(far_truth)}")
    check("Artifact rim detections are filtered out", precision >= 0.85,
          f"precision={precision:.2f}, kept={len(df)}")


def test_mbp_negative_control():
    """MBP must refuse to fabricate measurements on cell-free fields but
    still measure genuinely confluent fields with real signal."""
    for name, field in (("gradient", synth_field()),
                        ("flat", synth_field(gradient=0.0))):
        try:
            app.analyze_mbp(field, app.DEFAULT_PARAMS["MBP"])
            raised = False
        except ValueError:
            raised = True
        check(f"MBP negative control ({name} blank) raises instead of "
              f"fabricating signal", raised)

    confluent = synth_field(gradient=10.0) + 60.0  # cells fill the field
    yy, xx = np.mgrid[0:512, 0:512].astype(np.float64)
    confluent[(yy - 256) ** 2 + (xx - 256) ** 2 <= 128 ** 2] += 250.0
    try:
        metrics, pos, cells = app.analyze_mbp(confluent,
                                              app.DEFAULT_PARAMS["MBP"])
        measured = metrics["mbp_positive_area_px"] > 0
    except ValueError as exc:
        measured = False
    check("Confluent field with real MBP signal is still measured",
          measured)


def test_calibration_tiebreak():
    """On a clean gold ROI where many settings tie at zero error, the
    calibration must pick strict settings that generalize."""
    tmp = tempfile.mkdtemp(prefix="puncta_tie_")
    roi = synth_field(noise=2.0, gradient=10.0)
    truth = add_puncta(roi, 40, sigma=2.0, amp=300.0)
    p_roi = os.path.join(tmp, "clean_roi.tif")
    tifffile.imwrite(p_roi, np.clip(roi, 0, 65535).astype(np.uint16))

    state, report, msg = app.cb_calibrate(
        [p_roi], str(len(truth)), "CDA", app.CHANNEL_CHOICES[0], {})
    p = state["CDA"]
    check("Calibration does not pick the most permissive tie",
          p["log_threshold"] > 0.02 or p["min_snr"] > 2.0,
          f"threshold={p['log_threshold']} snr={p['min_snr']} "
          f"sim={p['sim_filter']}")

    routine = synth_field()
    truth2 = add_puncta(routine, 40, sigma=2.0, amp=120.0)
    df, _ = app.detect_puncta(routine, p)
    recall, precision = match_stats(truth2, df)
    check("Calibrated params generalize to a dimmer routine image",
          recall >= 0.85 and precision >= 0.8,
          f"recall={recall:.2f} precision={precision:.2f}")


def main():
    print("Running antigen puncta app self-test\n")
    test_cda_ptau()
    test_syn_small_puncta()
    test_mbp()
    test_io_roundtrip()
    test_calibration()
    test_multichannel_layouts()
    test_nonfinite_and_lzw()
    test_confusers_and_artifacts()
    test_mbp_negative_control()
    test_calibration_tiebreak()
    test_callbacks_and_ui()
    print()
    if FAILURES:
        print(f"FAILED: {len(FAILURES)} check(s): {FAILURES}")
        raise SystemExit(1)
    print("All checks passed.")


if __name__ == "__main__":
    main()