import numpy as np
from leonida.classify import CLASSES, PALETTE, class_shares, classify, drop_map_labels

def block(color, size=8):
    return np.full((size, size, 3), color, dtype=np.uint8)

def test_pure_colors_hit_their_class():
    for name, color in PALETTE.items():
        labels = classify(block(color))
        assert (labels == CLASSES[name]).all(), name

def test_jpeg_noise_still_classifies():
    rng = np.random.default_rng(0)
    noisy = block(PALETTE["strasse"], 32).astype(np.int16)
    noisy += rng.integers(-10, 11, noisy.shape)          # JPEG-Rauschen
    labels = classify(np.clip(noisy, 0, 255).astype(np.uint8))
    assert (labels == CLASSES["strasse"]).mean() > 0.98

def test_unknown_pixels_are_filled_from_neighbours():
    img = block(PALETTE["stadt"], 16)
    img[7:9, 7:9] = (255, 0, 255)                         # Beschriftung, in keiner Klasse
    labels = classify(img)
    assert (labels == CLASSES["stadt"]).all()

def test_reports_the_share_of_unknown_pixels():
    img = block(PALETTE["stadt"], 16)
    img[7:9, 7:9] = (255, 0, 255)              # 4 von 256 Pixeln passen zu keiner Farbe
    labels, share = classify(img, report_unknown=True)
    assert (labels == CLASSES["stadt"]).all()
    assert abs(share - 4 / 256) < 1e-9

def test_shares_sum_to_one():
    img = np.concatenate([block(PALETTE["wasser"], 8), block(PALETTE["land"], 8)], axis=1)
    shares = class_shares(classify(img))
    assert abs(sum(shares.values()) - 1.0) < 1e-9
    assert abs(shares["wasser"] - 0.5) < 1e-9

def test_letter_shaped_island_becomes_water():
    """Ein Buchstabe im Wasser ist Beschriftung, keine Stadt: keine natuerlichen Merkmale."""
    labels = np.full((200, 200), CLASSES["wasser"], dtype=np.uint8)
    labels[40:60, 40:120] = CLASSES["stadt"]          # Querbalken
    labels[40:120, 40:60] = CLASSES["strasse"]        # Laengsbalken -> zusammen ein "L"
    out = drop_map_labels(labels)
    assert (out == CLASSES["wasser"]).all()

def test_real_island_survives():
    """Eine Insel mit Gruen und Sand bleibt, auch wenn sie klein und vom Festland getrennt ist."""
    labels = np.full((200, 200), CLASSES["wasser"], dtype=np.uint8)
    labels[40:120, 40:120] = CLASSES["stadt"]
    labels[50:70, 50:70] = CLASSES["gruen"]
    labels[100:115, 45:115] = CLASSES["sand"]
    out = drop_map_labels(labels)
    assert (out[40:120, 40:120] != CLASSES["wasser"]).all()

def test_mainland_is_never_dropped():
    """Die groesste Nicht-Wasser-Komponente bleibt immer, auch voellig ohne Gruen und Sand."""
    labels = np.full((200, 200), CLASSES["wasser"], dtype=np.uint8)
    labels[10:190, 10:100] = CLASSES["stadt"]
    out = drop_map_labels(labels)
    assert (out[10:190, 10:100] == CLASSES["stadt"]).all()