"""Camera-space placement, 3DGS PLY export and orbit rendering for Lift4D.

The object->camera convention is taken verbatim from Lift4D's
``sam3d/run_inference.py:render_gs_frame``:

    p_p3d = compose_transform(scale, quaternion_to_matrix(q_l2c), trans)
                .transform_points(p_local)                # PyTorch3D row-vector
    p_cam = p_p3d @ diag(-1, -1, 1)

PyTorch3D's ``Transform3d.rotate(R)`` applies ``p_row @ R``, i.e. R^T in the
usual column-vector sense -- which is why run_inference composes the Gaussian
orientations with ``quaternion_invert(q_l2c)``.  Written as one rigid map:

    R_o2c = diag(-1,-1,1) @ R(q_l2c)^T
    t_o2c = diag(-1,-1,1) @ t
    p_cam = R_o2c @ (scale * p_local) + t_o2c

Camera space is the usual OpenCV/3DGS frame (+x right, +y down, +z forward).
"""
from __future__ import annotations

import math
from typing import List

import numpy as np
import torch

SH_C0 = 0.28209479177387814


def _flip(device):
    """diag(-1, -1, 1): PyTorch3D -> camera, as in run_inference.py."""
    return torch.tensor([[-1.0, 0.0, 0.0], [0.0, -1.0, 0.0], [0.0, 0.0, 1.0]],
                        dtype=torch.float32, device=device)


def _mat_to_quat(R: torch.Tensor) -> torch.Tensor:
    from pytorch3d.transforms import matrix_to_quaternion
    return matrix_to_quaternion(R)


def to_camera_space(gs, scale_t, trans_t, rot_t, device="cuda"):
    """Return (means, quats, scales, colors, opacities) in camera space."""
    from pytorch3d.transforms import quaternion_to_matrix, quaternion_multiply

    xyz = gs.get_xyz.to(device).float()
    quats = gs.get_rotation.to(device).float()
    scales = gs.get_scaling.to(device).float()
    opac = gs.get_opacity.to(device).float().reshape(-1)
    dc = gs._features_dc.to(device).float()
    if dc.dim() == 3:
        dc = dc.squeeze(1)
    colors = torch.clamp(dc * SH_C0 + 0.5, 0.0, 1.0)

    s = scale_t.to(device).float().reshape(-1)
    if s.numel() == 1:
        s = s.repeat(3)
    t = trans_t.to(device).float().reshape(3)
    q = rot_t.to(device).float().reshape(4)

    M = _flip(device)
    R_o2c = M @ quaternion_to_matrix(q).transpose(0, 1)
    t_o2c = M @ t
    q_o2c = _mat_to_quat(R_o2c)

    means = (R_o2c @ (xyz * s).transpose(0, 1)).transpose(0, 1) + t_o2c
    quats = quaternion_multiply(q_o2c.unsqueeze(0).expand(quats.shape[0], -1), quats)
    scales = scales * s
    return means, quats, scales, colors, opac


# --------------------------------------------------------------------------
# 3DGS PLY export for gr.Model3D
# --------------------------------------------------------------------------

# camera space (y down, z forward) -> viewer space (y up, z toward viewer)
_VIEW_FLIP = np.array([[1.0, 0.0, 0.0], [0.0, -1.0, 0.0], [0.0, 0.0, -1.0]], dtype=np.float32)
_VIEW_QUAT = np.array([0.0, 1.0, 0.0, 0.0], dtype=np.float32)  # wxyz, Rx(180 deg)


def _quat_mul_np(a, b):
    aw, ax, ay, az = a
    bw, bx, by, bz = b[:, 0], b[:, 1], b[:, 2], b[:, 3]
    return np.stack([
        aw * bw - ax * bx - ay * by - az * bz,
        aw * bx + ax * bw + ay * bz - az * by,
        aw * by - ax * bz + ay * bw + az * bx,
        aw * bz + ax * by - ay * bx + az * bw,
    ], axis=-1)


