#!/usr/bin/env python3
"""Split a predicted point cloud into one OBJ per target class.

The prediction carries one of SUM's 13 fine classes per point. This folds them
into the four the project cares about, plus a bucket for everything the model
could not place:

    ground    terrain
    building  facade_surface, roof_surface, chimney, dormer, balcony,
              roof_installation, wall
    tree      high_vegetation
    vehicle   car, boat
    unseen    unclassified, water

water sits in `unseen` deliberately. The Seosan site has essentially no water,
so every point the model calls water is a misread, not a class we can trust.
Filing it as ground would bake that error into the terrain; filing it as its
own bucket keeps it visible.

OBJ carries no per-point class, so each file is written as vertex-only geometry
(`v x y z r g b`) with one file per class. Viewers that read vertex colour show
the photo texture; the rest show the points.

Usage:
    python split_by_class.py pred.ply OUTDIR [--source rgb.ply]
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent))
from sumparts_palette import (  # noqa: E402
    CLASSES as FINE, GROUPS, read_colours, read_point_classes,
)

def write_obj(path: Path, xyz: np.ndarray, rgb: np.ndarray | None) -> None:
    """OBJ, vertex-only.

    Kept because it was asked for, but be aware: an OBJ with no `f` lines is a
    mesh with zero faces, and most viewers render exactly that - nothing. Use
    the .ply next to it to actually look at the points.
    """
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        f.write(f"# {path.stem}: {len(xyz):,} points\n")
        f.write("# vertex-only (no faces) -- most viewers show nothing.\n")
        f.write("# open the .ply beside this file instead.\n")
        if rgb is None:
            for p in xyz:
                f.write(f"v {p[0]:.4f} {p[1]:.4f} {p[2]:.4f}\n")
        else:
            for p, c in zip(xyz, rgb):
                f.write(f"v {p[0]:.4f} {p[1]:.4f} {p[2]:.4f} "
                        f"{c[0]/255:.4f} {c[1]/255:.4f} {c[2]/255:.4f}\n")


def write_ply(path: Path, xyz: np.ndarray, rgb: np.ndarray | None) -> None:
    """Binary PLY - the format point-cloud viewers actually open.

    float32 coordinates and uint8 red/green/blue, which is the spelling
    CloudCompare, MeshLab and Mapple all read without coaxing.
    """
    from plyfile import PlyData, PlyElement

    dtype = [("x", "f4"), ("y", "f4"), ("z", "f4")]
    if rgb is not None:
        dtype += [("red", "u1"), ("green", "u1"), ("blue", "u1")]

    arr = np.empty(len(xyz), dtype=dtype)
    arr["x"], arr["y"], arr["z"] = xyz[:, 0], xyz[:, 1], xyz[:, 2]
    if rgb is not None:
        arr["red"], arr["green"], arr["blue"] = rgb[:, 0], rgb[:, 1], rgb[:, 2]

    path.parent.mkdir(parents=True, exist_ok=True)
    PlyData([PlyElement.describe(arr, "vertex")], text=False).write(str(path))


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("pred", type=Path, help="prediction ply")
    ap.add_argument("outdir", type=Path, help="root directory for the class folders")
    ap.add_argument("--source", type=Path, default=None,
                    help="pre-inference ply, to carry photo colour into the OBJs")
    ap.add_argument("--name", default="seosan_BlockYBA_tile0",
                    help="base filename inside each class folder")
    args = ap.parse_args()

    if not args.pred.exists():
        raise SystemExit(f"{args.pred}: not found")

    xyz, cls = read_point_classes(args.pred)
    print(f"prediction : {args.pred.name}  {len(xyz):,} points")

    rgb = None
    if args.source and args.source.exists():
        rgb = read_colours(args.source, len(xyz))
        print(f"colour     : {'from ' + args.source.name if rgb is not None else 'unavailable (count mismatch)'}")

    print()
    print(f"  {'folder':<10} {'points':>10} {'share':>8}  from")
    total = 0
    for group, ids in GROUPS.items():
        mask = np.isin(cls, ids)
        n = int(mask.sum())
        total += n
        members = ", ".join(FINE[i] for i in ids)
        sub_xyz = xyz[mask]
        sub_rgb = rgb[mask] if rgb is not None else None
        write_ply(args.outdir / group / f"{args.name}.ply", sub_xyz, sub_rgb)
        write_obj(args.outdir / group / f"{args.name}.obj", sub_xyz, sub_rgb)
        print(f"  {group:<10} {n:>10,} {100*n/len(xyz):>7.2f}%  {members}")

    print()
    print(f"  total      {total:>10,}  ({'all points accounted for' if total == len(xyz) else 'MISMATCH'})")
    print(f"  written to {args.outdir}")


if __name__ == "__main__":
    main()