# Adapted from TencentARC/Pixal3D, Copyright (c) 2026 Tencent, MIT (see LICENSE-PIXAL3D).

import torch
from torch import nn
from torch.nn import functional as F


def project_points_to_image_batch(
    points_3d: torch.Tensor,
    transform_matrix: torch.Tensor,
    camera_angle_x: torch.Tensor,
    resolution: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Project 3D points to pixel coordinates with a pinhole camera (32 mm sensor width).

    Args:
        points_3d: B N 3 points.
        transform_matrix: B 4 4 camera-to-world matrices.
        camera_angle_x: B horizontal fields of view, radians.
        resolution: Square image resolution in pixels.

    Returns:
        Pixel coordinates (B N 2), depth (B N) and a mask of points inside the view (B N).
    """
    batch, count = points_3d.shape[:2]
    ones = torch.ones(batch, count, 1, device=points_3d.device, dtype=points_3d.dtype)
    points_homogeneous = torch.cat([points_3d, ones], dim=-1)
    world_to_camera = torch.linalg.inv(transform_matrix.float()).to(transform_matrix.dtype)
    points_camera = torch.bmm(points_homogeneous, world_to_camera.transpose(-2, -1))[..., :3]
    x_cam, y_cam, z_cam = points_camera.unbind(-1)
    depth = -z_cam
    sensor_width = 32.0
    focal_length = 16.0 / torch.tan(camera_angle_x / 2.0)
    focal_length_pixels = (focal_length * resolution / sensor_width).unsqueeze(1)
    x_ndc = focal_length_pixels * x_cam / (-z_cam + 1e-08)
    y_ndc = focal_length_pixels * y_cam / (-z_cam + 1e-08)
    x_pixel = x_ndc + resolution / 2.0
    y_pixel = -y_ndc + resolution / 2.0
    valid_mask = (x_pixel >= 0) & (x_pixel < resolution) & (y_pixel >= 0) & (y_pixel < resolution) & (depth > 0)
    return torch.stack([x_pixel, y_pixel], dim=-1), depth, valid_mask


class ProjGrid(nn.Module):
    """Project a 3D grid (or selected grid cells) into the front view and sample image features there."""

    def __init__(self, grid_resolution: int, image_resolution: int):
        super().__init__()
        self.image_resolution = image_resolution
        one_dim = torch.linspace(-1, 1, grid_resolution)
        x, y, z = torch.meshgrid(one_dim, one_dim, one_dim, indexing="ij")
        rotation_matrix = torch.tensor([[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]])
        grid_points = torch.matmul(torch.stack((x, y, z), dim=-1), rotation_matrix.T).reshape(-1, 3)
        self.register_buffer("grid_points", grid_points)
        front_view_transform_matrix = torch.tensor(
            [
                [1.0, 0.0, 0.0, 0.0],
                [0.0, 0.0, -1.0, -2.0],
                [0.0, 1.0, 0.0, 0.0],
                [0.0, 0.0, 0.0, 1.0],
            ]
        )
        self.register_buffer("front_view_transform_matrix", front_view_transform_matrix)

    def forward(
        self,
        features_map: torch.Tensor,
        camera_angle_x: torch.Tensor,
        distance: torch.Tensor,
        mesh_scale: torch.Tensor,
        grid_indices: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Sample a B C H W feature map at the projected grid points; returns B N C."""
        batch = features_map.shape[0]
        grid_points = self.grid_points
        if grid_indices is not None:
            grid_points = grid_points[grid_indices]
        grid_points = grid_points.expand(batch, -1, -1)
        grid_points = grid_points / mesh_scale.unsqueeze(-1).unsqueeze(-1) / 2
        transform_matrix = self.front_view_transform_matrix.expand(batch, -1, -1).clone()
        transform_matrix[:, 1, 3] = -distance
        image_points, _, _ = project_points_to_image_batch(
            grid_points, transform_matrix, camera_angle_x, self.image_resolution
        )
        image_points_norm = (image_points + 0.5) / self.image_resolution * 2 - 1
        sample_grid = image_points_norm.view(batch, -1, 1, 2)
        sampled = F.grid_sample(features_map, sample_grid, mode="bilinear", align_corners=False, padding_mode="border")
        return sampled.squeeze(-1).permute(0, 2, 1)