content: # The codes are partly from IBRNet
  
  import torch
  import torch.nn.functional as F
  from models.render_utils import sample_ptsFeatures_from_featureMaps, sample_ptsFeatures_from_featureVolume
  
  def safe_l2_normalize(x, dim=None, eps=1e-6):
      return F.normalize(x, p=2, dim=dim, eps=eps)
  
  class Projector():
      """
      Obtain features from geometryVolume and rendering_feature_maps for generalized rendering
      """
  
      def compute_angle(self, xyz, query_c2w, supporting_c2ws):
          """
  
          :param xyz: [N_rays, n_samples,3 ]
          :param query_c2w: [1,4,4]
          :param supporting_c2ws: [n,4,4]
          :return:
          """
          N_rays, n_samples, _ = xyz.shape
          num_views = supporting_c2ws.shape[0]
          xyz = xyz.reshape(-1, 3)
  
          ray2tar_pose = (query_c2w[:, :3, 3].unsqueeze(1) - xyz.unsqueeze(0))
          ray2tar_pose /= (torch.norm(ray2tar_pose, dim=-1, keepdim=True) + 1e-6)
          ray2support_pose = (supporting_c2ws[:, :3, 3].unsqueeze(1) - xyz.unsqueeze(0))
          ray2support_pose /= (torch.norm(ray2support_pose, dim=-1, keepdim=True) + 1e-6)
          ray_diff = ray2tar_pose - ray2support_pose
          ray_diff_norm = torch.norm(ray_diff, dim=-1, keepdim=True)
          ray_diff_dot = torch.sum(ray2tar_pose * ray2support_pose, dim=-1, keepdim=True)
          ray_diff_direction = ray_diff / torch.clamp(ray_diff_norm, min=1e-6)
          ray_diff = torch.cat([ray_diff_direction, ray_diff_dot], dim=-1)
          ray_diff = ray_diff.reshape((num_views, N_rays, n_samples, 4))  # the last dimension (4) is dot-product
          return ray_diff.detach()
  
  
      def compute_angle_view_independent(self, xyz, surface_normals, supporting_c2ws):
          """
  
          :param xyz: [N_rays, n_samples,3 ]
          :param surface_normals: [N_rays, n_samples,3 ]
          :param supporting_c2ws: [n,4,4]
          :return:
          """
          N_rays, n_samples, _ = xyz.shape
          num_views = supporting_c2ws.shape[0]
          xyz = xyz.reshape(-1, 3)
  
          ray2tar_pose = surface_normals
          ray2support_pose = (supporting_c2ws[:, :3, 3].unsqueeze(1) - xyz.unsqueeze(0))
          ray2support_pose /= (torch.norm(ray2support_pose, dim=-1, keepdim=True) + 1e-6)
          ray_diff = ray2tar_pose - ray2support_pose
          ray_diff_norm = torch.norm(ray_diff, dim=-1, keepdim=True)
          ray_diff_dot = torch.sum(ray2tar_pose * ray2support_pose, dim=-1, keepdim=True)
          ray_diff_direction = ray_diff / torch.clamp(ray_diff_norm, min=1e-6)
          ray_diff = torch.cat([ray_diff_direction, ray_diff_dot], dim=-1)
          ray_diff = ray_diff.reshape((num_views, N_rays, n_samples, 4))  # the last dimension (4) is dot-product, 
                                                                          # and the first three dimensions is the normalized ray diff vector
          return ray_diff.detach()
  
      @torch.no_grad()
      def compute_z_diff(self, xyz, w2cs, intrinsics, pred_depth_values):
          """
          compute the depth difference of query pts projected on the image and the predicted depth values of the image
          :param xyz:  [N_rays, n_samples,3 ]
          :param w2cs:  [N_views, 4, 4]
          :param intrinsics: [N_views, 3, 3]
          :param pred_depth_values: [N_views, N_rays, n_samples,1 ]
          :param pred_depth_masks: [N_views, N_rays, n_samples]
          :return:
          """
          device = xyz.device
          N_views = w2cs.shape[0]
          N_rays, n_samples, _ = xyz.shape
          proj_matrix = torch.matmul(intrinsics, w2cs[:, :3, :])
  
          proj_rot = proj_matrix[:, :3, :3]
          proj_trans = proj_matrix[:, :3, 3:]
  
          batch_xyz = xyz.permute(2, 0, 1).contiguous().view(1, 3, N_rays * n_samples).repeat(N_views, 1, 1)
  
          proj_xyz = proj_rot.bmm(batch_xyz) + proj_trans
  
          # X = proj_xyz[:, 0]
          # Y = proj_xyz[:, 1]
          Z = proj_xyz[:, 2].clamp(min=1e-3)  # [N_views, N_rays*n_samples]
          proj_z = Z.view(N_views, N_rays, n_samples, 1)
  
          z_diff = proj_z - pred_depth_values  # [N_views, N_rays, n_samples,1 ]
  
          return z_diff
  
      def compute(self,
                  pts,
                  # * 3d geometry feature volumes
                  geometryVolume=None,
                  geometryVolumeMask=None,
                  vol_dims=None,
                  partial_vol_origin=None,
                  vol_size=None,
                  # * 2d rendering feature maps
                  rendering_feature_maps=None,
                  color_maps=None,
                  w2cs=None,
                  intrinsics=None,
                  img_wh=None,
                  query_img_idx=0,  # the index of the N_views dim for rendering
                  query_c2w=None,
                  pred_depth_maps=None,   # no use here
                  pred_depth_masks=None   # no use here
                  ):
          """
          extract features of pts for rendering
          :param pts:
          :param geometryVolume:
          :param vol_dims:
          :param partial_vol_origin:
          :param vol_size:
          :param rendering_feature_maps:
          :param color_maps:
          :param w2cs:
          :param intrinsics:
          :param img_wh:
          :param rendering_img_idx: by default, we render the first view of w2cs
          :return:
          """
          device = pts.device
          c2ws = torch.inverse(w2cs)
  
          if len(pts.shape) == 2:
              pts = pts[None, :, :]
  
          N_rays, n_samples, _ = pts.shape
          N_views = rendering_feature_maps.shape[0]  # shape (N_views, C, H, W)
  
          supporting_img_idxs = torch.LongTensor([x for x in range(N_views) if x != query_img_idx]).to(device)
          query_img_idx = torch.LongTensor([query_img_idx]).to(device)
  
          if query_c2w is None and query_img_idx > -1:
              query_c2w = torch.index_select(c2ws, 0, query_img_idx)
              supporting_c2ws = torch.index_select(c2ws, 0, supporting_img_idxs)
              supporting_w2cs = torch.index_select(w2cs, 0, supporting_img_idxs)
              supporting_rendering_feature_maps = torch.index_select(rendering_feature_maps, 0, supporting_img_idxs)
              supporting_color_maps = torch.index_select(color_maps, 0, supporting_img_idxs)
              supporting_intrinsics = torch.index_select(intrinsics, 0, supporting_img_idxs)
  
              if pred_depth_maps is not None:
                  supporting_depth_maps = torch.index_select(pred_depth_maps, 0, supporting_img_idxs)
                  supporting_depth_masks = torch.index_select(pred_depth_masks, 0, supporting_img_idxs)
              # print("N_supporting_views: ", N_views - 1)
              N_supporting_views = N_views - 1
          else:
              supporting_c2ws = c2ws
              supporting_w2cs = w2cs
              supporting_rendering_feature_maps = rendering_feature_maps
              supporting_color_maps = color_maps
              supporting_intrinsics = intrinsics
              supporting_depth_maps = pred_depth_masks
              supporting_depth_masks = pred_depth_masks
              # print("N_supporting_views: ", N_views)
              N_supporting_views = N_views
          # import ipdb; ipdb.set_trace()
          if geometryVolume is not None:
              # * sample feature of pts from 3D feature volume
              pts_geometry_feature, pts_geometry_masks_0 = sample_ptsFeatures_from_featureVolume(
                  pts, geometryVolume, vol_dims,
                  partial_vol_origin, vol_size)  # [N_rays, n_samples, C], [N_rays, n_samples]
  
              if len(geometryVolumeMask.shape) == 3:
                  geometryVolumeMask = geometryVolumeMask[None, :, :, :]
  
              pts_geometry_masks_1, _ = sample_ptsFeatures_from_featureVolume(
                  pts, geometryVolumeMask.to(geometryVolume.dtype), vol_dims,
                  partial_vol_origin, vol_size)  # [N_rays, n_samples, C]
  
              pts_geometry_masks = pts_geometry_masks_0 & (pts_geometry_masks_1[..., 0] > 0)
          else:
              pts_geometry_feature = None
              pts_geometry_masks = None
  
          # * sample feature of pts from 2D feature maps
          pts_rendering_feats, pts_rendering_mask = sample_ptsFeatures_from_featureMaps(
              pts, supporting_rendering_feature_maps, supporting_w2cs,
              supporting_intrinsics, img_wh,
              return_mask=True)  # [N_views, C, N_rays, n_samples], # [N_views, N_rays, n_samples]
          # import ipdb; ipdb.set_trace()
          # * size (N_views, N_rays*n_samples, c)
          pts_rendering_feats = pts_rendering_feats.permute(0, 2, 3, 1).contiguous()
  
          pts_rendering_colors = sample_ptsFeatures_from_featureMaps(pts, supporting_color_maps, supporting_w2cs,
                                                                     supporting_intrinsics, img_wh)
          # * size (N_views, N_rays*n_samples, c)
          pts_rendering_colors = pts_rendering_colors.permute(0, 2, 3, 1).contiguous()
  
          rgb_feats = torch.cat([pts_rendering_colors, pts_rendering_feats], dim=-1)  # [N_views, N_rays, n_samples, 3+c]
  
  
          ray_diff = self.compute_angle(pts, query_c2w, supporting_c2ws)  # [N_views, N_rays, n_samples, 4]
          # import ipdb; ipdb.set_trace()
          if pts_geometry_masks is not None:
              final_mask = pts_geometry_masks[None, :, :].repeat(N_supporting_views, 1, 1) & \
                           pts_rendering_mask  # [N_views, N_rays, n_samples]
          else:
              final_mask = pts_rendering_mask
          # import ipdb; ipdb.set_trace()
          z_diff, pts_pred_depth_masks = None, None
          
          if pred_depth_maps is not None:
              pts_pred_depth_values = sample_ptsFeatures_from_featureMaps(pts, supporting_depth_maps, supporting_w2cs,
                                                                          supporting_intrinsics, img_wh)
              pts_pred_depth_values = pts_pred_depth_values.permute(0, 2, 3,
                                                                    1).contiguous()  # (N_views, N_rays*n_samples, 1)
  
              # - pts_pred_depth_masks are critical than final_mask,
              # - the ray containing few invalid pts will be treated invalid
              pts_pred_depth_masks = sample_ptsFeatures_from_featureMaps(pts, supporting_depth_masks.float(),
                                                                         supporting_w2cs,
                                                                         supporting_intrinsics, img_wh)
              
              pts_pred_depth_masks = pts_pred_depth_masks.permute(0, 2, 3, 1).contiguous()[:, :, :,
                                     0]  # (N_views, N_rays*n_samples)
  
              z_diff = self.compute_z_diff(pts, supporting_w2cs, supporting_intrinsics, pts_pred_depth_values)
          # import ipdb; ipdb.set_trace()
          return pts_geometry_feature, rgb_feats, ray_diff, final_mask, z_diff, pts_pred_depth_masks
  
  
      def compute_view_independent(   
                                      self,
                                      pts,
                                      # * 3d geometry feature volumes
                                      geometryVolume=None,
                                      geometryVolumeMask=None,
                                      sdf_network=None,
                                      lod=0,
                                      vol_dims=None,
                                      partial_vol_origin=None,
                                      vol_size=None,
                                      # * 2d rendering feature maps
                                      rendering_feature_maps=None,
                                      color_maps=None,
                                      w2cs=None,
                                      target_candidate_w2cs=None,
                                      intrinsics=None,
                                      img_wh=None,
                                      query_img_idx=0,  # the index of the N_views dim for rendering
                                      query_c2w=None,
                                      pred_depth_maps=None,   # no use here
                                      pred_depth_masks=None   # no use here
                                  ):
          """
          extract features of pts for rendering
          :param pts:
          :param geometryVolume:
          :param vol_dims:
          :param partial_vol_origin:
          :param vol_size:
          :param rendering_feature_maps:
          :param color_maps:
          :param w2cs:
          :param intrinsics:
          :param img_wh:
          :param rendering_img_idx: by default, we render the first view of w2cs
          :return:
          """
          device = pts.device
          c2ws = torch.inverse(w2cs)
  
          if len(pts.shape) == 2:
              pts = pts[None, :, :]
  
          N_rays, n_samples, _ = pts.shape
          N_views = rendering_feature_maps.shape[0]  # shape (N_views, C, H, W)
  
          supporting_img_idxs = torch.LongTensor([x for x in range(N_views) if x != query_img_idx]).to(device)
          query_img_idx = torch.LongTensor([query_img_idx]).to(device)
  
          if query_c2w is None and query_img_idx > -1:
              query_c2w = torch.index_select(c2ws, 0, query_img_idx)
              supporting_c2ws = torch.index_select(c2ws, 0, supporting_img_idxs)
              supporting_w2cs = torch.index_select(w2cs, 0, supporting_img_idxs)
              supporting_rendering_feature_maps = torch.index_select(rendering_feature_maps, 0, supporting_img_idxs)
              supporting_color_maps = torch.index_select(color_maps, 0, supporting_img_idxs)
              supporting_intrinsics = torch.index_select(intrinsics, 0, supporting_img_idxs)
  
              if pred_depth_maps is not None:
                  supporting_depth_maps = torch.index_select(pred_depth_maps, 0, supporting_img_idxs)
                  supporting_depth_masks = torch.index_select(pred_depth_masks, 0, supporting_img_idxs)
              # print("N_supporting_views: ", N_views - 1)
              N_supporting_views = N_views - 1
          else:
              supporting_c2ws = c2ws
              supporting_w2cs = w2cs
              supporting_rendering_feature_maps = rendering_feature_maps
              supporting_color_maps = color_maps
              supporting_intrinsics = intrinsics
              supporting_depth_maps = pred_depth_masks
              supporting_depth_masks = pred_depth_masks
              # print("N_supporting_views: ", N_views)
              N_supporting_views = N_views
          # import ipdb; ipdb.set_trace()
          if geometryVolume is not None:
              # * sample feature of pts from 3D feature volume
              pts_geometry_feature, pts_geometry_masks_0 = sample_ptsFeatures_from_featureVolume(
                  pts, geometryVolume, vol_dims,
                  partial_vol_origin, vol_size)  # [N_rays, n_samples, C], [N_rays, n_samples]
  
              if len(geometryVolumeMask.shape) == 3:
                  geometryVolumeMask = geometryVolumeMask[None, :, :, :]
  
              pts_geometry_masks_1, _ = sample_ptsFeatures_from_featureVolume(
                  pts, geometryVolumeMask.to(geometryVolume.dtype), vol_dims,
                  partial_vol_origin, vol_size)  # [N_rays, n_samples, C]
  
              pts_geometry_masks = pts_geometry_masks_0 & (pts_geometry_masks_1[..., 0] > 0)
          else:
              pts_geometry_feature = None
              pts_geometry_masks = None
  
          # * sample feature of pts from 2D feature maps
          pts_rendering_feats, pts_rendering_mask = sample_ptsFeatures_from_featureMaps(
              pts, supporting_rendering_feature_maps, supporting_w2cs,
              supporting_intrinsics, img_wh,
              return_mask=True)  # [N_views, C, N_rays, n_samples], # [N_views, N_rays, n_samples]
  
          # * size (N_views, N_rays*n_samples, c)
          pts_rendering_feats = pts_rendering_feats.permute(0, 2, 3, 1).contiguous()
  
          pts_rendering_colors = sample_ptsFeatures_from_featureMaps(pts, supporting_color_maps, supporting_w2cs,
                                                                     supporting_intrinsics, img_wh)
          # * size (N_views, N_rays*n_samples, c)
          pts_rendering_colors = pts_rendering_colors.permute(0, 2, 3, 1).contiguous()
  
          rgb_feats = torch.cat([pts_rendering_colors, pts_rendering_feats], dim=-1)  # [N_views, N_rays, n_samples, 3+c]
          
          # import ipdb; ipdb.set_trace()
          
          gradients = sdf_network.gradient(
                                          pts.reshape(-1, 3), # pts.squeeze(0),
                                          geometryVolume.unsqueeze(0), 
                                          lod=lod
                                          ).squeeze()
          
          surface_normals = safe_l2_normalize(gradients, dim=-1) # [npts, 3]
          # input normals
          ren_ray_diff = self.compute_angle_view_independent(
                           xyz=pts,
                           surface_normals=surface_normals,
                           supporting_c2ws=supporting_c2ws   
          )
  
          # # choose closest target view direction from 32 candidate views
          # # choose the closest source view as view direction instead of the normals vectors
          # pts2src_centers = safe_l2_normalize((supporting_c2ws[:, :3, 3].unsqueeze(1) - pts)) # [N_views, npts, 3]
  
          # cosine_distance = torch.sum(pts2src_centers * surface_normals, dim=-1, keepdim=True) # [N_views, npts, 1]
          # # choose the largest cosine distance as the view direction
          # max_idx = torch.argmax(cosine_distance, dim=0) # [npts, 1]
  
          # chosen_view_direction = pts2src_centers[max_idx.squeeze(), torch.arange(pts.shape[1]), :] # [npts, 3]
          # ren_ray_diff = self.compute_angle_view_independent(
          #                  xyz=pts,
          #                  surface_normals=chosen_view_direction,
          #                  supporting_c2ws=supporting_c2ws   
          # )
  
  
  
          # # choose closest target view direction from 8 candidate views
          # # choose the closest source view as view direction instead of the normals vectors
          # target_candidate_c2ws = torch.inverse(target_candidate_w2cs)
          # pts2src_centers = safe_l2_normalize((target_candidate_c2ws[:, :3, 3].unsqueeze(1) - pts)) # [N_views, npts, 3]
  
          # cosine_distance = torch.sum(pts2src_centers * surface_normals, dim=-1, keepdim=True) # [N_views, npts, 1]
          # # choose the largest cosine distance as the view direction
          # max_idx = torch.argmax(cosine_distance, dim=0) # [npts, 1]
  
          # chosen_view_direction = pts2src_centers[max_idx.squeeze(), torch.arange(pts.shape[1]), :] # [npts, 3]
          # ren_ray_diff = self.compute_angle_view_independent(
          #                  xyz=pts,
          #                  surface_normals=chosen_view_direction,
          #                  supporting_c2ws=supporting_c2ws   
          # )
  
  
          # ray_diff = self.compute_angle(pts, query_c2w, supporting_c2ws)  # [N_views, N_rays, n_samples, 4]
          # import ipdb; ipdb.set_trace()
  
  
          # input_directions = safe_l2_normalize(pts)
          # ren_ray_diff = self.compute_angle_view_independent(
          #                  xyz=pts,
          #                  surface_normals=input_directions,
          #                  supporting_c2ws=supporting_c2ws   
          # )
  
          if pts_geometry_masks is not None:
              final_mask = pts_geometry_masks[None, :, :].repeat(N_supporting_views, 1, 1) & \
                           pts_rendering_mask  # [N_views, N_rays, n_samples]
          else:
              final_mask = pts_rendering_mask
          # import ipdb; ipdb.set_trace()
          z_diff, pts_pred_depth_masks = None, None
          
          if pred_depth_maps is not None:
              pts_pred_depth_values = sample_ptsFeatures_from_featureMaps(pts, supporting_depth_maps, supporting_w2cs,
                                                                          supporting_intrinsics, img_wh)
              pts_pred_depth_values = pts_pred_depth_values.permute(0, 2, 3,
                                                                    1).contiguous()  # (N_views, N_rays*n_samples, 1)
  
              # - pts_pred_depth_masks are critical than final_mask,
              # - the ray containing few invalid pts will be treated invalid
              pts_pred_depth_masks = sample_ptsFeatures_from_featureMaps(pts, supporting_depth_masks.float(),
                                                                         supporting_w2cs,
                                                                         supporting_intrinsics, img_wh)
              
              pts_pred_depth_masks = pts_pred_depth_masks.permute(0, 2, 3, 1).contiguous()[:, :, :,
                                     0]  # (N_views, N_rays*n_samples)
  
              z_diff = self.compute_z_diff(pts, supporting_w2cs, supporting_intrinsics, pts_pred_depth_values)
          # import ipdb; ipdb.set_trace()
          return pts_geometry_feature, rgb_feats, ren_ray_diff, final_mask, z_diff, pts_pred_depth_masks
  