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()}