#!/usr/bin/env python3
"""Equirect → überlappende Perspektiv-Ansichten für COLMAP.

12 Ansichten bei Pitch 0° (60° FOV, 30°-Schritte) + 4 bei Pitch -35°
(75° FOV) pro Pano. Ausgabe 1280px, Dateiname pN_yYAW_pPITCH.jpg.
"""
from __future__ import annotations

import sys
from pathlib import Path

import numpy as np
from PIL import Image

PANOS_DIR = Path("/home/panomity.de/vr.panomity.com/panos")
OUT_DIR = Path(sys.argv[1] if len(sys.argv) > 1 else "/home/panomity.de/gsplat-apartment/images2")

SIZE = 1280


def e2p(equi: np.ndarray, fov_deg: float, yaw_deg: float, pitch_deg: float, size: int) -> np.ndarray:
    """Equirectangular → perspektivische Ansicht (Pinhole, quadratisch)."""
    h, w, _ = equi.shape
    fov = np.deg2rad(fov_deg)
    yaw = np.deg2rad(yaw_deg)
    pitch = np.deg2rad(pitch_deg)

    f = (size / 2) / np.tan(fov / 2)
    xs = np.arange(size) - (size - 1) / 2
    ys = np.arange(size) - (size - 1) / 2
    px, py = np.meshgrid(xs, ys)

    # Kamerastrahlen (x rechts, y unten, z vorwärts)
    dirs = np.stack([px / f, py / f, np.ones_like(px)], axis=-1)
    dirs /= np.linalg.norm(dirs, axis=-1, keepdims=True)

    # Pitch (um x) dann Yaw (um y)
    cp, sp = np.cos(pitch), np.sin(pitch)
    rot_x = np.array([[1, 0, 0], [0, cp, -sp], [0, sp, cp]])
    cy, sy = np.cos(yaw), np.sin(yaw)
    rot_y = np.array([[cy, 0, sy], [0, 1, 0], [-sy, 0, cy]])
    dirs = dirs @ rot_x.T @ rot_y.T

    dx, dy, dz = dirs[..., 0], dirs[..., 1], dirs[..., 2]
    lon = np.arctan2(dx, dz)          # -pi..pi
    lat = np.arcsin(np.clip(dy, -1, 1))  # -pi/2..pi/2

    u = (lon / (2 * np.pi) + 0.5) * (w - 1)
    v = (lat / np.pi + 0.5) * (h - 1)

    u0 = np.floor(u).astype(int) % w
    u1 = (u0 + 1) % w
    v0 = np.clip(np.floor(v).astype(int), 0, h - 1)
    v1 = np.clip(v0 + 1, 0, h - 1)
    fu = (u - np.floor(u))[..., None]
    fv = (v - np.floor(v))[..., None]

    img = (equi[v0, u0] * (1 - fu) * (1 - fv)
           + equi[v0, u1] * fu * (1 - fv)
           + equi[v1, u0] * (1 - fu) * fv
           + equi[v1, u1] * fu * fv)
    return img.astype(np.uint8)


def main() -> None:
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    # Einheitliches 60°-FOV, damit COLMAP eine einzige Pinhole-Kamera nutzen kann
    # (f = 640 / tan(30°) = 1108.51 bei 1280px).
    views = [(60.0, yaw, 0.0) for yaw in range(0, 360, 30)]
    views += [(60.0, yaw, -35.0) for yaw in range(0, 360, 90)]

    for tiles in sorted(PANOS_DIR.glob("p*.tiles"), key=lambda d: int(d.name[1:-6])):
        pano = tiles.name.replace(".tiles", "")
        eq_path = tiles / "equirect.jpg"
        if not eq_path.is_file():
            continue
        done = [OUT_DIR / f"{pano}_y{int(y):03d}_p{int(p):+03d}.jpg" for _, y, p in views]
        if all(d.is_file() for d in done):
            print(f"{pano}: vorhanden")
            continue
        equi = np.asarray(Image.open(eq_path).convert("RGB"))
        for fov, yaw, pitch in views:
            out = OUT_DIR / f"{pano}_y{int(yaw):03d}_p{int(pitch):+03d}.jpg"
            if out.is_file():
                continue
            img = e2p(equi, fov, yaw, pitch, SIZE)
            Image.fromarray(img).save(out, quality=90)
        print(f"{pano}: {len(views)} Ansichten")


if __name__ == "__main__":
    main()