/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *  * Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 *  * Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *  * Neither the name of NVIDIA CORPORATION nor the names of its
 *    contributors may be used to endorse or promote products derived
 *    from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
 * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

/*
    Bicubic filtering
    See GPU Gems 2: "Fast Third-Order Texture Filtering", Sigg & Hadwiger
    https://developer.nvidia.com/gpugems/gpugems2/part-iii-high-quality-rendering/chapter-20-fast-third-order-texture-filtering

    Reformulation thanks to Keenan Crane
*/

#ifndef _BICUBICTEXTURE_KERNEL_CUH_
#define _BICUBICTEXTURE_KERNEL_CUH_

enum Mode {
  MODE_NEAREST,
  MODE_BILINEAR,
  MODE_BICUBIC,
  MODE_FAST_BICUBIC,
  MODE_CATROM
};

cudaTextureObject_t texObjPoint, texObjLinear;

// w0, w1, w2, and w3 are the four cubic B-spline basis functions
__host__ __device__ float w0(float a) {
  //    return (1.0f/6.0f)*(-a*a*a + 3.0f*a*a - 3.0f*a + 1.0f);
  return (1.0f / 6.0f) * (a * (a * (-a + 3.0f) - 3.0f) + 1.0f);  // optimized
}

__host__ __device__ float w1(float a) {
  //    return (1.0f/6.0f)*(3.0f*a*a*a - 6.0f*a*a + 4.0f);
  return (1.0f / 6.0f) * (a * a * (3.0f * a - 6.0f) + 4.0f);
}

__host__ __device__ float w2(float a) {
  //    return (1.0f/6.0f)*(-3.0f*a*a*a + 3.0f*a*a + 3.0f*a + 1.0f);
  return (1.0f / 6.0f) * (a * (a * (-3.0f * a + 3.0f) + 3.0f) + 1.0f);
}

__host__ __device__ float w3(float a) { return (1.0f / 6.0f) * (a * a * a); }

// g0 and g1 are the two amplitude functions
__device__ float g0(float a) { return w0(a) + w1(a); }

__device__ float g1(float a) { return w2(a) + w3(a); }

// h0 and h1 are the two offset functions
__device__ float h0(float a) {
  // note +0.5 offset to compensate for CUDA linear filtering convention
  return -1.0f + w1(a) / (w0(a) + w1(a)) + 0.5f;
}

__device__ float h1(float a) { return 1.0f + w3(a) / (w2(a) + w3(a)) + 0.5f; }

// filter 4 values using cubic splines
template <class T>
__device__ T cubicFilter(float x, T c0, T c1, T c2, T c3) {
  T r;
  r = c0 * w0(x);
  r += c1 * w1(x);
  r += c2 * w2(x);
  r += c3 * w3(x);
  return r;
}

// slow but precise bicubic lookup using 16 texture lookups
template <class T, class R>  // texture data type, return type
__device__ R tex2DBicubic(const cudaTextureObject_t tex, float x, float y) {
  x -= 0.5f;
  y -= 0.5f;
  float px = floorf(x);
  float py = floorf(y);
  float fx = x - px;
  float fy = y - py;

  return cubicFilter<R>(
      fy, cubicFilter<R>(
              fx, tex2D<R>(tex, px - 1, py - 1), tex2D<R>(tex, px, py - 1),
              tex2D<R>(tex, px + 1, py - 1), tex2D<R>(tex, px + 2, py - 1)),
      cubicFilter<R>(fx, tex2D<R>(tex, px - 1, py), tex2D<R>(tex, px, py),
                     tex2D<R>(tex, px + 1, py), tex2D<R>(tex, px + 2, py)),
      cubicFilter<R>(fx, tex2D<R>(tex, px - 1, py + 1),
                     tex2D<R>(tex, px, py + 1), tex2D<R>(tex, px + 1, py + 1),
                     tex2D<R>(tex, px + 2, py + 1)),
      cubicFilter<R>(fx, tex2D<R>(tex, px - 1, py + 2),
                     tex2D<R>(tex, px, py + 2), tex2D<R>(tex, px + 1, py + 2),
                     tex2D<R>(tex, px + 2, py + 2)));
}

