#!/usr/bin/env python3
"""SHARP-Batch für vr.panomity.com.

Erzeugt pro Szene aus den 4 horizontalen Cube-Faces (f/r/b/l) je ein
3D-Gaussian-Splat (Apple SHARP) und merged sie über die bekannten
Yaw-Rotationen zu einem 360°-Splat (PLY, 3DGS-Format).

Usage:
  sharp_batch.py --panos 7,8,9 [--faces f,r,b,l] [--keep-face-plys]
"""
from __future__ import annotations

import argparse
import subprocess
import sys
from pathlib import Path

import numpy as np

SHARP_BIN = "/home/panomity.de/sharp/.venv/bin/sharp"
PANOS_DIR = Path("/home/panomity.de/vr.panomity.com/panos")
OUT_DIR = Path("/home/panomity.de/sharp-out")

# Yaw je Cube-Face (krpano: f=0°, r=90°, b=180°, l=270°; Drehung um Y-Achse).
# hXXX = Diagonal-Ansichten, on-the-fly aus dem Equirect gerendert, um die
# Nahtstellen zwischen den Cube-Faces zu füllen.
FACE_YAW = {
    "f": 0.0, "r": 90.0, "b": 180.0, "l": 270.0,
    "h045": 45.0, "h135": 135.0, "h225": 225.0, "h315": 315.0,
}
DEFAULT_FACES = "f,h045,r,h135,b,h225,l,h315"


def yaw_matrix(deg: float) -> np.ndarray:
    a = np.deg2rad(deg)
    c, s = np.cos(a), np.sin(a)
    # Rotation um die Y-Achse (COLMAP/OpenCV-Konvention: x rechts, y unten, z vor)
    return np.array([[c, 0.0, s], [0.0, 1.0, 0.0], [-s, 0.0, c]], dtype=np.float64)


def rotmat_to_quat(m: np.ndarray) -> np.ndarray:
    """Rotationsmatrix → Quaternion (w, x, y, z)."""
    t = np.trace(m)
    if t > 0:
        s = np.sqrt(t + 1.0) * 2
        w = 0.25 * s
        x = (m[2, 1] - m[1, 2]) / s
        y = (m[0, 2] - m[2, 0]) / s
        z = (m[1, 0] - m[0, 1]) / s
    elif m[0, 0] > m[1, 1] and m[0, 0] > m[2, 2]:
        s = np.sqrt(1.0 + m[0, 0] - m[1, 1] - m[2, 2]) * 2
        w = (m[2, 1] - m[1, 2]) / s
        x = 0.25 * s
        y = (m[0, 1] + m[1, 0]) / s
        z = (m[0, 2] + m[2, 0]) / s
    elif m[1, 1] > m[2, 2]:
        s = np.sqrt(1.0 + m[1, 1] - m[0, 0] - m[2, 2]) * 2
        w = (m[0, 2] - m[2, 0]) / s
        x = (m[0, 1] + m[1, 0]) / s
        y = 0.25 * s
        z = (m[1, 2] + m[2, 1]) / s
    else:
        s = np.sqrt(1.0 + m[2, 2] - m[0, 0] - m[1, 1]) * 2
        w = (m[1, 0] - m[0, 1]) / s
        x = (m[0, 2] + m[2, 0]) / s
        y = (m[1, 2] + m[2, 1]) / s
        z = 0.25 * s
    return np.array([w, x, y, z], dtype=np.float64)


def quat_mul(q1: np.ndarray, q2: np.ndarray) -> np.ndarray:
    w1, x1, y1, z1 = q1[..., 0], q1[..., 1], q1[..., 2], q1[..., 3]
    w2, x2, y2, z2 = q2[..., 0], q2[..., 1], q2[..., 2], q2[..., 3]
    return np.stack([
        w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2,
        w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2,
        w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2,
        w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2,
    ], axis=-1)


def read_gs_ply(path: Path):
    """Liest das vertex-Element eines binären 3DGS-PLY (SHARP hängt weitere
    Elemente wie extrinsic/intrinsic an, die hier ignoriert werden).
    Rückgabe: (names, data[N, C])."""
    with open(path, "rb") as f:
        header = b""
        while not header.endswith(b"end_header\n"):
            header += f.readline()
        lines = header.decode("ascii", errors="replace").splitlines()
        n = 0
        names = []
        current_element = None
        for line in lines:
            parts = line.split()
            if not parts:
                continue
            if parts[0] == "element":
                current_element = parts[1]
                if current_element == "vertex":
                    n = int(parts[2])
            elif parts[0] == "property" and current_element == "vertex":
                if parts[1] != "float":
                    raise ValueError(f"non-float vertex property in {path}: {line}")
                names.append(parts[2])
        data = np.fromfile(f, dtype=np.float32, count=n * len(names)).reshape(n, len(names))
    return names, data


