import numpy as np

from cellquant.labels import ALL_LABELS
from cellquant.postprocess import PostConfig, analyze, logits_to_probs, render_overlay, segment_cells


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


def _probs(h=96, w=96):
    """Two touching discs; the left one antibody+ with a punctum inside; one external punctum."""
    d1, d2 = _disc(h, w, 48, 30, 14), _disc(h, w, 48, 56, 14)
    fg = d1 | d2
    inner = _disc(h, w, 48, 30, 11) | _disc(h, w, 48, 56, 11)
    cell = np.zeros((3, h, w), np.float32)
    cell[1][inner] = 0.9
    cell[2][fg & ~inner] = 0.9
    cell[0] = 1 - cell[1] - cell[2]
    antibody = np.where(d1, 0.95, 0.05).astype(np.float32)
    cellnos = np.full((h, w), 0.1, np.float32)
    nos = np.zeros((3, h, w), np.float32)
    p_in = _disc(h, w, 48, 30, 2)
    p_ext = _disc(h, w, 10, 80, 2)
    nos[1][p_in] = 0.9
    nos[2][p_ext] = 0.9
    nos[0] = 1 - nos[1] - nos[2]
    return {"cell": cell, "antibody": antibody, "cellnos": cellnos, "nos": nos}


def test_touching_discs_split_into_two_instances():
    inst = segment_cells(_probs()["cell"], PostConfig())
    assert inst.max() == 2
    areas = np.bincount(inst.ravel())[1:]
    assert all(500 <= a <= 700 for a in areas)  # pi*14^2 ~ 616


def test_analyze_labels_and_counts():
    res = analyze(_probs(), PostConfig())
    assert len(res.cells) == 2 and len(res.nostrils) == 2
    labels = sorted(c["label"] for c in res.cells)
    assert labels == ["Cell, Antibody -", "Cell, Nostril, Antibody +"]
    assert sorted(n["label"] for n in res.nostrils) == ["Nostril External of Nucleus", "Nostril in Nucleus"]
    assert sum(res.counts.values()) == 4 and set(res.counts) == set(ALL_LABELS)
    inside = [n for n in res.nostrils if n["label"] == "Nostril in Nucleus"][0]
    assert inside["cell_id"] == [c for c in res.cells if c["label"].startswith("Cell, Nostril")][0]["id"]


def test_cellnos_rules():
    p = _probs()
    head_only = analyze(p, PostConfig(cellnos_rule="head"))
    assert all(not c["label"].startswith("Cell, Nostril") for c in head_only.cells)
    contains = analyze(p, PostConfig(cellnos_rule="contains"))
    assert sum(c["label"].startswith("Cell, Nostril") for c in contains.cells) == 1


def test_geometric_nostril_rule_overrides_head():
    p = _probs()
    # make the network call the inside punctum "external"; geometric rule should still say inside
    p["nos"][1], p["nos"][2] = p["nos"][2].copy(), p["nos"][1].copy()
    head = analyze(p, PostConfig(nostril_rule="head"))
    geo = analyze(p, PostConfig(nostril_rule="geometric"))
    assert head.counts["Nostril in Nucleus"] == 1 and geo.counts["Nostril in Nucleus"] == 1
    # with head rule the punctum inside the disc is now labelled external
    inside_disc = [n for n in head.nostrils if n["cell_id"] > 0][0]
    assert inside_disc["label"] == "Nostril External of Nucleus"
    assert [n for n in geo.nostrils if n["cell_id"] > 0][0]["label"] == "Nostril in Nucleus"


def test_min_area_filters_small_instances():
    p = _probs()
    assert len(analyze(p, PostConfig(min_cell_area=1000)).cells) == 0


def test_logits_to_probs_and_overlay():
    logits = np.random.default_rng(0).normal(size=(8, 40, 40)).astype(np.float32)
    pr = logits_to_probs(logits)
    assert np.allclose(pr["cell"].sum(0), 1, atol=1e-5) and np.allclose(pr["nos"].sum(0), 1, atol=1e-5)
    assert pr["antibody"].min() >= 0 and pr["antibody"].max() <= 1
    res = analyze(_probs(), PostConfig())
    img = np.zeros((96, 96, 3), np.uint8)
    ov = render_overlay(img, res)
    assert ov.shape == img.shape and ov.any()


def test_empty_image_gives_no_objects():
    h = w = 64
    p = {"cell": np.stack([np.ones((h, w)), np.zeros((h, w)), np.zeros((h, w))]).astype(np.float32),
         "antibody": np.zeros((h, w), np.float32), "cellnos": np.zeros((h, w), np.float32),
         "nos": np.stack([np.ones((h, w)), np.zeros((h, w)), np.zeros((h, w))]).astype(np.float32)}
    res = analyze(p, PostConfig())
    assert res.cells == [] and res.nostrils == [] and sum(res.counts.values()) == 0