"""Tests for the validation section: object-by-object comparison with the annotation."""
import os
import zipfile

import numpy as np
import pytest
from PIL import Image

from cellquant.annotations import RasterAnnotation
from cellquant.postprocess import Result
from cellquant.validation import (COMPARISON_HEADERS, DIFF_HEADERS, DIFF_COLORS, ValidationSet, compare, gt_result,
                                  legend_html, write_bundle)

XML = """<?xml version="1.0" encoding="utf-8"?>
<annotations><version>1.1</version><meta><task><size>1</size></task></meta>
<image id="0" name="{name}" width="{w}" height="{h}">
<polyline label="Cell, Antibody +" source="manual" occluded="0" points="10,10;30,10;30,30;10,30;10,10" z_order="0"></polyline>
<polyline label="Nostril External of Nucleus" source="manual" occluded="0" points="40,40;44,40;44,44;40,44;40,40" z_order="0"></polyline>
</image></annotations>"""


def _disc(h, w, cy, cx, r):
    yy, xx = np.mgrid[:h, :w]
    return (yy - cy) ** 2 + (xx - cx) ** 2 <= r * r


def _case():
    """Four annotated nuclei: 1 matched with the same label, 2 matched with a different label,
    3 and 4 merged into one prediction (3 matches, 4 is missed); one extra predicted nucleus.
    Puncta: one matched, one missed, one extra."""
    H = W = 128
    gt = np.zeros((H, W), np.int32)
    gt[_disc(H, W, 30, 30, 12)] = 1
    gt[_disc(H, W, 30, 90, 12)] = 2
    gt[_disc(H, W, 90, 30, 12)] = 3
    gt[_disc(H, W, 90, 52, 10)] = 4
    gn = np.zeros((H, W), np.int32)
    gn[_disc(H, W, 30, 30, 2)] = 1
    gn[_disc(H, W, 110, 110, 2)] = 2
    R = RasterAnnotation(gt, ["", "Cell, Antibody +", "Cell, Antibody +", "Cell, Antibody +", "Cell, Antibody -"], gn,
                         ["", "Nostril in Nucleus", "Nostril External of Nucleus"])
    pr = np.zeros((H, W), np.int32)
    pr[_disc(H, W, 30, 30, 12)] = 1
    pr[_disc(H, W, 30, 90, 12)] = 2
    pr[_disc(H, W, 90, 30, 12) | _disc(H, W, 90, 52, 10)] = 3
    pr[_disc(H, W, 110, 100, 8)] = 4
    pn = np.zeros((H, W), np.int32)
    pn[_disc(H, W, 30, 30, 2)] = 1
    pn[_disc(H, W, 60, 60, 2)] = 2
    res = Result(pr, pn, [{"id": 1, "label": "Cell, Antibody +"}, {"id": 2, "label": "Cell, Antibody -"},
                          {"id": 3, "label": "Cell, Antibody -"}, {"id": 4, "label": "Cell, Antibody +"}],
                 [{"id": 1, "label": "Nostril in Nucleus"}, {"id": 2, "label": "Nostril External of Nucleus"}], {}, {})
    img = np.zeros((H, W, 3), np.uint8)
    return img, R, res


