# SPDX-License-Identifier: Apache-2.0
# Copyright 2026 Lucasgitguy-co
"""Activated Gaussian attributes and standard degree-zero 3DGS PLY IO.

Coordinates are preserved as supplied. Idol's *exported PLY* uses its y-up frame;
we do not apply Idol's internal z-up conversion a second time.
"""
from dataclasses import dataclass
from pathlib import Path

import numpy as np
from plyfile import PlyData, PlyElement

C0 = 0.28209479177387814


@dataclass
class Splats:
    xyz: np.ndarray
    rgb: np.ndarray
    scale: np.ndarray
    quat: np.ndarray  # w,x,y,z
    opacity: np.ndarray

    def __post_init__(self):
        n = len(self.xyz)
        if n == 0:
            raise ValueError("No Gaussians: refusing to write or train on an empty object.")
        for name, width in [("xyz", 3), ("rgb", 3), ("scale", 3), ("quat", 4), ("opacity", 1)]:
            a = np.asarray(getattr(self, name), dtype=np.float32).reshape(-1, width)
            if a.shape != (n, width) or not np.isfinite(a).all():
                raise ValueError(f"Invalid {name} shape or nonfinite values")
            setattr(self, name, a)
        if (self.scale <= 0).any():
            raise ValueError("Gaussian scales must be positive")
        if (self.opacity < 0).any() or (self.opacity > 1).any():
            raise ValueError("Opacity must lie in [0,1]")
        norm = np.linalg.norm(self.quat, axis=1, keepdims=True)
        if (norm < 1e-8).any():
            raise ValueError("Zero quaternion")
        self.quat = self.quat / norm

    def take(self, indices):
        return Splats(**{k: getattr(self, k)[indices] for k in self.__dataclass_fields__})


def read_ply(path):
    v = PlyData.read(str(path))["vertex"].data
    required = ["x", "y", "z", "f_dc_0", "f_dc_1", "f_dc_2", "opacity",
                "scale_0", "scale_1", "scale_2", "rot_0", "rot_1", "rot_2", "rot_3"]
    if not set(required).issubset(v.dtype.names):
        raise ValueError("Expected a Gaussian-splat PLY, not a mesh or ordinary point-cloud PLY")
    col = lambda names: np.column_stack([v[k] for k in names]).astype(np.float32)
    return Splats(col(["x", "y", "z"]), 0.5 + C0 * col([f"f_dc_{i}" for i in range(3)]),
                  np.exp(np.clip(col([f"scale_{i}" for i in range(3)]), -20, 5)),
                  col([f"rot_{i}" for i in range(4)]),
                  1 / (1 + np.exp(-np.clip(col(["opacity"]), -30, 30))))


def write_ply(path, splats):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    names = ["x", "y", "z", "nx", "ny", "nz", "f_dc_0", "f_dc_1", "f_dc_2", "opacity",
             "scale_0", "scale_1", "scale_2", "rot_0", "rot_1", "rot_2", "rot_3"]
    opacity = np.clip(splats.opacity, 1e-6, 1 - 1e-6)
    values = np.concatenate([splats.xyz, np.zeros_like(splats.xyz), (splats.rgb - 0.5) / C0,
                             np.log(opacity / (1 - opacity)), np.log(splats.scale), splats.quat], axis=1)
    table = np.empty(len(values), dtype=[(k, "f4") for k in names])
    for i, name in enumerate(names):
        table[name] = values[:, i]
    PlyData([PlyElement.describe(table, "vertex")], text=False).write(str(path))


def rotation_matrices(q):
    w, x, y, z = q.T
    return np.stack([1-2*(y*y+z*z), 2*(x*y-z*w), 2*(x*z+y*w),
                     2*(x*y+z*w), 1-2*(x*x+z*z), 2*(y*z-x*w),
                     2*(x*z-y*w), 2*(y*z+x*w), 1-2*(x*x+y*y)], axis=-1).reshape(-1, 3, 3)


def preview(splats, size=256, yaw=0.0):
    """CPU orthographic Gaussian preview; debugging only, not an evaluation renderer."""
    from PIL import Image
    c, s = np.cos(yaw), np.sin(yaw)
    view = np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]], dtype=np.float32)
    xyz = splats.xyz @ view.T
    rot = rotation_matrices(splats.quat)
    cov = (rot * (splats.scale ** 2)[:, None, :]) @ rot.transpose(0, 2, 1)
    cov = view[None] @ cov @ view.T[None]
    pixels = size / 1.25
    canvas = np.zeros((size, size, 3), dtype=np.float32)
    # Back to front, camera on +z. Image y points down.
    for i in np.argsort(xyz[:, 2]):
        center = np.array([xyz[i, 0], -xyz[i, 1]]) * pixels + size / 2
        flip = np.diag([1., -1.])
        covariance = flip @ cov[i, :2, :2] @ flip * pixels ** 2 + np.eye(2) * 0.3
        radius = min(size, int(np.ceil(3 * np.sqrt(np.linalg.eigvalsh(covariance).max()))))
        lo = np.maximum(0, np.floor(center - radius).astype(int))
        hi = np.minimum(size, np.ceil(center + radius + 1).astype(int))
        if np.any(hi <= lo):
            continue
        yy, xx = np.mgrid[lo[1]:hi[1], lo[0]:hi[0]]
        delta = np.stack([xx - center[0], yy - center[1]], axis=-1)
        exponent = np.einsum("...i,ij,...j->...", delta, np.linalg.inv(covariance), delta)
        alpha = (splats.opacity[i, 0] * np.exp(-0.5 * exponent)).clip(0, .99)[..., None]
        patch = canvas[lo[1]:hi[1], lo[0]:hi[0]]
        patch[:] = patch * (1-alpha) + np.clip(splats.rgb[i], 0, 1) * alpha
    return Image.fromarray((canvas.clip(0, 1) * 255).astype(np.uint8))