import torch
import torch.nn.functional as F


def coords_grid(b, h, w, homogeneous=False, device=None):
    y, x = torch.meshgrid(
        torch.arange(h), torch.arange(w), indexing="ij"
    )  # [H, W]

    stacks = [x, y]

    if homogeneous:
        ones = torch.ones_like(x)  # [H, W]
        stacks.append(ones)

    grid = torch.stack(stacks, dim=0).float()  # [2, H, W] or [3, H, W]

    grid = grid[None].expand(b, -1, -1, -1)  # [B, 2, H, W] or [B, 3, H, W]

    if device is not None:
        grid = grid.to(device)

    return grid


def warp_with_pose_depth_candidates(
    feature1,
    ref_intrinsics,
    tgt_intrinsics,
    pose,
    depth,
    clamp_min_depth=1e-3,
    grid_sample_disable_cudnn=False,
):
    """
    feature1: [B, C, H, W]
    ref_intrinsics: [B, 3, 3], used to backproject reference pixels
    tgt_intrinsics: [B, 3, 3], used to reproject into target-view pixels
    pose: [B, 4, 4]
    depth: [B, D, H, W]
    """

    assert ref_intrinsics.size(1) == ref_intrinsics.size(2) == 3
    assert tgt_intrinsics.size(1) == tgt_intrinsics.size(2) == 3
    assert pose.size(1) == pose.size(2) == 4
    assert depth.dim() == 4

    b, d, h, w = depth.size()
    c = feature1.size(1)
    feature_dtype = feature1.dtype

    with torch.no_grad(), torch.autocast(
        device_type=depth.device.type, enabled=False
    ):
        ref_intrinsics = ref_intrinsics.float()
        tgt_intrinsics = tgt_intrinsics.float()
        pose = pose.float()
        depth = depth.float()
        # pixel coordinates
        grid = coords_grid(
            b, h, w, homogeneous=True, device=depth.device
        )  # [B, 3, H, W]
        # back project to 3D and transform viewpoint
        points = torch.linalg.inv(ref_intrinsics).bmm(
            grid.reshape(b, 3, -1)
        )  # [B, 3, H*W]
        points = torch.bmm(pose[:, :3, :3], points).unsqueeze(2) * depth.reshape(
            b, 1, d, h * w
        )  # [B, 3, D, H*W]
        points = points + pose[:, :3, -1:].unsqueeze(-1)  # [B, 3, D, H*W]
        # reproject to 2D image plane
        points = torch.bmm(tgt_intrinsics, points.view(b, 3, -1)).view(
            b, 3, d, h * w
        )  # [B, 3, D, H*W]
        pixel_coords = points[:, :2] / points[:, -1:].clamp(
            min=clamp_min_depth
        )  # [B, 2, D, H*W]

        # normalize to [-1, 1]
        x_grid = 2 * pixel_coords[:, 0] / (w - 1) - 1
        y_grid = 2 * pixel_coords[:, 1] / (h - 1) - 1

        grid = torch.stack([x_grid, y_grid], dim=-1)  # [B, D, H*W, 2]

    # Disable cuDNN for large grid_sample inputs; this avoids a known PyTorch
    # issue while keeping the default path for smaller inputs.
    # Ref: https://github.com/pytorch/pytorch/issues/88380
    if feature1.numel() > 1000000:
        grid_sample_disable_cudnn = True
    with torch.backends.cudnn.flags(enabled=not grid_sample_disable_cudnn):
        # grid_sample requires the feature and grid dtypes to match. Compute the
        # warp in FP32 for sub-pixel grid accuracy, then return to the feature dtype.
        with torch.autocast(device_type=feature1.device.type, enabled=False):
            warped_feature = F.grid_sample(
                feature1.float(),
                grid.view(b, d * h, w, 2),
                mode="bilinear",
                padding_mode="zeros",
                align_corners=True,
            ).view(
                b, c, d, h, w
            )  # [B, C, D, H, W]

    return warped_feature.to(feature_dtype)