import numpy as np
from scipy import ndimage as ndi
from skimage.color import rgb2lab

PALETTE = {
    "wasser":     (48, 144, 192),   # gemessen im Mosaik; tiles.py nennt die Kompositionsfarbe
    "land":       (192, 216, 128),
    "gruen":      (139, 195, 74),
    "sand":       (232, 216, 136),
    "stadt":      (216, 216, 216),
    "gebaeude":   (176, 176, 176),
    "strasse":    (80, 80, 80),
    "weg":        (192, 152, 96),
    "spekulativ": (198, 96, 96),
}
CLASSES = {name: i for i, name in enumerate(PALETTE)}
UNKNOWN = 255
MAX_DELTA_E = 22.0     # darüber unbekannt; deckt den Tiefenverlauf des Wassers mit ab


def classify(rgb: np.ndarray, report_unknown: bool = False):
    """Ordnet jedes Pixel der nächsten Palettenfarbe zu; Ausreißer werden aufgefüllt.

    Mit report_unknown=True kommt zusätzlich der Anteil der Pixel zurück, die vor dem
    Auffüllen zu keiner Legendenfarbe passten. Aufgabe 6 bricht darüber bei über 5 % ab.
    """
    lab = rgb2lab(rgb.astype(np.float32) / 255.0)
    ref = rgb2lab(np.array([list(PALETTE.values())], dtype=np.float32) / 255.0)[0]
    dist = np.linalg.norm(lab[:, :, None, :] - ref[None, None, :, :], axis=3)
    labels = dist.argmin(axis=2).astype(np.uint8)
    unknown = dist.min(axis=2) > MAX_DELTA_E
    share = float(unknown.mean())
    if unknown.any():
        # Nächstes bekanntes Pixel übernehmen.
        idx = ndi.distance_transform_edt(unknown, return_distances=False, return_indices=True)
        labels = labels[tuple(idx)]
    return (labels, share) if report_unknown else labels


def class_shares(labels: np.ndarray) -> dict[str, float]:
    total = labels.size
    counts = np.bincount(labels.ravel(), minlength=len(PALETTE))
    return {name: float(counts[i]) / total for name, i in CLASSES.items()}


NATURAL_CLASSES = ("gruen", "sand", "land")
LABEL_MAX_NATURAL = 0.02     # Anteil natuerlicher Merkmale, ab dem eine Flaeche als echt gilt
# Ab dieser Groesse zaehlt die groesste Komponente als Festland; die Ausnahme gilt nur ihr,
# ein Buchstabe kann sie also nur erben, wenn die Szene sonst gar kein Land enthaelt. Jeder
# Wert zwischen dem groessten gemessenen Buchstaben (14.760 px) und der echten Hauptlandmasse
# (5.852.263 px) aendert am realen Mosaik nichts; 15.000 liegt in dieser Spanne und bleibt
# zugleich unter den 16.200 px aus test_mainland_is_never_dropped.
MAINLAND_MIN_SIZE = 15_000


def drop_map_labels(labels: np.ndarray) -> np.ndarray:
    """Beschriftungen der Originalkarte auf Wasser zuruecksetzen.

    Die Schriftzuege der Karte werden pixelweise als stadt/strasse/gebaeude klassifiziert und
    sind farblich von echten Flaechen nicht zu unterscheiden. Trennbar sind sie erst als
    Region: ein Schriftzug im Wasser bildet eine eigene, vom Festland getrennte Nicht-Wasser-
    Komponente ohne jedes natuerliche Merkmal, waehrend eine echte Insel Gruen oder Sand
    traegt. Die groesste Komponente ist das Festland und bleibt unangetastet, sofern sie
    ueberhaupt gross genug ist, um als Festland infrage zu kommen (MAINLAND_MIN_SIZE); sonst
    wird auch sie nur ueber den Anteil natuerlicher Merkmale beurteilt.
    """
    solid = labels != CLASSES["wasser"]
    if not solid.any():
        return labels
    components, count = ndi.label(solid, structure=np.ones((3, 3)))
    sizes = ndi.sum(solid, components, range(1, count + 1))
    largest = int(np.argmax(sizes)) + 1
    mainland = largest if sizes[largest - 1] >= MAINLAND_MIN_SIZE else 0
    natural = np.zeros(labels.shape, dtype=bool)
    for name in NATURAL_CLASSES:
        natural |= labels == CLASSES[name]
    shares = ndi.mean(natural.astype(np.float32), components, range(1, count + 1))
    drop = np.zeros(count + 1, dtype=bool)
    for index in range(1, count + 1):
        drop[index] = index != mainland and shares[index - 1] <= LABEL_MAX_NATURAL
    out = labels.copy()
    out[drop[components]] = CLASSES["wasser"]
    return out