"""Shared lossless mask taxonomy for the painting vision dataset."""
import numpy as np

CLASSES = ("other", "wall_unpainted", "wall_painted", "wall_uncertain", "skirting",
           "switch_outlet", "ac_unit", "door", "window", "wall_obstacle")
IGNORE = 255
UNKNOWN = 254

# Visible surface material of wall pixels, predicted by the material head
# independently of paint state. Wallpaper is the visible covering, so a
# wallpapered gypsum wall is ``wallpaper``, not ``drywall``. Older datasets that
# only use 0/1 stay valid; 255 marks unknown material or non-wall pixels.
MATERIAL_CLASSES = ("other_wall_material", "drywall", "wallpaper")
LEGACY_MATERIAL_CLASSES = MATERIAL_CLASSES[:2]
# A smooth finished coat over board or plaster: the material is genuinely one of
# the two and the photo cannot tell which. Trained toward 50/50 drywall/other, so
# the model's honest answer there is low confidence (predict.py: "unknown").
MATERIAL_DRYWALL_OR_OTHER = 254
MATERIAL_VALUES = tuple(range(len(MATERIAL_CLASSES))) + (MATERIAL_DRYWALL_OR_OTHER, IGNORE)

# Surface condition of decided wall pixels, from a third head: what has to be
# prepared before paint goes on. ``paint_failure`` covers peeling, flaking,
# blistering and cracked coatings; ``rust`` covers corrosion and rust staining.
# Masks live in condition_masks/<split>/<stem>.png; 255 = unknown or not wall.
# Datasets without condition masks are simply unsupervised for this head.
CONDITION_CLASSES = ("sound", "paint_failure", "rust")
CONDITION_VALUES = tuple(range(len(CONDITION_CLASSES))) + (IGNORE,)

# Detection classes for the equipment around the wall: what has to be on site to
# do the job, and what must not be painted over. These are *objects*, not pixel
# labels - they never appear in the ten-class mask, so ``CLASSES`` above (and any
# trained checkpoint) is unaffected by adding or reordering them. ``site_objects.py``
# detects them without a model, from shape and colour cues.
OBJECT_CLASSES = ("paint_can", "paint_tray", "brush", "roller", "ladder", "drop_sheet")

# RGB values for lossless color-coded PNG masks.
RGB_PALETTE = {
    (0, 0, 0): 0,
    (0, 0, 255): 1,
    (0, 255, 0): 2,
    (255, 255, 0): 3,
    (255, 0, 255): 4,
    (255, 0, 0): 5,
    (0, 255, 255): 6,
    (255, 128, 0): 7,
    (128, 0, 255): 8,
    (0, 128, 128): 9,
}


def mask_to_labels(mask):
    m = np.asarray(mask)
    if m.ndim == 3:
        rgb = m[..., :3]
        labels = np.full(m.shape[:2], UNKNOWN, dtype=np.uint8)
        for color, label in RGB_PALETTE.items():
            labels[np.all(rgb == color, axis=2)] = label
        labels[np.all(rgb == (255, 255, 255), axis=2)] = IGNORE
    else:
        labels = m.astype(np.uint8)
    valid = (labels == IGNORE) | np.isin(labels, np.arange(len(CLASSES)))
    if not valid.all():
        raise ValueError(f"mask contains unrecognized/invalid values {np.unique(labels[~valid]).tolist()}")
    return labels