#!/usr/bin/env python3
"""Dump what a prediction PLY actually contains.

check_pred.py recovers the class by matching RGB against the palette. If that
reports one class for every point, the question is whether the model really
collapsed or whether the colour decode is wrong — so read the raw fields
instead of interpreting them.

Usage:
    python dump_pred.py pred.ply
"""
from __future__ import annotations

import sys
from pathlib import Path

import numpy as np
from plyfile import PlyData

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.),
])


def main() -> None:
    if len(sys.argv) != 2:
        print(__doc__)
        raise SystemExit(2)

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

    print(f"file  : {path.name}")
    print(f"points: {len(v):,}")
    print(f"props : {props}")
    print()

    if "label" in props:
        lab = np.asarray(v["label"])
        u, c = np.unique(lab, return_counts=True)
        print(f"label field: dtype={lab.dtype}")
        for k, n in sorted(zip(u.tolist(), c.tolist()), key=lambda t: -t[1]):
            name = CLASSES[k] if 0 <= k < len(CLASSES) else f"?{k}"
            print(f"  {k:>3} {name:<20} {n:>10,}  {100*n/len(lab):>6.2f}%")
        print()

    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:
        rgb = np.stack([np.asarray(v[c]) for c in rgb_set], axis=1)
        uniq, cnt = np.unique(rgb.reshape(-1, 3), axis=0, return_counts=True)
        print(f"colour field {rgb_set}: dtype={rgb.dtype}, {len(uniq)} distinct")
        order = np.argsort(-cnt)
        for i in order[:15]:
            col = uniq[i]
            d = ((COLOR_MAP - col.astype(np.float64)) ** 2).sum(axis=1)
            k = int(d.argmin())
            exact = "exact" if d[k] < 1 else f"nearest (dist {np.sqrt(d[k]):.0f})"
            name = CLASSES[k] if k < len(CLASSES) else f"?{k}"
            print(f"  {tuple(int(x) for x in col)!s:<20} {cnt[i]:>10,} "
                  f"{100*cnt[i]/len(rgb):>6.2f}%  -> {name} ({exact})")

    print()
    if "label" in props and rgb_set:
        lab = np.asarray(v["label"])
        rgb = np.stack([np.asarray(v[c]) for c in rgb_set], axis=1)
        d = ((rgb[:, None, :].astype(np.float64) - COLOR_MAP[None, :, :]) ** 2).sum(axis=2)
        from_colour = d.argmin(axis=1)
        agree = int((from_colour == lab).sum())
        print(f"label vs colour agreement: {agree:,}/{len(lab):,} "
              f"({100*agree/len(lab):.2f}%)")
        if agree < len(lab):
            print("  -> the two disagree; the label field is authoritative")


if __name__ == "__main__":
    main()