import torch
import torch.nn.functional as F


def estimate_surface_normal(
    points: torch.Tensor, d: int = 2, mode: str = "closest"
) -> torch.Tensor:
    # estimate surface normal from coordinated point clouds
    # re-implemented the following codes with pytorch:
    # https://github.com/wkentaro/morefusion/blob/master/morefusion/geometry/estimate_pointcloud_normals.py
    # https://github.com/jmccormac/pySceneNetRGBD/blob/master/calculate_surface_normals.py

    assert points.dim() == 4, f"expected (B,3,H,W), but got {points.shape}"
    B, C, H, W = points.shape
    assert C == 3, f"expected C==3, but got {C}"
    device = points.device

    # points = F.pad(points, (0, 0, d, d), mode="constant", value=float("inf"))
    points = F.pad(points, (0, 0, d, d), mode="replicate")
    points = F.pad(points, (d, d, 0, 0), mode="circular")
    points = points.permute(0, 2, 3, 1)  # (B,H,W,3)

    # 8 adjacent offsets
    #  -----------
    # | 7 | 6 | 5 |
    #  -----------
    # | 0 |   | 4 |
    #  -----------
    # | 1 | 2 | 3 |
    #  -----------
    offsets = torch.tensor(
        [
            # (dh,dw)
            (-d, 0),  # 0
            (-d, d),  # 1
            (0, d),  # 2
            (d, d),  # 3
            (d, 0),  # 4
            (d, -d),  # 5
            (0, -d),  # 6
            (-d, -d),  # 7
        ],
        device=device,
    )

    # (B,H,W) indices
    b = torch.arange(B, device=device)[:, None, None]
    h = torch.arange(H, device=device)[None, :, None]
    w = torch.arange(W, device=device)[None, None, :]
    k = torch.arange(8, device=device)

    # anchor points
    b1 = b[:, None]  # (B,1,1,1)
    h1 = h[:, None] + d  # (1,1,H,1)
    w1 = w[:, None] + d  # (1,1,1,W)
    anchors = points[b1, h1, w1]  # (B,H,W,3) -> (B,1,H,W,3)

    # neighbor points
    offset = offsets[k]  # (8,2)
    b2 = b1
    h2 = h1 + offset[None, :, 0, None, None]  # (1,8,H,1)
    w2 = w1 + offset[None, :, 1, None, None]  # (1,8,1,W)
    points1 = points[b2, h2, w2]  # (B,8,H,W,3)

    # anothor neighbor points
    offset = offsets[(k + 2) % 8]
    b3 = b1
    h3 = h1 + offset[None, :, 0, None, None]
    w3 = w1 + offset[None, :, 1, None, None]
    points2 = points[b3, h3, w3]  # (B,8,H,W,3)

    if mode == "closest":
        # find the closest neighbor pair
        diff = torch.norm(points1 - anchors, dim=4)
        diff = diff + torch.norm(points2 - anchors, dim=4)
        i = torch.argmin(diff, dim=1)  # (B,H,W)
        # get normals by cross product
        anchors = anchors[b, 0, h, w]  # (B,H,W,3)
        points1 = points1[b, i, h, w]  # (B,H,W,3)
        points2 = points2[b, i, h, w]  # (B,H,W,3)
        vector1 = points1 - anchors
        vector2 = points2 - anchors
        normals = torch.cross(vector1, vector2, dim=-1)  # (B,H,W,3)
    elif mode == "mean":
        # get normals by cross product
        vector1 = points1 - anchors
        vector2 = points2 - anchors
        normals = torch.cross(vector1, vector2, dim=-1)  # (B,8,H,W,3)
        normals = normals.mean(dim=1)  # (B,H,W,3)
    else:
        raise NotImplementedError(mode)

    normals = normals / (torch.norm(normals, dim=3, keepdim=True) + 1e-8)
    normals = normals.permute(0, 3, 1, 2)  # (B,3,H,W)

    return normals