"""
Pure-PyTorch differentiable point-based Gaussian renderer.
No CUDA compilation needed — works with standard PyTorch CUDA ops.

Simplified splat renderer:
- Projects 3D points to 2D screen space
- Renders each point as a 2D Gaussian splat
- Supports per-point color, opacity, and extra feature channels
- Differentiable w.r.t. point positions, colors, opacities, scales
"""
import torch
import torch.nn.functional as F


def project_points(means3d, viewmat, K, height, width):
    """Project 3D points to 2D screen coordinates."""
    R = viewmat[:3, :3]
    t = viewmat[:3, 3]
    cam_points = means3d @ R.T + t

    depths = cam_points[:, 2]
    visible = depths > 0.01

    fx, fy = K[0, 0], K[1, 1]
    cx, cy = K[0, 2], K[1, 2]

    safe_depths = torch.where(depths > 0.01, depths, torch.ones_like(depths))
    points_2d = torch.stack([
        fx * cam_points[:, 0] / safe_depths + cx,
        fy * cam_points[:, 1] / safe_depths + cy,
    ], dim=-1)

    return points_2d, depths, visible


def render_splat(points_2d, depths, colors, opacities, scales, height, width,
                 radius_factor=3.0):
    """Render points as 2D Gaussian splats using sorted alpha blending.

    Fully functional (no in-place ops) for differentiability.
    Gradients flow through colors, opacities, scales, and point positions.
    The alpha compositing buffer is detached (controls order only).
    """
    N = points_2d.shape[0]
    C = colors.shape[1]
    device = points_2d.device

    margin = 50
    valid = (depths > 0.01) & \
            (points_2d[:, 0] > -margin) & (points_2d[:, 0] < width + margin) & \
            (points_2d[:, 1] > -margin) & (points_2d[:, 1] < height + margin)

    if valid.sum() == 0:
        return torch.zeros(height, width, C, device=device)

    pts = points_2d[valid]
    cols = colors[valid]
    ops = opacities[valid]
    scs = scales[valid]
    deps = depths[valid]

    # Sort by depth (back-to-front)
    sort_idx = torch.argsort(deps, descending=True)
    pts = pts[sort_idx]
    cols = cols[sort_idx]
    ops = ops[sort_idx]
    scs = scs[sort_idx]

    radii = torch.clamp(scs * radius_factor, min=0.5, max=20.0)

    # Detached alpha compositing buffer — purely for ordering
    alpha_accum = torch.ones(height, width, device=device)

    # Collect all contributions, then sum functionally at the end
    image = torch.zeros(height, width, C, device=device)

    M = pts.shape[0]
    for i in range(M):
        px = pts[i, 0]
        py = pts[i, 1]
        r_int = int(radii[i].item()) + 1
        if r_int < 1:
            continue

        x0 = max(0, int(px.item()) - r_int)
        x1 = min(width, int(px.item()) + r_int + 1)
        y0 = max(0, int(py.item()) - r_int)
        y1 = min(height, int(py.item()) + r_int + 1)
        if x1 <= x0 or y1 <= y0:
            continue

        ys = torch.arange(y0, y1, device=device).float()
        xs = torch.arange(x0, x1, device=device).float()
        yy, xx = torch.meshgrid(ys, xs, indexing='ij')

        sigma = radii[i]
        gauss = torch.exp(-((xx - px) ** 2 + (yy - py) ** 2) / (2 * sigma ** 2 + 1e-8))
        alpha = gauss * ops[i]  # differentiable w.r.t. ops, sigma, px, py

        # Read current alpha (detached)
        current_alpha = alpha_accum[y0:y1, x0:x1].detach()
        blend_alpha = alpha * current_alpha  # differentiable

        # Add contribution to image via functional padding (no in-place)
        contrib = blend_alpha.unsqueeze(-1) * cols[i].unsqueeze(0).unsqueeze(0)
        # contrib: [h_region, w_region, C]
        pad_top = y0
        pad_bottom = height - y1
        pad_left = x0
        pad_right = width - x1
        padded = F.pad(contrib.permute(2, 0, 1),
                       (pad_left, pad_right, pad_top, pad_bottom),
                       mode='constant', value=0).permute(1, 2, 0)
        image = image + padded

        # Update detached alpha buffer (clone to avoid in-place modification)
        new_region = current_alpha * (1 - alpha.detach())
        alpha_new = alpha_accum.clone()
        alpha_new[y0:y1, x0:x1] = new_region
        alpha_accum = alpha_new.detach()

    return image


def rasterize_gaussians(means3d, quats, scales, opacities, colors,
                        viewmat, K, height, width):
    """Main rasterization entry point."""
    points_2d, depths, visible = project_points(means3d, viewmat, K, height, width)

    fx = K[0, 0]
    safe_depths = torch.where(depths > 0.01, depths, torch.ones_like(depths))
    screen_scales = (fx * scales.mean(dim=-1) / safe_depths).clamp(min=0.5, max=50.0)

    image = render_splat(points_2d, depths, colors, opacities, screen_scales,
                         height, width)

    return image