#!/usr/bin/env python3
# Published copy. Host names and paths are replaced by <placeholders> or left out of comments, nothing else is changed. The sha256 of the original is in sha256.txt.
"""Pure PyTorch 3D Gaussian splat renderer (no CUDA extension, no nvcc needed) for Idol PLY files.
Usage: render_turntable.py <ply> <camera.json|-> <out_dir> [--frames 240] [--size 720] [--sheet-only] [--bg 12,14,18]
Writes frames/f_0000.png ... and sheet.png (4 views: input-view, +90, back, -90). PLY = standard 3DGS (y points down).
Per frame: project, EWA 2D covariance (+0.3 px dilation), fragments per gaussian footprint, sort by (pixel, depth), alpha-composite in float64."""
import argparse, json, math, os, sys, time

import numpy as np
import torch
from plyfile import PlyData
from PIL import Image

C0 = 0.28209479177387814
DEV = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def load_ply(path):
    v = PlyData.read(path)["vertex"]
    xyz = np.stack([v["x"], v["y"], v["z"]], 1)
    fdc = np.stack([v["f_dc_0"], v["f_dc_1"], v["f_dc_2"]], 1)
    op = 1 / (1 + np.exp(-np.asarray(v["opacity"])))
    sc = np.exp(np.stack([v["scale_0"], v["scale_1"], v["scale_2"]], 1))
    q = np.stack([v["rot_0"], v["rot_1"], v["rot_2"], v["rot_3"]], 1)
    q = q / np.linalg.norm(q, axis=1, keepdims=True)
    col = np.clip(0.5 + C0 * fdc, 0, 1)
    t = lambda a: torch.from_numpy(np.ascontiguousarray(a)).float().to(DEV)
    return t(xyz), t(col), t(op), t(sc), t(q)


def quat_to_rot(q):
    w, x, y, z = q.unbind(1)
    R = torch.stack([1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y),
                     2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x),
                     2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y)], 1)
    return R.view(-1, 3, 3)


def look_at(pos, up=(0.0, -1.0, 0.0)):
    pos = np.asarray(pos, dtype=np.float64)
    f = -pos / np.linalg.norm(pos)
    up = np.asarray(up, dtype=np.float64)
    r = np.cross(f, up); r /= np.linalg.norm(r)
    u = np.cross(r, f)   # camera "up"
    # image y points down: use -u as camera y axis if up given is the world up; here world up = -y (PLY convention), so camera y(down) = -u
    W = np.stack([r, -u, f], 0)
    return torch.from_numpy(W).float().to(DEV), torch.from_numpy(pos).float().to(DEV)


