from cellquant.labels import (ALL_LABELS, CELL_LABELS, COUNT_COLUMNS, NOSTRIL_LABELS, TOTAL_NUCLEI, TOTAL_PUNCTA, cell_attributes,
                              cell_label, count_row, nostril_in_nucleus, nostril_label)


def test_cell_label_roundtrip():
    for lab in CELL_LABELS:
        a, n = cell_attributes(lab)
        assert cell_label(a, n) == lab


def test_nostril_label_roundtrip():
    for lab in NOSTRIL_LABELS:
        assert nostril_label(nostril_in_nucleus(lab)) == lab


def test_label_sets_disjoint():
    assert len(ALL_LABELS) == 6 and set(CELL_LABELS).isdisjoint(NOSTRIL_LABELS)


def test_count_columns_order():
    # each group of labels is followed by its total, nuclei first
    assert COUNT_COLUMNS == CELL_LABELS + (TOTAL_NUCLEI,) + NOSTRIL_LABELS + (TOTAL_PUNCTA,)
    assert len(COUNT_COLUMNS) == 8 and COUNT_COLUMNS.index(TOTAL_NUCLEI) == 4 and COUNT_COLUMNS.index(TOTAL_PUNCTA) == 7
    assert [c for c in COUNT_COLUMNS if c not in (TOTAL_NUCLEI, TOTAL_PUNCTA)] == list(ALL_LABELS)


def test_count_row_sums_each_group():
    counts = {"Cell, Antibody +": 3, "Cell, Antibody -": 5, "Cell, Nostril, Antibody +": 7, "Cell, Nostril, Antibody -": 11,
              "Nostril in Nucleus": 13, "Nostril External of Nucleus": 17}
    row = count_row(counts)
    assert row == [3, 5, 7, 11, 26, 13, 17, 30]
    assert len(row) == len(COUNT_COLUMNS)
    assert row[COUNT_COLUMNS.index(TOTAL_NUCLEI)] == sum(counts[lab] for lab in CELL_LABELS)
    assert row[COUNT_COLUMNS.index(TOTAL_PUNCTA)] == sum(counts[lab] for lab in NOSTRIL_LABELS)
    assert all(isinstance(v, int) for v in row)


def test_count_row_missing_labels_count_as_zero():
    assert count_row({}) == [0] * 8
    row = count_row({"Cell, Antibody -": 2, "Nostril External of Nucleus": 4, "unknown label": 99})
    assert row == [0, 2, 0, 0, 2, 0, 4, 4]
    assert count_row({"Cell, Antibody +": 1.0}) == [1, 0, 0, 0, 1, 0, 0, 0]   # numpy / float counts become ints