def write_viewer_ply(path, means, quats, scales, colors, opacities, center=None):
    """Write a standard 3DGS .ply that gr.Model3D / gsplat.js can display.

    Gaussians come in camera space; they are rotated into the viewer's y-up
    frame and recentred so the object appears upright and framed.
    """
    m = means.detach().cpu().numpy().astype(np.float32)
    q = quats.detach().cpu().numpy().astype(np.float32)
    s = scales.detach().cpu().numpy().astype(np.float32)
    c = colors.detach().cpu().numpy().astype(np.float32)
    o = opacities.detach().cpu().numpy().astype(np.float32).reshape(-1)

    if center is None:
        center = np.median(m, axis=0)
    m = (m - center) @ _VIEW_FLIP.T
    q = _quat_mul_np(_VIEW_QUAT, q)

    eps = 1e-6
    f_dc = (c - 0.5) / SH_C0
    opacity = np.log(np.clip(o, eps, 1 - eps) / (1 - np.clip(o, eps, 1 - eps)))
    scale_log = np.log(np.clip(s, eps, None))
    normals = np.zeros_like(m)

    data = np.concatenate([m, normals, f_dc, opacity[:, None], scale_log, q], axis=1).astype(np.float32)

    header = (
        "ply\nformat binary_little_endian 1.0\n"
        f"element vertex {data.shape[0]}\n"
        "property float x\nproperty float y\nproperty float z\n"
        "property float nx\nproperty float ny\nproperty float nz\n"
        "property float f_dc_0\nproperty float f_dc_1\nproperty float f_dc_2\n"
        "property float opacity\n"
        "property float scale_0\nproperty float scale_1\nproperty float scale_2\n"
        "property float rot_0\nproperty float rot_1\nproperty float rot_2\nproperty float rot_3\n"
        "end_header\n"
    )
    with open(path, "wb") as fh:
        fh.write(header.encode("ascii"))
        fh.write(data.tobytes())
    return str(path)


# --------------------------------------------------------------------------
# orbit rendering with the Inria 3DGS rasterizer
# --------------------------------------------------------------------------

def _proj_matrix(znear, zfar, fovx, fovy, device):
    tx, ty = math.tan(fovx / 2), math.tan(fovy / 2)
    P = torch.zeros(4, 4, dtype=torch.float32, device=device)
    P[0, 0] = 1.0 / tx
    P[1, 1] = 1.0 / ty
    P[2, 2] = zfar / (zfar - znear)
    P[2, 3] = -(zfar * znear) / (zfar - znear)
    P[3, 2] = 1.0
    return P


def _look_at(cam_pos, target, device):
    """World->camera rotation/translation in the OpenCV (y-down) convention."""
    f = target - cam_pos
    f = f / torch.norm(f).clamp_min(1e-8)
    up_hint = torch.tensor([0.0, -1.0, 0.0], device=device)
    if torch.abs(torch.dot(f, up_hint)) > 0.999:
        up_hint = torch.tensor([0.0, 0.0, 1.0], device=device)
    r = torch.cross(f, up_hint, dim=0)
    r = r / torch.norm(r).clamp_min(1e-8)
    d = torch.cross(f, r, dim=0)
    R_wc = torch.stack([r, d, f], dim=0)
    t_wc = -R_wc @ cam_pos
    return R_wc, t_wc


