"""
Patch Projector
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from models.render_utils import sample_ptsFeatures_from_featureMaps


class PatchProjector():
    def __init__(self, patch_size):
        self.h_patch_size = patch_size
        self.offsets = build_patch_offset(patch_size)  # the warping patch offsets index

        self.z_axis = torch.tensor([0, 0, 1]).float()

        self.plane_dist_thresh = 0.001

    # * correctness checked
    def pixel_warp(self, pts, imgs, intrinsics,
                   w2cs, img_wh=None):
        """

        :param pts: [N_rays, n_samples, 3]
        :param imgs: [N_views, 3, H, W]
        :param intrinsics: [N_views, 4, 4]
        :param c2ws: [N_views, 4, 4]
        :param img_wh:
        :return:
        """
        if img_wh is None:
            N_views, _, sizeH, sizeW = imgs.shape
            img_wh = [sizeW, sizeH]

        pts_color, valid_mask = sample_ptsFeatures_from_featureMaps(
            pts, imgs, w2cs, intrinsics, img_wh,
            proj_matrix=None, return_mask=True)  # [N_views, c, N_rays, n_samples], [N_views, N_rays, n_samples]

        pts_color = pts_color.permute(2, 3, 0, 1)
        valid_mask = valid_mask.permute(1, 2, 0)

        return pts_color, valid_mask  # [N_rays, n_samples, N_views,  3] , [N_rays, n_samples, N_views]

    def patch_warp(self, pts, uv, normals, src_imgs,
                   ref_intrinsic, src_intrinsics,
                   ref_c2w, src_c2ws, img_wh=None
                   ):
        """

        :param pts: [N_rays, n_samples, 3]
        :param uv : [N_rays, 2]  normalized in (-1, 1)
        :param normals: [N_rays, n_samples, 3]  The normal of pt in world space
        :param src_imgs: [N_src, 3, h, w]
        :param ref_intrinsic: [4,4]
        :param src_intrinsics: [N_src, 4, 4]
        :param ref_c2w: [4,4]
        :param src_c2ws: [N_src, 4, 4]
        :return:
        """
        device = pts.device

        N_rays, n_samples, _ = pts.shape
        N_pts = N_rays * n_samples

        N_src, _, sizeH, sizeW = src_imgs.shape

        if img_wh is not None:
            sizeW, sizeH = img_wh[0], img_wh[1]

        # scale uv from (-1, 1) to (0, W/H)
        uv[:, 0] = (uv[:, 0] + 1) / 2. * (sizeW - 1)
        uv[:, 1] = (uv[:, 1] + 1) / 2. * (sizeH - 1)

        ref_intr = ref_intrinsic[:3, :3]
        inv_ref_intr = torch.inverse(ref_intr)
        src_intrs = src_intrinsics[:, :3, :3]
        inv_src_intrs = torch.inverse(src_intrs)

        ref_pose = ref_c2w
        inv_ref_pose = torch.inverse(ref_pose)
        src_poses = src_c2ws
        inv_src_poses = torch.inverse(src_poses)

        ref_cam_loc = ref_pose[:3, 3].unsqueeze(0)  # [1, 3]
        sampled_dists = torch.norm(pts - ref_cam_loc, dim=-1)  # [N_pts, 1]

        relative_proj = inv_src_poses @ ref_pose
        R_rel = relative_proj[:, :3, :3]
        t_rel = relative_proj[:, :3, 3:]
        R_ref = inv_ref_pose[:3, :3]
        t_ref = inv_ref_pose[:3, 3:]

        pts = pts.view(-1, 3)
        normals = normals.view(-1, 3)

        with torch.no_grad():
            rot_normals = R_ref @ normals.unsqueeze(-1)  # [N_pts, 3, 1]
            points_in_ref = R_ref @ pts.unsqueeze(
                -1) + t_ref  # [N_pts, 3, 1]  points in the reference frame coordiantes system
            d1 = torch.sum(rot_normals * points_in_ref, dim=1).unsqueeze(
                1)  # distance from the plane to ref camera center

            d2 = torch.sum(rot_normals.unsqueeze(1) * (-R_rel.transpose(1, 2) @ t_rel).unsqueeze(0),
                           dim=2)  # distance from the plane to src camera center
            valid_hom = (torch.abs(d1) > self.plane_dist_thresh) & (
                    torch.abs(d1 - d2) > self.plane_dist_thresh) & ((d2 / d1) < 1)

            d1 = d1.squeeze()
            sign = torch.sign(d1)
            sign[sign == 0] = 1
            d = torch.clamp(torch.abs(d1), 1e-8) * sign

            H = src_intrs.unsqueeze(1) @ (
                    R_rel.unsqueeze(1) + t_rel.unsqueeze(1) @ rot_normals.view(1, N_pts, 1, 3) / d.view(1,
                                                                                                        N_pts,
                                                                                                        1, 1)
            ) @ inv_ref_intr.view(1, 1, 3, 3)

            # replace invalid homs with fronto-parallel homographies
            H_invalid = src_intrs.unsqueeze(1) @ (
                    R_rel.unsqueeze(1) + t_rel.unsqueeze(1) @ self.z_axis.to(device).view(1, 1, 1, 3).expand(-1, N_pts,
                                                                                                             -1,
                                                                                                             -1) / sampled_dists.view(
                1, N_pts, 1, 1)
            ) @ inv_ref_intr.view(1, 1, 3, 3)
            tmp_m = ~valid_hom.view(-1, N_src).t()
            H[tmp_m] = H_invalid[tmp_m]

        pixels = uv.view(N_rays, 1, 2) + self.offsets.float().to(device)
        Npx = pixels.shape[1]
        grid, warp_mask_full = self.patch_homography(H, pixels)

        warp_mask_full = warp_mask_full & (grid[..., 0] < (sizeW - self.h_patch_size)) & (
                grid[..., 1] < (sizeH - self.h_patch_size)) & (grid >= self.h_patch_size).all(dim=-1)
        warp_mask_full = warp_mask_full.view(N_src, N_rays, n_samples, Npx)

        grid = torch.clamp(normalize(grid, sizeH, sizeW), -10, 10)

        sampled_rgb_val = F.grid_sample(src_imgs, grid.view(N_src, -1, 1, 2), align_corners=True).squeeze(
            -1).transpose(1, 2)
        sampled_rgb_val = sampled_rgb_val.view(N_src, N_rays, n_samples, Npx, 3)

        warp_mask_full = warp_mask_full.permute(1, 2, 0, 3).contiguous()  # (N_rays, n_samples, N_src, Npx)
        sampled_rgb_val = sampled_rgb_val.permute(1, 2, 0, 3, 4).contiguous()  # (N_rays, n_samples, N_src, Npx, 3)

        return sampled_rgb_val, warp_mask_full

    def patch_homography(self, H, uv):
        N, Npx = uv.shape[:2]
        Nsrc = H.shape[0]
        H = H.view(Nsrc, N, -1, 3, 3)
        hom_uv = add_hom(uv)

        # einsum is 30 times faster
        # tmp = (H.view(Nsrc, N, -1, 1, 3, 3) @ hom_uv.view(1, N, 1, -1, 3, 1)).squeeze(-1).view(Nsrc, -1, 3)
        tmp = torch.einsum("vprik,pok->vproi", H, hom_uv).reshape(Nsrc, -1, 3)

        grid = tmp[..., :2] / torch.clamp(tmp[..., 2:], 1e-8)
        mask = tmp[..., 2] > 0
        return grid, mask


def add_hom(pts):
    try:
        dev = pts.device
        ones = torch.ones(pts.shape[:-1], device=dev).unsqueeze(-1)
        return torch.cat((pts, ones), dim=-1)

    except AttributeError:
        ones = np.ones((pts.shape[0], 1))
        return np.concatenate((pts, ones), axis=1)


def normalize(flow, h, w, clamp=None):
    # either h and w are simple float or N torch.tensor where N batch size
    try:
        h.device

    except AttributeError:
        h = torch.tensor(h, device=flow.device).float().unsqueeze(0)
        w = torch.tensor(w, device=flow.device).float().unsqueeze(0)

    if len(flow.shape) == 4:
        w = w.unsqueeze(1).unsqueeze(2)
        h = h.unsqueeze(1).unsqueeze(2)
    elif len(flow.shape) == 3:
        w = w.unsqueeze(1)
        h = h.unsqueeze(1)
    elif len(flow.shape) == 5:
        w = w.unsqueeze(0).unsqueeze(2).unsqueeze(2)
        h = h.unsqueeze(0).unsqueeze(2).unsqueeze(2)

    res = torch.empty_like(flow)
    if res.shape[-1] == 3:
        res[..., 2] = 1

    # for grid_sample with align_corners=True
    # https://github.com/pytorch/pytorch/blob/c371542efc31b1abfe6f388042aa3ab0cef935f2/aten/src/ATen/native/GridSampler.h#L33
    res[..., 0] = 2 * flow[..., 0] / (w - 1) - 1
    res[..., 1] = 2 * flow[..., 1] / (h - 1) - 1

    if clamp:
        return torch.clamp(res, -clamp, clamp)
    else:
        return res


def build_patch_offset(h_patch_size):
    offsets = torch.arange(-h_patch_size, h_patch_size + 1)
    return torch.stack(torch.meshgrid(offsets, offsets, indexing="ij")[::-1], dim=-1).view(1, -1, 2)  # nb_pixels_patch * 2