def test_compare_counts_and_differences():
    img, R, res = _case()
    c = compare("T", img, R, res)
    rows = {r[0]: r for r in c.rows}
    assert list(rows) == ["Cell, Antibody +", "Cell, Antibody -", "Cell, Nostril, Antibody +", "Cell, Nostril, Antibody -",
                          "All nuclei", "Nostril in Nucleus", "Nostril External of Nucleus", "All puncta"]
    assert len(c.rows[0]) == len(COMPARISON_HEADERS)
    # All nuclei: 4 annotated, 4 predicted, 1 matched same label, 2 label differs, 1 missed, 1 extra
    assert rows["All nuclei"][1:] == [4, 4, 0, 1, 2, 1, 1]
    assert rows["Cell, Antibody +"][1:] == [3, 2, -1, 1, 2, 0, 1]
    assert rows["All puncta"][1:] == [2, 2, 0, 1, 0, 1, 1]
    kinds = [(d[1], d[2]) for d in c.diffs]
    assert kinds == [("missed", "nucleus"), ("extra", "nucleus"), ("label differs", "nucleus"), ("label differs", "nucleus"),
                     ("missed", "punctum"), ("extra", "punctum")]
    assert [d[0] for d in c.diffs] == [1, 2, 3, 4, 5, 6] and len(c.diffs[0]) == len(DIFF_HEADERS)
    missed = c.diffs[0]
    assert missed[3] == "Cell, Antibody -" and abs(missed[5] - 52) < 1 and abs(missed[6] - 90) < 1
    assert "merged" in missed[8] and "predicted nucleus 3" in missed[8]
    extra = c.diffs[1]
    assert extra[4] == "Cell, Antibody +" and abs(extra[5] - 100) < 1 and extra[8] == "no annotated nucleus here"
    assert c.metrics["nuclei"]["f1"] == 0.75 and c.metrics["n_missed_merged"] == 1
    assert "4 annotated nuclei, 4 predicted" in c.summary and "nucleus F1 75.0%" in c.summary
    for ov in (c.overlay_gt, c.overlay_pred, c.overlay_diff, c.overlay_missed, c.overlay_extra):
        assert ov.shape == img.shape and ov.dtype == np.uint8 and ov.any()
    # the missed nucleus is outlined in the 'missed' colour, the extra one in the 'extra' colour
    assert (c.overlay_diff[_disc(128, 128, 90, 52, 10)] == DIFF_COLORS["missed"]).all(axis=-1).any()
    assert (c.overlay_diff[_disc(128, 128, 110, 100, 8)] == DIFF_COLORS["extra"]).all(axis=-1).any()
    # the missed-only panel shows the missed nucleus and punctum and nothing else; extra-only the reverse
    missed_px = (c.overlay_missed == DIFF_COLORS["missed"]).all(axis=-1)
    extra_px = (c.overlay_extra == DIFF_COLORS["extra"]).all(axis=-1)
    # (the number drawn next to the missed punctum at (110, 110) reaches into the right half of the extra nucleus,
    #  so only its left half must be free of the missed colour)
    left_half = _disc(128, 128, 110, 100, 8) & (np.arange(128)[None, :] < 104)
    assert missed_px[_disc(128, 128, 90, 52, 10)].any() and not missed_px[left_half].any()
    assert missed_px[_disc(128, 128, 110, 110, 5)].any()                      # ring around the missed punctum
    assert not (c.overlay_missed == DIFF_COLORS["extra"]).all(axis=-1).any()  # no extra objects drawn
    assert extra_px[_disc(128, 128, 110, 100, 8)].any() and not extra_px[_disc(128, 128, 90, 52, 10)].any()
    assert extra_px[_disc(128, 128, 60, 60, 5)].any()                         # ring around the extra punctum
    assert not (c.overlay_extra == DIFF_COLORS["missed"]).all(axis=-1).any()
    # matched objects are untouched on both panels (the matched nucleus keeps its raw pixels)
    assert (c.overlay_missed[_disc(128, 128, 30, 30, 9)] == img[_disc(128, 128, 30, 30, 9)]).all()


def test_perfect_prediction_has_no_differences():
    img, R, _ = _case()
    res = gt_result(R)
    c = compare("P", img, R, res)
    assert c.diffs == [] and c.metrics["nuclei"]["f1"] == 1.0 and c.metrics["puncta"]["f1"] == 1.0
    assert all(r[3] == 0 and r[6] == 0 and r[7] == 0 for r in c.rows)


def test_empty_prediction_lists_every_annotated_object_as_missed():
    img, R, _ = _case()
    res = Result(np.zeros_like(R.cell_inst), np.zeros_like(R.nostril_inst), [], [], {}, {})
    c = compare("E", img, R, res)
    assert sum(1 for d in c.diffs if d[1] == "missed") == R.n_cells + R.n_nostrils == len(c.diffs)
    assert all(d[8] == "no prediction here" for d in c.diffs if d[2] == "nucleus")