def render_orbit(frames, size=512, orbit_steps=48, fov_deg=42.0,
                 elevation_deg=12.0, bg=(1.0, 1.0, 1.0), device="cuda"):
    """Render a 360-degree orbit that also plays the reconstructed sequence.

    ``frames`` is a list of (means, quats, scales, colors, opacities) tuples in
    camera space, one per reconstructed video frame.  Step ``t`` of the orbit
    shows the frame at ``t / orbit_steps`` of the sequence, so the video shows
    both the object's deformation and every side of the geometry.

    Each frame is recentred on its own median so the turntable stays
    object-centric: the subject's bulk translation through the scene (a running
    horse crosses far more than its own body length) would otherwise fling it
    out of frame.

    The orbit is phased so azimuth 0 starts broadside to the subject's longest
    horizontal axis (its length): for an elongated subject like the horse, an
    end-on start reads as an odd, foreshortened perspective, whereas presenting
    the long side first gives the natural side-profile turntable.
    """
    from diff_gaussian_rasterization import (GaussianRasterizationSettings,
                                             GaussianRasterizer)
    from pytorch3d.transforms import matrix_to_quaternion, quaternion_multiply

    centers = [f[0].median(dim=0).values for f in frames]
    rad = max(
        float(torch.quantile((f[0] - c).norm(dim=1), 0.97))
        for f, c in zip(frames, centers)
    )
    rad = max(rad, 1e-3)
    center = torch.zeros(3, dtype=torch.float32, device=device)

    # Phase the orbit so it starts broadside to the subject's longest horizontal
    # axis. Camera space is y-down, so the horizontal (ground) plane is x-z; the
    # dominant x-z eigenvector of the recentred means is the subject's length.
    # The horizontal view direction at azimuth a is [-sin a, cos a]; choosing
    # az_phase = atan2(lz, lx) makes that direction perpendicular to the length
    # axis at a = 0 (a side-on view) instead of looking straight down its length.
    xz = torch.cat([(f[0] - c)[:, [0, 2]] for f, c in zip(frames, centers)], dim=0)
    xz = (xz - xz.mean(dim=0)).to(torch.float32)
    cov = (xz.transpose(0, 1) @ xz) / max(xz.shape[0] - 1, 1)
    evals, evecs = torch.linalg.eigh(cov.cpu())
    length_axis = evecs[:, int(torch.argmax(evals))]
    az_phase = math.atan2(float(length_axis[1]), float(length_axis[0]))

    fov = math.radians(fov_deg)
    dist = rad / math.tan(fov / 2.0) * 1.12
    el = math.radians(elevation_deg)
    bg_t = torch.tensor(bg, dtype=torch.float32, device=device)
    P = _proj_matrix(0.01, 100.0, fov, fov, device)
    tan_half = math.tan(fov / 2.0)

    out = []
    n = len(frames)
    with torch.no_grad():
        for t in range(orbit_steps):
            az = az_phase + 2.0 * math.pi * t / orbit_steps
            direction = torch.tensor([
                math.sin(az) * math.cos(el),
                -math.sin(el),
                -math.cos(az) * math.cos(el),
            ], dtype=torch.float32, device=device)
            cam_pos = center + dist * direction
            R_wc, t_wc = _look_at(cam_pos, center, device)

            # Gaussians are moved into the orbit camera's frame here, so the
            # rasterizer gets an identity view matrix and a pure camera->clip
            # projection.
            fi = min(int(t * n / orbit_steps), n - 1)
            means, quats, scales, colors, opac = frames[fi]
            means = means - centers[fi]
            q_wc = matrix_to_quaternion(R_wc)
            quats_v = quaternion_multiply(q_wc.unsqueeze(0).expand(quats.shape[0], -1), quats)
            means_v = (R_wc @ means.transpose(0, 1)).transpose(0, 1) + t_wc

            settings = GaussianRasterizationSettings(
                image_height=size, image_width=size,
                tanfovx=tan_half, tanfovy=tan_half,
                bg=bg_t, scale_modifier=1.0,
                viewmatrix=torch.eye(4, dtype=torch.float32, device=device),
                projmatrix=P.transpose(0, 1).contiguous(),
                sh_degree=0,
                campos=torch.zeros(3, dtype=torch.float32, device=device),
                prefiltered=False, debug=False,
            )
            rasterizer = GaussianRasterizer(raster_settings=settings)
            screen = torch.zeros_like(means_v)
            image, _ = rasterizer(
                means3D=means_v.contiguous(),
                means2D=screen,
                shs=None,
                colors_precomp=colors.contiguous(),
                opacities=opac.reshape(-1, 1).contiguous(),
                scales=scales.contiguous(),
                rotations=quats_v.contiguous(),
                cov3D_precomp=None,
            )
            frame = (image.clamp(0, 1).permute(1, 2, 0) * 255).to(torch.uint8).cpu().numpy()
            out.append(frame)
    return out


def save_video(frames: List[np.ndarray], path, fps=12):
    import imageio.v2 as imageio
    imageio.mimwrite(str(path), frames, fps=fps, quality=8, macro_block_size=1,
                     codec="libx264", output_params=["-pix_fmt", "yuv420p"])
    return str(path)