content: """
  pytorch grid_sample doesn't support second-order derivative
  implement custom version
  """
  
  import torch
  import torch.nn.functional as F
  import numpy as np
  
  
  def grid_sample_2d(image, optical):
      N, C, IH, IW = image.shape
      _, H, W, _ = optical.shape
  
      ix = optical[..., 0]
      iy = optical[..., 1]
  
      ix = ((ix + 1) / 2) * (IW - 1);
      iy = ((iy + 1) / 2) * (IH - 1);
      with torch.no_grad():
          ix_nw = torch.floor(ix);
          iy_nw = torch.floor(iy);
          ix_ne = ix_nw + 1;
          iy_ne = iy_nw;
          ix_sw = ix_nw;
          iy_sw = iy_nw + 1;
          ix_se = ix_nw + 1;
          iy_se = iy_nw + 1;
  
      nw = (ix_se - ix) * (iy_se - iy)
      ne = (ix - ix_sw) * (iy_sw - iy)
      sw = (ix_ne - ix) * (iy - iy_ne)
      se = (ix - ix_nw) * (iy - iy_nw)
  
      with torch.no_grad():
          torch.clamp(ix_nw, 0, IW - 1, out=ix_nw)
          torch.clamp(iy_nw, 0, IH - 1, out=iy_nw)
  
          torch.clamp(ix_ne, 0, IW - 1, out=ix_ne)
          torch.clamp(iy_ne, 0, IH - 1, out=iy_ne)
  
          torch.clamp(ix_sw, 0, IW - 1, out=ix_sw)
          torch.clamp(iy_sw, 0, IH - 1, out=iy_sw)
  
          torch.clamp(ix_se, 0, IW - 1, out=ix_se)
          torch.clamp(iy_se, 0, IH - 1, out=iy_se)
  
      image = image.view(N, C, IH * IW)
  
      nw_val = torch.gather(image, 2, (iy_nw * IW + ix_nw).long().view(N, 1, H * W).repeat(1, C, 1))
      ne_val = torch.gather(image, 2, (iy_ne * IW + ix_ne).long().view(N, 1, H * W).repeat(1, C, 1))
      sw_val = torch.gather(image, 2, (iy_sw * IW + ix_sw).long().view(N, 1, H * W).repeat(1, C, 1))
      se_val = torch.gather(image, 2, (iy_se * IW + ix_se).long().view(N, 1, H * W).repeat(1, C, 1))
  
      out_val = (nw_val.view(N, C, H, W) * nw.view(N, 1, H, W) +
                 ne_val.view(N, C, H, W) * ne.view(N, 1, H, W) +
                 sw_val.view(N, C, H, W) * sw.view(N, 1, H, W) +
                 se_val.view(N, C, H, W) * se.view(N, 1, H, W))
  
      return out_val
  
  
  # - checked for correctness
  def grid_sample_3d(volume, optical):
      """
      bilinear sampling cannot guarantee continuous first-order gradient
      mimic pytorch grid_sample function
      The 8 corner points of a volume noted as: 4 points (front view); 4 points (back view)
      fnw (front north west) point
      bse (back south east) point
      :param volume: [B, C, X, Y, Z]
      :param optical: [B, x, y, z, 3]
      :return:
      """
      N, C, ID, IH, IW = volume.shape
      _, D, H, W, _ = optical.shape
  
      ix = optical[..., 0]
      iy = optical[..., 1]
      iz = optical[..., 2]
  
      ix = ((ix + 1) / 2) * (IW - 1)
      iy = ((iy + 1) / 2) * (IH - 1)
      iz = ((iz + 1) / 2) * (ID - 1)
  
      mask_x = (ix > 0) & (ix < IW)
      mask_y = (iy > 0) & (iy < IH)
      mask_z = (iz > 0) & (iz < ID)
  
      mask = mask_x & mask_y & mask_z  # [B, x, y, z]
      mask = mask[:, None, :, :, :].repeat(1, C, 1, 1, 1)  # [B, C, x, y, z]
  
      with torch.no_grad():
          # back north west
          ix_bnw = torch.floor(ix)
          iy_bnw = torch.floor(iy)
          iz_bnw = torch.floor(iz)
  
          ix_bne = ix_bnw + 1
          iy_bne = iy_bnw
          iz_bne = iz_bnw
  
          ix_bsw = ix_bnw
          iy_bsw = iy_bnw + 1
          iz_bsw = iz_bnw
  
          ix_bse = ix_bnw + 1
          iy_bse = iy_bnw + 1
          iz_bse = iz_bnw
  
          # front view
          ix_fnw = ix_bnw
          iy_fnw = iy_bnw
          iz_fnw = iz_bnw + 1
  
          ix_fne = ix_bnw + 1
          iy_fne = iy_bnw
          iz_fne = iz_bnw + 1
  
          ix_fsw = ix_bnw
          iy_fsw = iy_bnw + 1
          iz_fsw = iz_bnw + 1
  
          ix_fse = ix_bnw + 1
          iy_fse = iy_bnw + 1
          iz_fse = iz_bnw + 1
  
      # back view
      bnw = (ix_fse - ix) * (iy_fse - iy) * (iz_fse - iz)  # smaller volume, larger weight
      bne = (ix - ix_fsw) * (iy_fsw - iy) * (iz_fsw - iz)
      bsw = (ix_fne - ix) * (iy - iy_fne) * (iz_fne - iz)
      bse = (ix - ix_fnw) * (iy - iy_fnw) * (iz_fnw - iz)
  
      # front view
      fnw = (ix_bse - ix) * (iy_bse - iy) * (iz - iz_bse)  # smaller volume, larger weight
      fne = (ix - ix_bsw) * (iy_bsw - iy) * (iz - iz_bsw)
      fsw = (ix_bne - ix) * (iy - iy_bne) * (iz - iz_bne)
      fse = (ix - ix_bnw) * (iy - iy_bnw) * (iz - iz_bnw)
  
      with torch.no_grad():
          # back view
          torch.clamp(ix_bnw, 0, IW - 1, out=ix_bnw)
          torch.clamp(iy_bnw, 0, IH - 1, out=iy_bnw)
          torch.clamp(iz_bnw, 0, ID - 1, out=iz_bnw)
  
          torch.clamp(ix_bne, 0, IW - 1, out=ix_bne)
          torch.clamp(iy_bne, 0, IH - 1, out=iy_bne)
          torch.clamp(iz_bne, 0, ID - 1, out=iz_bne)
  
          torch.clamp(ix_bsw, 0, IW - 1, out=ix_bsw)
          torch.clamp(iy_bsw, 0, IH - 1, out=iy_bsw)
          torch.clamp(iz_bsw, 0, ID - 1, out=iz_bsw)
  
          torch.clamp(ix_bse, 0, IW - 1, out=ix_bse)
          torch.clamp(iy_bse, 0, IH - 1, out=iy_bse)
          torch.clamp(iz_bse, 0, ID - 1, out=iz_bse)
  
          # front view
          torch.clamp(ix_fnw, 0, IW - 1, out=ix_fnw)
          torch.clamp(iy_fnw, 0, IH - 1, out=iy_fnw)
          torch.clamp(iz_fnw, 0, ID - 1, out=iz_fnw)
  
          torch.clamp(ix_fne, 0, IW - 1, out=ix_fne)
          torch.clamp(iy_fne, 0, IH - 1, out=iy_fne)
          torch.clamp(iz_fne, 0, ID - 1, out=iz_fne)
  
          torch.clamp(ix_fsw, 0, IW - 1, out=ix_fsw)
          torch.clamp(iy_fsw, 0, IH - 1, out=iy_fsw)
          torch.clamp(iz_fsw, 0, ID - 1, out=iz_fsw)
  
          torch.clamp(ix_fse, 0, IW - 1, out=ix_fse)
          torch.clamp(iy_fse, 0, IH - 1, out=iy_fse)
          torch.clamp(iz_fse, 0, ID - 1, out=iz_fse)
  
      # xxx = volume[:, :, iz_bnw.long(), iy_bnw.long(), ix_bnw.long()]
      volume = volume.view(N, C, ID * IH * IW)
      # yyy = volume[:, :, (iz_bnw * ID + iy_bnw * IW + ix_bnw).long()]
  
      # back view
      bnw_val = torch.gather(volume, 2,
                             (iz_bnw * ID ** 2 + iy_bnw * IW + ix_bnw).long().view(N, 1, D * H * W).repeat(1, C, 1))
      bne_val = torch.gather(volume, 2,
                             (iz_bne * ID ** 2 + iy_bne * IW + ix_bne).long().view(N, 1, D * H * W).repeat(1, C, 1))
      bsw_val = torch.gather(volume, 2,
                             (iz_bsw * ID ** 2 + iy_bsw * IW + ix_bsw).long().view(N, 1, D * H * W).repeat(1, C, 1))
      bse_val = torch.gather(volume, 2,
                             (iz_bse * ID ** 2 + iy_bse * IW + ix_bse).long().view(N, 1, D * H * W).repeat(1, C, 1))
  
      # front view
      fnw_val = torch.gather(volume, 2,
                             (iz_fnw * ID ** 2 + iy_fnw * IW + ix_fnw).long().view(N, 1, D * H * W).repeat(1, C, 1))
      fne_val = torch.gather(volume, 2,
                             (iz_fne * ID ** 2 + iy_fne * IW + ix_fne).long().view(N, 1, D * H * W).repeat(1, C, 1))
      fsw_val = torch.gather(volume, 2,
                             (iz_fsw * ID ** 2 + iy_fsw * IW + ix_fsw).long().view(N, 1, D * H * W).repeat(1, C, 1))
      fse_val = torch.gather(volume, 2,
                             (iz_fse * ID ** 2 + iy_fse * IW + ix_fse).long().view(N, 1, D * H * W).repeat(1, C, 1))
  
      out_val = (
          # back
              bnw_val.view(N, C, D, H, W) * bnw.view(N, 1, D, H, W) +
              bne_val.view(N, C, D, H, W) * bne.view(N, 1, D, H, W) +
              bsw_val.view(N, C, D, H, W) * bsw.view(N, 1, D, H, W) +
              bse_val.view(N, C, D, H, W) * bse.view(N, 1, D, H, W) +
              # front
              fnw_val.view(N, C, D, H, W) * fnw.view(N, 1, D, H, W) +
              fne_val.view(N, C, D, H, W) * fne.view(N, 1, D, H, W) +
              fsw_val.view(N, C, D, H, W) * fsw.view(N, 1, D, H, W) +
              fse_val.view(N, C, D, H, W) * fse.view(N, 1, D, H, W)
  
      )
  
      # * zero padding
      out_val = torch.where(mask, out_val, torch.zeros_like(out_val).float().to(out_val.device))
  
      return out_val
  
  
  # Interpolation kernel
  def get_weight(s, a=-0.5):
      mask_0 = (torch.abs(s) >= 0) & (torch.abs(s) <= 1)
      mask_1 = (torch.abs(s) > 1) & (torch.abs(s) <= 2)
      mask_2 = torch.abs(s) > 2
  
      weight = torch.zeros_like(s).to(s.device)
      weight = torch.where(mask_0, (a + 2) * (torch.abs(s) ** 3) - (a + 3) * (torch.abs(s) ** 2) + 1, weight)
      weight = torch.where(mask_1,
                           a * (torch.abs(s) ** 3) - (5 * a) * (torch.abs(s) ** 2) + (8 * a) * torch.abs(s) - 4 * a,
                           weight)
  
      # if (torch.abs(s) >= 0) & (torch.abs(s) <= 1):
      #     return (a + 2) * (torch.abs(s) ** 3) - (a + 3) * (torch.abs(s) ** 2) + 1
      #
      # elif (torch.abs(s) > 1) & (torch.abs(s) <= 2):
      #     return a * (torch.abs(s) ** 3) - (5 * a) * (torch.abs(s) ** 2) + (8 * a) * torch.abs(s) - 4 * a
      # return 0
  
      return weight
  
  
  def cubic_interpolate(p, x):
      """
      one dimensional cubic interpolation
      :param p: [N, 4]  (4) should be in order
      :param x: [N]
      :return:
      """
      return p[:, 1] + 0.5 * x * (p[:, 2] - p[:, 0] + x * (
              2.0 * p[:, 0] - 5.0 * p[:, 1] + 4.0 * p[:, 2] - p[:, 3] + x * (
              3.0 * (p[:, 1] - p[:, 2]) + p[:, 3] - p[:, 0])))
  
  
  def bicubic_interpolate(p, x, y, if_batch=True):
      """
      two dimensional cubic interpolation
      :param p: [N, 4, 4]
      :param x: [N]
      :param y: [N]
      :return:
      """
      num = p.shape[0]
  
      if not if_batch:
          arr0 = cubic_interpolate(p[:, 0, :], x)  # [N]
          arr1 = cubic_interpolate(p[:, 1, :], x)
          arr2 = cubic_interpolate(p[:, 2, :], x)
          arr3 = cubic_interpolate(p[:, 3, :], x)
          return cubic_interpolate(torch.stack([arr0, arr1, arr2, arr3], dim=-1), y)  # [N]
      else:
          x = x[:, None].repeat(1, 4).view(-1)
          p = p.contiguous().view(num * 4, 4)
          arr = cubic_interpolate(p, x)
          arr = arr.view(num, 4)
  
          return cubic_interpolate(arr, y)
  
  
  def tricubic_interpolate(p, x, y, z):
      """
      three dimensional cubic interpolation
      :param p: [N,4,4,4]
      :param x: [N]
      :param y: [N]
      :param z: [N]
      :return:
      """
      num = p.shape[0]
  
      arr0 = bicubic_interpolate(p[:, 0, :, :], x, y)  # [N]
      arr1 = bicubic_interpolate(p[:, 1, :, :], x, y)
      arr2 = bicubic_interpolate(p[:, 2, :, :], x, y)
      arr3 = bicubic_interpolate(p[:, 3, :, :], x, y)
  
      return cubic_interpolate(torch.stack([arr0, arr1, arr2, arr3], dim=-1), z)  # [N]
  
  
  def cubic_interpolate_batch(p, x):
      """
      one dimensional cubic interpolation
      :param p: [B, N, 4]  (4) should be in order
      :param x: [B, N]
      :return:
      """
      return p[:, :, 1] + 0.5 * x * (p[:, :, 2] - p[:, :, 0] + x * (
              2.0 * p[:, :, 0] - 5.0 * p[:, :, 1] + 4.0 * p[:, :, 2] - p[:, :, 3] + x * (
              3.0 * (p[:, :, 1] - p[:, :, 2]) + p[:, :, 3] - p[:, :, 0])))
  
  
  def bicubic_interpolate_batch(p, x, y):
      """
      two dimensional cubic interpolation
      :param p: [B, N, 4, 4]
      :param x: [B, N]
      :param y: [B, N]
      :return:
      """
      B, N, _, _ = p.shape
  
      x = x[:, :, None].repeat(1, 1, 4).view(B, N * 4)  # [B, N*4]
      arr = cubic_interpolate_batch(p.contiguous().view(B, N * 4, 4), x)
      arr = arr.view(B, N, 4)
      return cubic_interpolate_batch(arr, y)  # [B, N]
  
  
  # * batch version cannot speed up training
  def tricubic_interpolate_batch(p, x, y, z):
      """
      three dimensional cubic interpolation
      :param p: [N,4,4,4]
      :param x: [N]
      :param y: [N]
      :param z: [N]
      :return:
      """
      N = p.shape[0]
  
      x = x[None, :].repeat(4, 1)
      y = y[None, :].repeat(4, 1)
  
      p = p.permute(1, 0, 2, 3).contiguous()
  
      arr = bicubic_interpolate_batch(p[:, :, :, :], x, y)  # [4, N]
  
      arr = arr.permute(1, 0).contiguous()  # [N, 4]
  
      return cubic_interpolate(arr, z)  # [N]
  
  
  def tricubic_sample_3d(volume, optical):
      """
      tricubic sampling; can guarantee continuous gradient  (interpolation border)
      :param volume: [B, C, ID, IH, IW]
      :param optical: [B, D, H, W, 3]
      :param sample_num:
      :return:
      """
  
      @torch.no_grad()
      def get_shifts(x):
          x1 = -1 * (1 + x - torch.floor(x))
          x2 = -1 * (x - torch.floor(x))
          x3 = torch.floor(x) + 1 - x
          x4 = torch.floor(x) + 2 - x
  
          return torch.stack([x1, x2, x3, x4], dim=-1)  # (B,d,h,w,4)
  
      N, C, ID, IH, IW = volume.shape
      _, D, H, W, _ = optical.shape
  
      device = volume.device
  
      ix = optical[..., 0]
      iy = optical[..., 1]
      iz = optical[..., 2]
  
      ix = ((ix + 1) / 2) * (IW - 1)  # (B,d,h,w)
      iy = ((iy + 1) / 2) * (IH - 1)
      iz = ((iz + 1) / 2) * (ID - 1)
  
      ix = ix.view(-1)
      iy = iy.view(-1)
      iz = iz.view(-1)
  
      with torch.no_grad():
          shifts_x = get_shifts(ix).view(-1, 4)  # (B*d*h*w,4)
          shifts_y = get_shifts(iy).view(-1, 4)
          shifts_z = get_shifts(iz).view(-1, 4)
  
          perm_weights = torch.ones([N * D * H * W, 4 * 4 * 4]).long().to(device)
          perm = torch.cumsum(perm_weights, dim=-1) - 1  # (B*d*h*w,64)
  
          perm_z = perm // 16  # [N*D*H*W, num]
          perm_y = (perm - perm_z * 16) // 4
          perm_x = (perm - perm_z * 16 - perm_y * 4)
  
          shifts_x = torch.gather(shifts_x, 1, perm_x)  # [N*D*H*W, num]
          shifts_y = torch.gather(shifts_y, 1, perm_y)
          shifts_z = torch.gather(shifts_z, 1, perm_z)
  
          ix_target = (ix[:, None] + shifts_x).long()  # [N*D*H*W, num]
          iy_target = (iy[:, None] + shifts_y).long()
          iz_target = (iz[:, None] + shifts_z).long()
  
          torch.clamp(ix_target, 0, IW - 1, out=ix_target)
          torch.clamp(iy_target, 0, IH - 1, out=iy_target)
          torch.clamp(iz_target, 0, ID - 1, out=iz_target)
  
      local_dist_x = ix - ix_target[:, 1]  # ! attention here is [:, 1]
      local_dist_y = iy - iy_target[:, 1 + 4]
      local_dist_z = iz - iz_target[:, 1 + 16]
  
      local_dist_x = local_dist_x.view(N, 1, D * H * W).repeat(1, C, 1).view(-1)
      local_dist_y = local_dist_y.view(N, 1, D * H * W).repeat(1, C, 1).view(-1)
      local_dist_z = local_dist_z.view(N, 1, D * H * W).repeat(1, C, 1).view(-1)
  
      # ! attention: IW is correct
      idx_target = iz_target * ID ** 2 + iy_target * IW + ix_target  # [N*D*H*W, num]
  
      volume = volume.view(N, C, ID * IH * IW)
  
      out = torch.gather(volume, 2,
                         idx_target.view(N, 1, D * H * W * 64).repeat(1, C, 1))
      out = out.view(N * C * D * H * W, 4, 4, 4)
  
      # - tricubic_interpolate() is a bit faster than tricubic_interpolate_batch()
      final = tricubic_interpolate(out, local_dist_x, local_dist_y, local_dist_z).view(N, C, D, H, W)  # [N,C,D,H,W]
  
      return final
  
  
  
  if __name__ == "__main__":
      # image = torch.Tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]).view(1, 3, 1, 3)
      #
      # optical = torch.Tensor([0.9, 0.5, 0.6, -0.7]).view(1, 1, 2, 2)
      #
      # print(grid_sample_2d(image, optical))
      #
      # print(F.grid_sample(image, optical, padding_mode='border', align_corners=True))
  
      from ops.generate_grids import generate_grid
  
      p = torch.tensor([x for x in range(4)]).view(1, 4).float()
  
      v = cubic_interpolate(p, torch.tensor([0.5]).view(1))
      # v = bicubic_interpolate(p, torch.tensor([2/3]).view(1) , torch.tensor([2/3]).view(1))
  
      vsize = 9
      volume = generate_grid([vsize, vsize, vsize], 1)  # [1,3,10,10,10]
      # volume = torch.tensor([x for x in range(1000)]).view(1, 1, 10, 10, 10).float()
      X, Y, Z = 0, 0, 6
      x = 2 * X / (vsize - 1) - 1
      y = 2 * Y / (vsize - 1) - 1
      z = 2 * Z / (vsize - 1) - 1
  
      # print(volume[:, :, Z, Y, X])
  
      # volume = volume.view(1, 3, -1)
      # xx = volume[:, :, Z * 9*9 + Y * 9 + X]
  
      optical = torch.Tensor([-0.6, -0.7, 0.5, 0.3, 0.5, 0.5]).view(1, 1, 1, 2, 3)
  
      print(F.grid_sample(volume, optical, padding_mode='border', align_corners=True))
      print(grid_sample_3d(volume, optical))
      print(tricubic_sample_3d(volume, optical))
      # target, relative_coords = implicit_sample_3d(volume, optical, 1)
      # print(target)
  