#!/usr/bin/env python3
"""Summarise a prediction PLY written by main.py's visualization step.

main.py colours predictions through SUMV2_Triangle_COLOR_MAP rather than
writing a label field, so recover the class by matching each point's RGB back
to that palette.

Usage:
    python check_pred.py .../seosan_BlockYBA_tile0_pred.ply
"""
from __future__ import annotations

import sys
from pathlib import Path

import numpy as np
from plyfile import PlyData

# openpoints/dataset/sumv2_triangle/sumv2_triangle.py
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])
    if not path.exists():
        print(f"{path}: not found")
        raise SystemExit(1)

    v = PlyData.read(str(path))["vertex"]
    props = [p.name for p in v.properties]
    print(f"=== {path.name} ===")
    print(f"  points     : {len(v):,}")
    print(f"  properties : {props}")

    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:
        print("  no colour channels -- cannot recover predicted class")
        raise SystemExit(1)

    rgb = np.stack([np.asarray(v[c], dtype=np.float64) for c in rgb_set], axis=1)
    if rgb.max() <= 1.0:
        rgb = rgb * 255.0

    # nearest palette entry per point
    d = ((rgb[:, None, :] - COLOR_MAP[None, :, :]) ** 2).sum(axis=2)
    cls = d.argmin(axis=1)
    resid = np.sqrt(d.min(axis=1))

    print(f"  palette fit: max residual {resid.max():.1f} "
          f"({'exact' if resid.max() < 1 else 'approximate'})")
    print()
    print(f"  {'class':<20} {'points':>10}  {'share':>7}")
    u, c = np.unique(cls, return_counts=True)
    order = np.argsort(-c)
    for i in order:
        k, n = int(u[i]), int(c[i])
        name = CLASSES[k] if k < len(CLASSES) else f"?{k}"
        print(f"  {name:<20} {n:>10,}  {100 * n / len(cls):>6.2f}%")

    print()
    print(f"  distinct classes predicted: {len(u)}")
    if len(u) == 1:
        print("  note: single class everywhere -- expected from a 1-epoch model")


if __name__ == "__main__":
    main()