def write_gs_ply(path: Path, names, data: np.ndarray) -> None:
    with open(path, "wb") as f:
        f.write(b"ply\nformat binary_little_endian 1.0\n")
        f.write(f"element vertex {data.shape[0]}\n".encode())
        for name in names:
            f.write(f"property float {name}\n".encode())
        f.write(b"end_header\n")
        data.astype(np.float32).tofile(f)


INPUT_SIZE = 1216  # Quell-Equirect liefert ~1200px pro Face; spart ~40% VRAM ggü. 1536


def prepare_input(src: Path, outdir: Path) -> Path:
    """Face-Bild für SHARP vorbereiten: auf INPUT_SIZE skalieren + EXIF.

    Die Cube-Faces haben 90° FOV; ohne EXIF nimmt SHARP 30mm KB-äquivalent
    (54° FOV) an — das verbiegt Wände. 90° bei 36x24-Diagonale ≙ 15,3mm:
    f_px = 15.3 * diag_px / 43.2666 = width/2 (unabhängig von der Auflösung).
    """
    import piexif
    from PIL import Image

    dst = outdir / ("exif_" + src.name)
    if dst.is_file():
        return dst
    outdir.mkdir(parents=True, exist_ok=True)
    img = Image.open(src).convert("RGB")
    if img.width != INPUT_SIZE:
        img = img.resize((INPUT_SIZE, INPUT_SIZE), Image.LANCZOS)
    img.save(dst, quality=92)
    exif_bytes = piexif.dump({"Exif": {piexif.ExifIFD.FocalLength: (153, 10)}})
    piexif.insert(exif_bytes, str(dst))
    return dst


def render_diagonal_view(tiles: Path, pano: str, face: str) -> Path | None:
    """Rendert eine 90°-FOV-Diagonalansicht aus dem Equirect (1536px)."""
    out = OUT_DIR / pano / f"view_{face}.jpg"
    if out.is_file():
        return out
    eq_path = tiles / "equirect.jpg"
    if not eq_path.is_file():
        return None
    sys.path.insert(0, str(Path(__file__).parent))
    from equirect_views import e2p
    from PIL import Image
    equi = np.asarray(Image.open(eq_path).convert("RGB"))
    img = e2p(equi, 90.0, FACE_YAW[face], 0.0, INPUT_SIZE)
    out.parent.mkdir(parents=True, exist_ok=True)
    Image.fromarray(img).save(out, quality=92)
    return out


