"""Tests for the peak-distance rejection rule (PostConfig.min_peak_dist) and the
per-image post-processing overrides of evaluate.py."""
import json

import numpy as np

from cellquant.data import distance_target
from cellquant.postprocess import PostConfig, analyze, drop_low_peaks, peak_dist, 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 _scene():
    """A bright annotated-looking nucleus (distance peak 1) and a dim blob the
    distance head only half believes in (peak 0.4); both are solid foreground."""
    h = w = 96
    d1, d2 = _disc(h, w, 30, 30, 14), _disc(h, w, 66, 66, 12)
    cell = np.zeros((3, h, w), np.float32)
    cell[1][d1 | d2] = 0.9
    cell[0] = 1 - cell[1]
    gt = np.zeros((h, w), np.int32)
    gt[d1] = 1
    gt[d2] = 2
    dist = distance_target(gt)
    dist[d2] *= 0.4
    return cell, dist, d1, d2


def test_peak_dist_and_drop_low_peaks():
    inst = np.zeros((10, 10), np.int32)
    inst[1:4, 1:4] = 1
    inst[6:9, 6:9] = 2
    dist = np.zeros((10, 10), np.float32)
    dist[2, 2] = 0.9
    dist[7, 7] = 0.3
    pk = peak_dist(inst, dist)
    assert pk.shape == (3,) and pk[0] == 0 and np.isclose(pk[1], 0.9) and np.isclose(pk[2], 0.3)
    out = drop_low_peaks(inst, dist, 0.5)
    assert set(np.unique(out)) == {0, 1} and (out[6:9, 6:9] == 0).all() and (inst[6:9, 6:9] == 2).all()  # input untouched
    assert drop_low_peaks(inst, dist, 0.0) is inst
    assert peak_dist(np.zeros((4, 4), np.int32), dist[:4, :4]).tolist() == [0.0]


def test_min_peak_dist_drops_only_low_peak_nuclei():
    cell, dist, d1, d2 = _scene()
    both = segment_cells(cell, PostConfig(min_peak_dist=0.0), dist)
    assert both.max() == 2
    assert PostConfig().min_peak_dist == 0.64  # shipped default (reports/cv_v3/rule_study.md)
    assert segment_cells(cell, PostConfig(), dist).max() == 1
    kept = segment_cells(cell, PostConfig(min_peak_dist=0.6), dist)
    assert kept.max() == 1 and (kept[d1] > 0).mean() > 0.95 and (kept[d2] == 0).all()
    # a threshold below the dim blob's peak keeps both; 0 switches the rule off
    assert segment_cells(cell, PostConfig(min_peak_dist=0.3), dist).max() == 2
    assert segment_cells(cell, PostConfig(min_peak_dist=0.0), dist).max() == 2
    # without a distance map (8-channel checkpoints) the rule is not applied
    assert segment_cells(cell, PostConfig(min_peak_dist=0.9)).max() == 2


def test_analyze_applies_rule_and_records_it():
    cell, dist, d1, d2 = _scene()
    h, w = cell.shape[1:]
    probs = {"cell": cell, "antibody": np.full((h, w), 0.9, 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), "dist": dist}
    res = analyze(probs, PostConfig(min_peak_dist=0.6))
    assert len(res.cells) == 1 and res.config["min_peak_dist"] == 0.6
    assert len(analyze(probs, PostConfig(min_peak_dist=0.0)).cells) == 2


def test_evaluate_post_config_overrides(tmp_path):
    import evaluate as ev
    assert ev.load_post_config("") == {}
    g = ev.post_config_for(ev.load_post_config('{"min_peak_dist": 0.63}'), "C100")
    assert g.min_peak_dist == 0.63 and g.min_cell_area == PostConfig().min_cell_area
    p = tmp_path / "pc.json"
    p.write_text(json.dumps({"C100": {"min_peak_dist": 0.61}, "F99": {"min_peak_dist": 0.6, "min_cell_area": 40}}))
    o = ev.load_post_config(str(p))
    assert ev.post_config_for(o, "C100").min_peak_dist == 0.61
    assert ev.post_config_for(o, "F99").min_cell_area == 40
    assert ev.post_config_for(o, "C110") == PostConfig()  # no entry -> defaults
    import pytest
    with pytest.raises(SystemExit):
        ev.post_config_for({"no_such_field": 1}, "C100")