#!/usr/bin/env python3
"""Dichter Szenen-Splat aus Equirect + DepthAnything-Tiefenkarte.

Projiziert jeden Pixel des 360°-Bildes per Kugelkoordinaten und der
(Disparitäts-)Tiefe in einen farbigen 3D-Splat. Ergebnis: eine dichte,
lückenlose Oberfläche der Szene vom Aufnahmepunkt aus.

Verbesserungen gegen Unschärfe & schwarze Löcher:
- höhere Auflösung + kleinere Splats (schärfer)
- Kantenbehandlung: an Tiefensprüngen (Türöffnungen) werden gedehnte
  "Gummiband"-Splats verworfen statt schwarz/verschmiert darzustellen
- Kappung unzuverlässiger Ferntiefen (Fenster/Himmel)

Ausgabe: antimatter15 .splat (32 Bytes/Splat), passend zum Web-Viewer.
"""
from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
from PIL import Image


def build(pano_dir: Path, out: Path, width: int, depth_scale: float,
          near_disp: float = 0.05, gap: float = 1.25,
          far_cap: float = 7.5, edge_ratio: float = 0.18) -> None:
    eq_path = pano_dir / "equirect.jpg"
    dm_path = pano_dir / "dmap.png"
    if not eq_path.is_file() or not dm_path.is_file():
        raise SystemExit(f"equirect/dmap fehlt in {pano_dir}")

    W = width
    H = W // 2
    color = Image.open(eq_path).convert("RGB").resize((W, H), Image.LANCZOS)
    depth = Image.open(dm_path).convert("L").resize((W, H), Image.BILINEAR)
    col = np.asarray(color, dtype=np.float32) / 255.0
    disp = np.asarray(depth, dtype=np.float32) / 255.0      # 0..1, hell=nah

    us = (np.arange(W) + 0.5) / W
    vs = (np.arange(H) + 0.5) / H
    lon = (us - 0.5) * 2.0 * np.pi
    lat = (0.5 - vs) * np.pi
    lon_g, lat_g = np.meshgrid(lon, lat)

    cos_lat = np.cos(lat_g)
    dir_x = cos_lat * np.sin(lon_g)
    dir_y = -np.sin(lat_g)
    dir_z = cos_lat * np.cos(lon_g)

    dist = depth_scale / (disp + near_disp)
    dist = np.clip(dist, 0.1, far_cap)

    pos = np.stack([dir_x * dist, dir_y * dist, dir_z * dist], axis=-1)

    dtheta = 2.0 * np.pi / W
    scale = (dist * dtheta * gap).astype(np.float32)

    # Kantenmaske: große relative Disparitätssprünge zu Nachbarn = Tiefenkante
    # (Türrahmen/Verdeckung). Solche Pixel erzeugen gedehnte Splats → verwerfen.
    d = disp
    grad = np.zeros_like(d)
    grad[:, 1:]  = np.maximum(grad[:, 1:],  np.abs(d[:, 1:]  - d[:, :-1]))
    grad[:, :-1] = np.maximum(grad[:, :-1], np.abs(d[:, :-1] - d[:, 1:]))
    grad[1:, :]  = np.maximum(grad[1:, :],  np.abs(d[1:, :]  - d[:-1, :]))
    grad[:-1, :] = np.maximum(grad[:-1, :], np.abs(d[:-1, :] - d[1:, :]))
    keep = (grad < edge_ratio)
    # Ganz ferne (unzuverlässige) Pixel ebenfalls verwerfen.
    keep &= (disp > 0.06)

    keep_flat = keep.reshape(-1)
    n = int(keep_flat.sum())
    idx = np.where(keep_flat)[0]

    positions = pos.reshape(-1, 3)[idx].astype(np.float32)
    scales = np.repeat(scale.reshape(-1, 1)[idx], 3, axis=1).astype(np.float32)
    rgb = col.reshape(-1, 3)[idx] * 255.0
    rgba = np.concatenate([rgb, np.full((n, 1), 255.0)], axis=1)
    rgba = np.clip(rgba, 0, 255).astype(np.uint8)
    quat = np.tile(np.array([255, 128, 128, 128], dtype=np.uint8), (n, 1))

    rec = np.zeros(n, dtype=[("pos", np.float32, 3), ("scale", np.float32, 3),
                            ("rgba", np.uint8, 4), ("quat", np.uint8, 4)])
    rec["pos"] = positions
    rec["scale"] = scales
    rec["rgba"] = rgba
    rec["quat"] = quat
    out.parent.mkdir(parents=True, exist_ok=True)
    rec.tofile(out)
    total = W * H
    print(f"{out}: {n}/{total} Splats ({n*32/1e6:.1f} MB), "
          f"{100*n/total:.0f}% behalten, Distanz {dist.min():.2f}–{dist.max():.2f}")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("pano_dir")
    ap.add_argument("out")
    ap.add_argument("--width", type=int, default=1600)
    ap.add_argument("--depth-scale", type=float, default=1.6)
    ap.add_argument("--gap", type=float, default=1.25)
    ap.add_argument("--far-cap", type=float, default=7.5)
    args = ap.parse_args()
    build(Path(args.pano_dir), Path(args.out), args.width, args.depth_scale,
          gap=args.gap, far_cap=args.far_cap)