def run_sharp(image: Path, outdir: Path, device: str = "cuda") -> Path | None:
    outdir.mkdir(parents=True, exist_ok=True)
    existing = list(outdir.glob("**/*.ply"))
    if existing:
        return existing[0]

    env = dict(**__import__("os").environ)
    env["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

    # Bei CUDA-OOM einmal kurz warten, dann CPU — die Device-Wahl pro Pano
    # (pick_device) sorgt dafür, dass eine freie GPU wieder genutzt wird.
    attempts = [device, "cpu"] if device == "cuda" else [device]
    for i, dev in enumerate(attempts):
        cmd = [SHARP_BIN, "predict", "-i", str(image), "-o", str(outdir), "--no-render", "--device", dev]
        res = subprocess.run(cmd, capture_output=True, text=True, timeout=1800, env=env)
        if res.returncode == 0:
            plys = list(outdir.glob("**/*.ply"))
            if plys:
                return plys[0]
        err = res.stderr[-300:]
        if "out of memory" in res.stderr.lower() and i < len(attempts) - 1:
            print(f"    VRAM knapp ({dev}), weiche auf CPU aus…", flush=True)
            continue
        print(f"    SHARP-Fehler ({dev}): {err}", file=sys.stderr)
    return None


def merge_faces(pano: str, face_plys: dict, out_path: Path,
                min_opacity: float = 0.05, max_dist: float = 25.0) -> None:
    merged = None
    names_ref = None
    for face, ply in face_plys.items():
        names, data = read_gs_ply(ply)
        if names_ref is None:
            names_ref = names
        idx = {name: i for i, name in enumerate(names)}

        # Ausdünnen: transparente Splats und Fernpunkte verwerfen.
        logit = np.log(min_opacity / (1.0 - min_opacity))
        keep = data[:, idx["opacity"]] > logit
        pos = data[:, [idx["x"], idx["y"], idx["z"]]]
        keep &= (np.linalg.norm(pos, axis=1) < max_dist)
        data = data[keep]

        # Auf Web-taugliche Größe reduzieren (~200k Splats je Face).
        max_per_face = 200000
        if data.shape[0] > max_per_face:
            n_before = data.shape[0]
            sel = np.random.RandomState(42).choice(n_before, max_per_face, replace=False)
            data = data[np.sort(sel)]
            # Verbleibende Splats leicht vergrößern, um Lücken zu schließen.
            grow = np.float32(np.log(np.sqrt(n_before / max_per_face) * 0.5 + 0.5))
            for k in ("scale_0", "scale_1", "scale_2"):
                data[:, idx[k]] += grow

        rot = yaw_matrix(FACE_YAW[face])
        rq = rotmat_to_quat(rot)

        # Positionen drehen
        xyz = data[:, [idx["x"], idx["y"], idx["z"]]].astype(np.float64)
        xyz = xyz @ rot.T
        data[:, idx["x"]] = xyz[:, 0]
        data[:, idx["y"]] = xyz[:, 1]
        data[:, idx["z"]] = xyz[:, 2]

        # Orientierungen (Quaternion rot_0..rot_3 = w,x,y,z) drehen
        if "rot_0" in idx:
            quats = data[:, [idx["rot_0"], idx["rot_1"], idx["rot_2"], idx["rot_3"]]].astype(np.float64)
            quats = quat_mul(np.broadcast_to(rq, quats.shape), quats)
            data[:, idx["rot_0"]] = quats[:, 0]
            data[:, idx["rot_1"]] = quats[:, 1]
            data[:, idx["rot_2"]] = quats[:, 2]
            data[:, idx["rot_3"]] = quats[:, 3]

        # View-abhängige SH-Koeffizienten (f_rest_*) nullen: sie müssten
        # sonst mitrotiert werden; DC-Farbe bleibt korrekt.
        for name in names:
            if name.startswith("f_rest_"):
                data[:, idx[name]] = 0.0

        merged = data if merged is None else np.vstack([merged, data])

    write_gs_ply(out_path, names_ref, merged)
    print(f"  merged: {out_path} ({merged.shape[0]} Splats)")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--panos", default="", help="Kommagetrennte Pano-IDs, leer = alle")
    ap.add_argument("--faces", default=DEFAULT_FACES)
    ap.add_argument("--device", default="auto", help="cuda | cpu | auto (cuda wenn >4GB VRAM frei)")
    args = ap.parse_args()

    def pick_device() -> str:
        if args.device != "auto":
            return args.device
        try:
            free = int(subprocess.run(
                ["nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits"],
                capture_output=True, text=True, timeout=10).stdout.strip().splitlines()[0])
            return "cuda" if free > 7300 else "cpu"
        except Exception:
            return "cpu"

    faces = [f.strip() for f in args.faces.split(",") if f.strip()]
    if args.panos:
        panos = ["p" + p.strip().lstrip("p") for p in args.panos.split(",")]
    else:
        panos = sorted(
            (d.name.replace(".tiles", "") for d in PANOS_DIR.glob("p*.tiles")),
            key=lambda s: int(s[1:]),
        )

    OUT_DIR.mkdir(parents=True, exist_ok=True)
    for pano in panos:
        tiles = PANOS_DIR / f"{pano}.tiles"
        final = OUT_DIR / f"{pano}_360.ply"
        if final.exists():
            print(f"{pano}: vorhanden, übersprungen")
            continue
        device = pick_device()
        print(f"{pano}: SHARP auf {len(faces)} Faces… (device={device})", flush=True)
        face_plys = {}
        for face in faces:
            if face.startswith("h"):
                img = render_diagonal_view(tiles, pano, face)
            else:
                img = tiles / "vr" / f"pano_{face}.jpg"
            if img is None or not img.is_file():
                continue
            img = prepare_input(img, OUT_DIR / pano / face)
            ply = run_sharp(img, OUT_DIR / pano / face, device)
            if ply is not None:
                face_plys[face] = ply
                print(f"    {face}: ok", flush=True)
        if face_plys:
            merge_faces(pano, face_plys, final)


if __name__ == "__main__":
    main()