def test_validation_set_and_bundle(tmp_path):
    name = "K7 F 35D SC Iba W001 - RGB.png"
    Image.fromarray(np.zeros((60, 50, 3), np.uint8)).save(tmp_path / name)
    with zipfile.ZipFile(tmp_path / "ann.zip", "w") as zf:
        zf.writestr("annotations.xml", XML.format(name=name, w=50, h=60))
    vs = ValidationSet(str(tmp_path))
    assert vs.keys == ["K7"] and vs.choices() == [("K7  (1 nuclei, 1 puncta annotated)", "K7")]
    assert vs.image("K7").shape == (60, 50, 3)
    R = vs.raster("K7")
    assert R.n_cells == 1 and R.n_nostrils == 1
    c = compare("K7", vs.image("K7"), R, gt_result(R))
    zp = write_bundle(c, str(tmp_path / "out"))
    with zipfile.ZipFile(zp) as zf:
        names = sorted(zf.namelist())
    assert names == ["K7_all_differences.png", "K7_comparison.csv", "K7_differences.csv", "K7_extra_only.png",
                     "K7_manual.png", "K7_missed_only.png", "K7_predicted.png"]
    assert "missed" in legend_html() and "extra" in legend_html()
    # thumbnails for the image strip
    g = vs.gallery(size=32)
    assert len(g) == 1 and os.path.exists(g[0][0]) and g[0][1] == "K7: 1 nuclei, 1 puncta annotated"
    assert max(Image.open(g[0][0]).size) == 32 and vs.thumbnail("K7") == g[0][0]
    assert vs.counts("K7") == (1, 1)
    # a missing or empty folder is not an error
    assert ValidationSet(str(tmp_path / "nothing")).keys == []
    assert ValidationSet("").keys == []


def test_app_builds_with_and_without_validation_bundle(tmp_path, monkeypatch):
    pytest.importorskip("gradio")
    import importlib
    import app
    monkeypatch.setattr(app, "VALIDATION", ValidationSet(str(tmp_path)))
    assert app.build_demo() is not None
    name = "K7 F 35D SC Iba W001 - RGB.png"
    Image.fromarray(np.zeros((60, 50, 3), np.uint8)).save(tmp_path / name)
    with zipfile.ZipFile(tmp_path / "ann.zip", "w") as zf:
        zf.writestr("annotations.xml", XML.format(name=name, w=50, h=60))
    monkeypatch.setattr(app, "VALIDATION", ValidationSet(str(tmp_path)))
    demo = app.build_demo()
    assert demo is not None
    with pytest.raises(Exception):
        app.validate_image("nope")
    assert app.COMPARISON_NOTE.startswith("**Matched**")
    if app.MODEL is not None:  # trained weights present: the whole section-2 path on the one-image bundle
        monkeypatch.delenv("GEMINI_API_KEY", raising=False)
        summary, manual, predicted, missed, rows, diffs, zp = app.validate_image("K7")
        assert summary.startswith("### K7: 1 annotated nuclei") and "Run time" in summary and "Gemini summary is off" in summary
        assert manual.label.startswith("1. Manual annotation of K7: 1 nuclei, 1 puncta")
        v = manual.value
        assert os.path.exists(v["path"] if isinstance(v, dict) else v)   # gradio wraps file values in a dict
        assert predicted.label.startswith("2. Model prediction") and missed.label.startswith("3. Missed only")
        assert rows[0][0] == "Cell, Antibody +" and len(rows[0]) == len(COMPARISON_HEADERS) and os.path.exists(zp)
    importlib.reload(app)


def test_gemini_prompts_and_disabled_paths(monkeypatch):
    from cellquant import gemini_summary as g
    monkeypatch.delenv("GEMINI_API_KEY", raising=False)
    assert not g.enabled()
    rows = [["All nuclei", 178, 185, 7, 169, 2, 7, 14], ["All puncta", 20, 20, 0, 16, 0, 4, 4]]
    txt = g._comparison_prompt("C100", rows, COMPARISON_HEADERS)
    assert "C100" in txt and "All nuclei, Annotated = 178, Predicted = 185" in txt.replace("Label = ", "") and "Missed = 7" in txt
    assert "do not round" in txt and "training set" in txt
    assert g.summarize_comparison("C100", rows, COMPARISON_HEADERS).startswith("Gemini summary disabled")
    assert g.summarize([{"image": "a", "counts": {"Cell, Antibody +": 1}}]).startswith("Gemini summary disabled")
    assert g._retryable(Exception("404 NOT_FOUND. model x")) and g._retryable(Exception("429 RESOURCE_EXHAUSTED. {'error': {}}"))
    assert not g._retryable(Exception("400 INVALID_ARGUMENT. bad key"))
    assert g._short_error("ClientError: 429 RESOURCE_EXHAUSTED. {'error': {'code': 429}}") == "ClientError: 429 RESOURCE_EXHAUSTED"
    cap = ("ClientError: 429 RESOURCE_EXHAUSTED. {'error': {'code': 429, 'message': 'Your billing account has exceeded its "
           "monthly spending cap. Please go to AI Studio. Learn more at https://x', 'status': 'RESOURCE_EXHAUSTED'}}")
    assert g._short_error(cap) == "ClientError: 429 RESOURCE_EXHAUSTED: Your billing account has exceeded its monthly spending cap. Please go to AI Studio."