// fast bicubic texture lookup using 4 bilinear lookups
// assumes texture is set to non-normalized coordinates, point sampling
template <class T, class R>  // texture data type, return type
__device__ R tex2DFastBicubic(const cudaTextureObject_t tex, float x, float y) {
  x -= 0.5f;
  y -= 0.5f;
  float px = floorf(x);
  float py = floorf(y);
  float fx = x - px;
  float fy = y - py;

  // note: we could store these functions in a lookup table texture, but maths
  // is cheap
  float g0x = g0(fx);
  float g1x = g1(fx);
  float h0x = h0(fx);
  float h1x = h1(fx);
  float h0y = h0(fy);
  float h1y = h1(fy);

  R r = g0(fy) * (g0x * tex2D<R>(tex, px + h0x, py + h0y) +
                  g1x * tex2D<R>(tex, px + h1x, py + h0y)) +
        g1(fy) * (g0x * tex2D<R>(tex, px + h0x, py + h1y) +
                  g1x * tex2D<R>(tex, px + h1x, py + h1y));
  return r;
}

// higher-precision 2D bilinear lookup
template <class T, class R>  // texture data type, return type
__device__ R tex2DBilinear(const cudaTextureObject_t tex, float x, float y) {
  x -= 0.5f;
  y -= 0.5f;
  float px = floorf(x);  // integer position
  float py = floorf(y);
  float fx = x - px;  // fractional position
  float fy = y - py;
  px += 0.5f;
  py += 0.5f;

  return lerp(lerp(tex2D<R>(tex, px, py), tex2D<R>(tex, px + 1.0f, py), fx),
              lerp(tex2D<R>(tex, px, py + 1.0f),
                   tex2D<R>(tex, px + 1.0f, py + 1.0f), fx),
              fy);
}

#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 200

/*
  bilinear 2D texture lookup using tex2Dgather function
  - tex2Dgather() returns the four neighbouring samples in a single texture
   lookup
  - it is only supported on the Fermi architecture
  - you can select which component to fetch using the "comp" parameter
  - it can be used to efficiently implement custom texture filters

  The samples are returned in a 4-vector in the following order:
  x: (0, 1)
  y: (1, 1)
  z: (1, 0)
  w: (0, 0)
*/

template <class T, class R>  // texture data type, return type
__device__ float tex2DBilinearGather(const cudaTextureObject_t tex, float x,
                                     float y, int comp = 0) {
  x -= 0.5f;
  y -= 0.5f;
  float px = floorf(x);  // integer position
  float py = floorf(y);
  float fx = x - px;  // fractional position
  float fy = y - py;

  R samples = tex2Dgather<R>(tex, px + 0.5f, py + 0.5f, comp);

  return lerp(lerp((float)samples.w, (float)samples.z, fx),
              lerp((float)samples.x, (float)samples.y, fx), fy);
}

#endif

// Catmull-Rom interpolation

__host__ __device__ float catrom_w0(float a) {
  // return -0.5f*a + a*a - 0.5f*a*a*a;
  return a * (-0.5f + a * (1.0f - 0.5f * a));
}

__host__ __device__ float catrom_w1(float a) {
  // return 1.0f - 2.5f*a*a + 1.5f*a*a*a;
  return 1.0f + a * a * (-2.5f + 1.5f * a);
}

__host__ __device__ float catrom_w2(float a) {
  // return 0.5f*a + 2.0f*a*a - 1.5f*a*a*a;
  return a * (0.5f + a * (2.0f - 1.5f * a));
}

__host__ __device__ float catrom_w3(float a) {
  // return -0.5f*a*a + 0.5f*a*a*a;
  return a * a * (-0.5f + 0.5f * a);
}

template <class T>
__device__ T catRomFilter(float x, T c0, T c1, T c2, T c3) {
  T r;
  r = c0 * catrom_w0(x);
  r += c1 * catrom_w1(x);
  r += c2 * catrom_w2(x);
  r += c3 * catrom_w3(x);
  return r;
}

