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
# Groesser als jeder gemessene Buchstabe (4.876-14.760 px am realen Mosaik) und weit kleiner
# als die Hauptlandmasse (5.852.263 px): ohne diese Untergrenze waere in einer Szene ohne
# echte Hauptlandmasse die groesste - und einzige - Komponente automatisch "Festland" und
# damit nie loeschbar, auch wenn sie selbst ein Schriftzug ist (siehe
# test_letter_shaped_island_becomes_water, wo genau das die einzige Komponente ist).
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 ueberhaupt 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