import torch
import numpy as np
from anydata.geometry.camera import Camera
from anydata.geometry.depth import calculate_normals, calc_dot_prod
from anydata.utils.decorators import iterate12


def apply_stride(val, stride, capped=False):
    """Return closest resolution based on stride value"""
    if stride is None or stride == 1:
        return val
    rem = val % stride
    val2 = val // stride * stride
    if rem > stride / 2:
        val2 += stride
    if capped and val2 > val:
        val2 -= stride
    return val2 


@iterate12
def depth_filter(depth, intrinsics, thr=0.1):
    """Filter depth map based on surface normal"""

    hw = depth.shape[1:]
    K = torch.tensor(intrinsics).unsqueeze(0).float()
    cam = Camera(hw=hw, K=K)

    depth_torch = torch.tensor(depth).unsqueeze(0)

    normals = calculate_normals(depth_torch, camera=cam)
    points = cam.reconstruct_depth_map(depth_torch)
    dotprod = calc_dot_prod(points, normals)
    dotprod[torch.isnan(dotprod)] = 0
    invalid = dotprod[:, 0].abs() < thr

    mask = np.ones_like(depth_torch)
    mask[invalid.unsqueeze(1)] = 0

    return depth * mask[0]