// Note - can't use bilinear trick here because of negative lobes
template <class T, class R>  // texture data type, return type
__device__ R tex2DCatRom(const cudaTextureObject_t tex, float x, float y) {
  x -= 0.5f;
  y -= 0.5f;
  float px = floorf(x);
  float py = floorf(y);
  float fx = x - px;
  float fy = y - py;

  return catRomFilter<R>(
      fy, catRomFilter<R>(
              fx, tex2D<R>(tex, px - 1, py - 1), tex2D<R>(tex, px, py - 1),
              tex2D<R>(tex, px + 1, py - 1), tex2D<R>(tex, px + 2, py - 1)),
      catRomFilter<R>(fx, tex2D<R>(tex, px - 1, py), tex2D<R>(tex, px, py),
                      tex2D<R>(tex, px + 1, py), tex2D<R>(tex, px + 2, py)),
      catRomFilter<R>(fx, tex2D<R>(tex, px - 1, py + 1),
                      tex2D<R>(tex, px, py + 1), tex2D<R>(tex, px + 1, py + 1),
                      tex2D<R>(tex, px + 2, py + 1)),
      catRomFilter<R>(fx, tex2D<R>(tex, px - 1, py + 2),
                      tex2D<R>(tex, px, py + 2), tex2D<R>(tex, px + 1, py + 2),
                      tex2D<R>(tex, px + 2, py + 2)));
}

// test functions

// render image using normal bilinear texture lookup
__global__ void d_render(uchar4 *d_output, uint width, uint height, float tx,
                         float ty, float scale, float cx, float cy,
                         cudaTextureObject_t texObj) {
  uint x = __umul24(blockIdx.x, blockDim.x) + threadIdx.x;
  uint y = __umul24(blockIdx.y, blockDim.y) + threadIdx.y;
  uint i = __umul24(y, width) + x;

  float u = (x - cx) * scale + cx + tx;
  float v = (y - cy) * scale + cy + ty;

  if ((x < width) && (y < height)) {
    // write output color
    float c = tex2D<float>(texObj, u, v);
    // float c = tex2DBilinear<uchar, float>(tex, u, v);
    // float c = tex2DBilinearGather<uchar, uchar4>(tex2, u, v, 0) / 255.0f;
    d_output[i] = make_uchar4(c * 0xff, c * 0xff, c * 0xff, 0);
  }
}

// render image using bicubic texture lookup
__global__ void d_renderBicubic(uchar4 *d_output, uint width, uint height,
                                float tx, float ty, float scale, float cx,
                                float cy, cudaTextureObject_t texObj) {
  uint x = __umul24(blockIdx.x, blockDim.x) + threadIdx.x;
  uint y = __umul24(blockIdx.y, blockDim.y) + threadIdx.y;
  uint i = __umul24(y, width) + x;

  float u = (x - cx) * scale + cx + tx;
  float v = (y - cy) * scale + cy + ty;

  if ((x < width) && (y < height)) {
    // write output color
    float c = tex2DBicubic<uchar, float>(texObj, u, v);
    d_output[i] = make_uchar4(c * 0xff, c * 0xff, c * 0xff, 0);
  }
}

// render image using fast bicubic texture lookup
__global__ void d_renderFastBicubic(uchar4 *d_output, uint width, uint height,
                                    float tx, float ty, float scale, float cx,
                                    float cy, cudaTextureObject_t texObj) {
  uint x = __umul24(blockIdx.x, blockDim.x) + threadIdx.x;
  uint y = __umul24(blockIdx.y, blockDim.y) + threadIdx.y;
  uint i = __umul24(y, width) + x;

  float u = (x - cx) * scale + cx + tx;
  float v = (y - cy) * scale + cy + ty;

  if ((x < width) && (y < height)) {
    // write output color
    float c = tex2DFastBicubic<uchar, float>(texObj, u, v);
    d_output[i] = make_uchar4(c * 0xff, c * 0xff, c * 0xff, 0);
  }
}

// render image using Catmull-Rom texture lookup
__global__ void d_renderCatRom(uchar4 *d_output, uint width, uint height,
                               float tx, float ty, float scale, float cx,
                               float cy, cudaTextureObject_t texObj) {
  uint x = __umul24(blockIdx.x, blockDim.x) + threadIdx.x;
  uint y = __umul24(blockIdx.y, blockDim.y) + threadIdx.y;
  uint i = __umul24(y, width) + x;

  float u = (x - cx) * scale + cx + tx;
  float v = (y - cy) * scale + cy + ty;

  if ((x < width) && (y < height)) {
    // write output color
    float c = tex2DCatRom<uchar, float>(texObj, u, v);
    d_output[i] = make_uchar4(c * 0xff, c * 0xff, c * 0xff, 0);
  }
}

#endif  // _BICUBICTEXTURE_KERNEL_CUH_