#!/usr/bin/env python3
"""Shared SUM Parts class table and a memory-safe colour decoder.

main.py writes predictions as palette colours rather than labels, so recovering
the class means matching RGB back to the palette. The obvious way to do that,

    d = ((rgb[:, None, :] - COLOR_MAP[None, :, :]) ** 2).sum(axis=2)

builds an (N, 13, 3) float64 temporary. At 800k points that is ~250 MB per
intermediate and several exist at once - enough to run a machine out of memory
on a full tile, and the failure looks like an unrelated crash.

Since main.py writes exact palette entries, an exact lookup handles nearly
every point in one pass with no large temporary. Only the leftovers fall back
to a distance search, and that runs in chunks.
"""
from __future__ import annotations

import numpy as np

CLASSES = ['unclassified', 'terrain', 'high_vegetation', 'facade_surface',
           'water', 'car', 'boat', 'roof_surface', 'chimney', 'dormer',
           'balcony', 'roof_installation', 'wall']

COLOR_MAP = np.array([
    (0, 0, 0), (170, 85, 0), (0, 255, 0), (255, 255, 0),
    (0, 255, 255), (255, 0, 255), (0, 0, 153), (85, 85, 127),
    (255, 50, 50), (85, 0, 127), (50, 125, 150), (50, 0, 50),
    (215, 160, 140),
], dtype=np.float64)

# what the project actually needs, folded from the 13
#
# water sits in `unseen`, not `ground`: the Seosan site has essentially no
# water, so every point called water is a misread. Filing it as ground would
# bake that error into the terrain surface.
GROUPS: dict[str, list[int]] = {
    "ground":   [1],
    "building": [3, 7, 8, 9, 10, 11, 12],
    "tree":     [2],
    "vehicle":  [5, 6],
    "unseen":   [0, 4],
}


def _pack(rgb_u8: np.ndarray) -> np.ndarray:
    """RGB triples -> one int32 key each, for hashing."""
    a = rgb_u8.astype(np.int32)
    return (a[:, 0] << 16) | (a[:, 1] << 8) | a[:, 2]


def decode_palette(rgb: np.ndarray, chunk: int = 200_000) -> np.ndarray:
    """Class index per point, without allocating an (N, 13, 3) temporary.

    Exact palette hits resolve through a dict; anything else (a resampled or
    recompressed file) falls back to a chunked nearest-colour search.
    """
    rgb = np.asarray(rgb)
    if rgb.dtype != np.uint8:
        scaled = rgb.astype(np.float64)
        if scaled.max() <= 1.0:
            scaled = scaled * 255.0
        rgb = np.clip(scaled, 0, 255).astype(np.uint8)

    lut = {int(k): i for i, k in enumerate(_pack(COLOR_MAP.astype(np.uint8)))}
    keys = _pack(rgb)

    out = np.full(len(rgb), -1, dtype=np.int64)
    uniq, inverse = np.unique(keys, return_inverse=True)
    resolved = np.array([lut.get(int(k), -1) for k in uniq], dtype=np.int64)
    out = resolved[inverse]

    missing = out < 0
    if missing.any():
        idx = np.flatnonzero(missing)
        for s in range(0, len(idx), chunk):
            part = idx[s:s + chunk]
            block = rgb[part].astype(np.float64)
            d = ((block[:, None, :] - COLOR_MAP[None, :, :]) ** 2).sum(axis=2)
            out[part] = d.argmin(axis=1)

    return out


def read_point_classes(path) -> tuple[np.ndarray, np.ndarray]:
    """(xyz float32, class index) from a prediction or ground-truth ply."""
    from plyfile import PlyData

    v = PlyData.read(str(path))["vertex"]
    props = [p.name for p in v.properties]

    xyz = np.stack([np.asarray(v["x"], dtype=np.float32),
                    np.asarray(v["y"], dtype=np.float32),
                    np.asarray(v["z"], dtype=np.float32)], axis=1)

    if "label" in props:
        lab = np.asarray(v["label"]).astype(np.int64)
        # a placeholder label (all -1, or a single constant on an unlabelled
        # tile) carries no information; fall through to the colours
        if lab.min() >= 0 and len(np.unique(lab)) > 1:
            return xyz, lab
        if lab.min() >= 0 and "red" not in props and "r" not in props:
            return xyz, lab

    rgb_set = next((s for s in (("red", "green", "blue"), ("r", "g", "b"))
                    if all(c in props for c in s)), None)
    if rgb_set is None:
        raise SystemExit(f"{path}: no usable label field and no colour to decode")

    rgb = np.stack([np.asarray(v[c]) for c in rgb_set], axis=1)
    return xyz, decode_palette(rgb)


def read_colours(path, n: int) -> np.ndarray | None:
    """uint8 RGB from a ply, or None if it does not line up point-for-point."""
    from plyfile import PlyData

    v = PlyData.read(str(path))["vertex"]
    props = [p.name for p in v.properties]
    rgb_set = next((s for s in (("red", "green", "blue"), ("r", "g", "b"))
                    if all(c in props for c in s)), None)
    if rgb_set is None or len(v) != n:
        return None

    rgb = np.stack([np.asarray(v[c]) for c in rgb_set], axis=1)
    if rgb.dtype != np.uint8:
        rgb = rgb.astype(np.float64)
        if rgb.max() <= 1.0:
            rgb = rgb * 255.0
        rgb = np.clip(rgb, 0, 255).astype(np.uint8)
    return rgb