import numpy as np

from cellquant.metrics import count_errors, evaluate_cells, evaluate_nostrils, match_by_centroid, match_by_iou, micro_average


def _two_squares():
    inst = np.zeros((50, 50), np.int32)
    inst[5:15, 5:15] = 1
    inst[30:45, 30:45] = 2
    return inst


def test_perfect_prediction_scores_one():
    inst = _two_squares()
    labels = ["", "Cell, Antibody +", "Cell, Antibody -"]
    r = evaluate_cells(inst, labels, inst.copy(), list(labels))
    assert r["detection"]["f1"] == 1.0 and r["attribute_accuracy"]["both"] == 1.0
    assert r["per_class"]["Cell, Antibody +"]["f1"] == 1.0


def test_wrong_label_counts_as_fp_and_fn_for_classes_but_detected():
    inst = _two_squares()
    gt = ["", "Cell, Antibody +", "Cell, Antibody -"]
    pred = ["", "Cell, Antibody +", "Cell, Antibody +"]
    r = evaluate_cells(inst, gt, inst.copy(), pred)
    assert r["detection"]["tp"] == 2
    assert r["per_class"]["Cell, Antibody -"]["fn"] == 1 and r["per_class"]["Cell, Antibody +"]["fp"] == 1
    assert r["attribute_accuracy"]["antibody"] == 0.5


def test_shifted_prediction_fails_iou_but_passes_centroid():
    gt = _two_squares()
    pred = np.zeros_like(gt)
    pred[9:19, 9:19] = 1          # IoU with square 1 = 36/164 < 0.5, centroid (14,14) is inside GT square 1
    pred[30:45, 30:45] = 2
    assert len(match_by_iou(gt, pred, 0.5)) == 1
    r = evaluate_cells(gt, ["", "Cell, Antibody +", "Cell, Antibody +"], pred, ["", "Cell, Antibody +", "Cell, Antibody +"])
    assert r["detection"]["tp"] == 1 and r["detection"]["centroid"]["tp"] == 2


def test_nostril_matching_by_distance():
    gt = np.zeros((40, 40), np.int32)
    gt[10:13, 10:13] = 1
    gt[30:33, 30:33] = 2
    pred = np.zeros_like(gt)
    pred[12:15, 12:15] = 1       # 2 px away -> match
    pred[5:8, 30:33] = 2         # 25 px away -> no match
    m = match_by_centroid(gt, pred, 6.0)
    assert len(m) == 1 and m[0][:2] == (1, 1)
    r = evaluate_nostrils(gt, ["", "Nostril in Nucleus", "Nostril External of Nucleus"], pred, ["", "Nostril in Nucleus", "Nostril in Nucleus"])
    assert r["detection"]["tp"] == 1 and r["detection"]["fp"] == 1 and r["detection"]["fn"] == 1
    assert r["type_accuracy"] == 1.0


def test_count_errors_and_micro_average():
    ce = count_errors({"Cell, Antibody +": 10}, {"Cell, Antibody +": 8})
    assert ce["Cell, Antibody +"]["abs_err"] == 2 and abs(ce["Cell, Antibody +"]["rel_err"] - 0.2) < 1e-9
    per = [{"A": {"tp": 1, "fp": 1, "fn": 0, "gt_count": 1, "pred_count": 2}}, {"A": {"tp": 3, "fp": 0, "fn": 1, "gt_count": 4, "pred_count": 3}}]
    m = micro_average(per, ["A"])
    assert m["A"]["tp"] == 4 and m["A"]["fp"] == 1 and m["A"]["fn"] == 1 and abs(m["A"]["f1"] - 0.8) < 1e-9