@torch.no_grad()
def render(g, W, pos, size=720, fov_deg=40.0, bg=(12, 14, 18), rmax=14):
    xyz, col, op, sc, q = g
    N = xyz.shape[0]
    f = 0.5 * size / math.tan(math.radians(fov_deg) / 2)
    cx = cy = size / 2
    pc = (xyz - pos[None]) @ W.T            # camera coords
    z = pc[:, 2]
    ok = z > 0.05
    R = quat_to_rot(q)
    M = R * sc[:, None, :]                  # R S
    cov3 = M @ M.transpose(1, 2)
    zz = z.clamp_min(0.05)
    J = torch.zeros(N, 2, 3, device=DEV)
    J[:, 0, 0] = f / zz; J[:, 0, 2] = -f * pc[:, 0] / zz ** 2
    J[:, 1, 1] = f / zz; J[:, 1, 2] = -f * pc[:, 1] / zz ** 2
    T = J @ W[None]
    cov2 = T @ cov3 @ T.transpose(1, 2)
    a = cov2[:, 0, 0] + 0.3; b = cov2[:, 0, 1]; d = cov2[:, 1, 1] + 0.3
    det = (a * d - b * b).clamp_min(1e-9)
    mid = 0.5 * (a + d)
    lam = mid + torch.sqrt((mid * mid - det).clamp_min(0.1))
    rad = torch.ceil(3.0 * torch.sqrt(lam)).clamp(1, rmax)
    u = f * pc[:, 0] / zz + cx
    v = f * pc[:, 1] / zz + cy
    ok &= (u + rad > 0) & (u - rad < size) & (v + rad > 0) & (v - rad < size) & (op > 1 / 255)
    idx = torch.nonzero(ok).squeeze(1)
    if idx.numel() == 0:
        return np.tile(np.array(bg, dtype=np.uint8), (size, size, 1))
    u, v, rad, a, b, d, det, z, col, op = u[idx], v[idx], rad[idx].long(), a[idx], b[idx], d[idx], det[idx], z[idx], col[idx], op[idx]
    # depth rank (0 = nearest)
    rank = torch.empty_like(idx); rank[torch.argsort(z)] = torch.arange(idx.numel(), device=DEV)
    side = 2 * rad + 1
    cnt = side * side
    total = int(cnt.sum())
    gid = torch.repeat_interleave(torch.arange(idx.numel(), device=DEV), cnt)
    start = torch.cumsum(cnt, 0) - cnt
    local = torch.arange(total, device=DEV) - start[gid]
    sg = side[gid]
    ox = local % sg - rad[gid]
    oy = local // sg - rad[gid]
    px = torch.floor(u[gid]).long() + ox
    py = torch.floor(v[gid]).long() + oy
    inside = (px >= 0) & (px < size) & (py >= 0) & (py < size)
    gid, px, py = gid[inside], px[inside], py[inside]
    dx = px.float() + 0.5 - u[gid]
    dy = py.float() + 0.5 - v[gid]
    ia = d[gid] / det[gid]; ib = -b[gid] / det[gid]; ic = a[gid] / det[gid]
    pw = -0.5 * (ia * dx * dx + 2 * ib * dx * dy + ic * dy * dy)
    alpha = (op[gid] * torch.exp(pw)).clamp(max=0.99)
    keep = alpha >= 1 / 255
    gid, px, py, alpha = gid[keep], px[keep], py[keep], alpha[keep]
    pix = py * size + px
    key = pix * (1 << 24) + rank[gid]
    order = torch.argsort(key)
    gid, pix, alpha = gid[order], pix[order], alpha[order]
    lt = torch.log1p(-alpha.double())
    cs = torch.cumsum(lt, 0)
    excl = cs - lt
    first = torch.ones_like(pix, dtype=torch.bool); first[1:] = pix[1:] != pix[:-1]
    seg = torch.cumsum(first.long(), 0) - 1
    base = excl[first][seg]
    Tt = torch.exp(excl - base).float()
    w = (alpha * Tt)[:, None] * col[gid]
    img = torch.zeros(size * size, 3, device=DEV)
    img.index_add_(0, pix, w)
    lt_pix = torch.zeros(size * size, dtype=torch.float64, device=DEV)
    lt_pix.index_add_(0, pix, lt)
    trans = torch.exp(lt_pix).float()
    bgc = torch.tensor(bg, dtype=torch.float32, device=DEV) / 255.0
    img = img + trans[:, None] * bgc[None]
    return (img.view(size, size, 3).clamp(0, 1) * 255).byte().cpu().numpy()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("ply"); ap.add_argument("camera"); ap.add_argument("out")
    ap.add_argument("--frames", type=int, default=240); ap.add_argument("--size", type=int, default=720)
    ap.add_argument("--sheet-only", action="store_true"); ap.add_argument("--bg", default="12,14,18")
    ap.add_argument("--radius", type=float, default=2.0); ap.add_argument("--fov", type=float, default=40.0)
    ap.add_argument("--elev", type=float, default=None, help="elevation in degrees (default: input pitch clipped to 8..25)")
    a = ap.parse_args()
    bg = tuple(int(x) for x in a.bg.split(","))
    g = load_ply(a.ply)
    az0, elev = 0.0, 15.0
    if a.camera != "-" and os.path.exists(a.camera):
        c = json.load(open(a.camera))
        # input camera is in the z-up frame (x right, y forward-ish); PLY frame P = (x, -z, y)
        elev = float(np.clip(math.degrees(c["pitch"]), 8, 25))
        dx, dy = math.sin(c["yaw"]), math.cos(c["yaw"])      # horizontal direction in z-up frame (before cos(pitch))
        # P frame: horizontal plane is (x, z_P) with z_P = y_old
        az0 = math.atan2(dx, dy)
    if a.elev is not None:
        elev = a.elev
    os.makedirs(a.out, exist_ok=True)

    def cam(az):
        e = math.radians(elev)
        pos = a.radius * np.array([math.cos(e) * math.sin(az), -math.sin(e), math.cos(e) * math.cos(az)])
        return look_at(pos)

    t0 = time.time()
    views = []
    for k, off in enumerate((0, 90, 180, 270)):
        W, pos = cam(az0 + math.radians(off))
        views.append(render(g, W, pos, a.size, a.fov, bg))
    Image.fromarray(np.concatenate(views, 1)).save(os.path.join(a.out, "sheet.png"))
    print("sheet", round(time.time() - t0, 2), "s", flush=True)
    if a.sheet_only:
        return
    fd = os.path.join(a.out, "frames"); os.makedirs(fd, exist_ok=True)
    for i in range(a.frames):
        W, pos = cam(az0 + 2 * math.pi * i / a.frames)
        Image.fromarray(render(g, W, pos, a.size, a.fov, bg)).save(os.path.join(fd, "f_%04d.png" % i))
    print("frames", a.frames, round(time.time() - t0, 1), "s", flush=True)


if __name__ == "__main__":
    main()