/* SPDX-FileCopyrightText: 2017-2022 Blender Foundation
 *
 * SPDX-License-Identifier: Apache-2.0 */

#pragma once

#include "kernel/globals.h"
#include "kernel/util/image_2d.h"
#include "util/defines.h"
#include "util/types_image.h"

CCL_NAMESPACE_BEGIN

ccl_device_inline float frac(const float x, ccl_private int *ix)
{
  int i = float_to_int(x) - ((x < 0.0f) ? 1 : 0);
  *ix = i;
  return x - (float)i;
}

/* w0, w1, w2, and w3 are the four cubic B-spline basis functions. */
ccl_device float cubic_w0(const float a)
{
  return (1.0f / 6.0f) * (a * (a * (-a + 3.0f) - 3.0f) + 1.0f);
}
ccl_device float cubic_w1(const float a)
{
  return (1.0f / 6.0f) * (a * a * (3.0f * a - 6.0f) + 4.0f);
}
ccl_device float cubic_w2(const float a)
{
  return (1.0f / 6.0f) * (a * (a * (-3.0f * a + 3.0f) + 3.0f) + 1.0f);
}
ccl_device float cubic_w3(const float a)
{
  return (1.0f / 6.0f) * (a * a * a);
}

/* g0 and g1 are the two amplitude functions. */
ccl_device float cubic_g0(const float a)
{
  return cubic_w0(a) + cubic_w1(a);
}
ccl_device float cubic_g1(const float a)
{
  return cubic_w2(a) + cubic_w3(a);
}

/* h0 and h1 are the two offset functions */
ccl_device float cubic_h0(const float a)
{
  return (cubic_w1(a) / cubic_g0(a)) - 1.0f;
}
ccl_device float cubic_h1(const float a)
{
  return (cubic_w3(a) / cubic_g1(a)) + 1.0f;
}

/* Fast bicubic texture lookup using 4 bilinear lookups, adapted from CUDA samples. */
template<typename T>
/* Workaround for an apparent compiler bug with CUDA 12.8 and Blackwell.
 * This appears to be fixed with CUDA 12.9, so this workaround can be removed
 * when the minimum toolkit version is increased. */
#if defined(__KERNEL_CUDA__) && (__CUDA_ARCH__ >= 1000)
ccl_device_inline
#else
ccl_device_noinline
#endif
T kernel_image_interp_bicubic(const ccl_global KernelImageInfo &info,
                                                const float2 uv)
{
  ccl_gpu_image_object_2D tex = (ccl_gpu_image_object_2D)info.data;

  const float x = (uv.x * (float)info.width) - 0.5f;
  const float y = (uv.y * (float)info.height) - 0.5f;

  const float px = floorf(x);
  const float py = floorf(y);
  const float fx = x - px;
  const float fy = y - py;

  const float g0x = cubic_g0(fx);
  const float g1x = cubic_g1(fx);
  /* Note +0.5 offset to compensate for CUDA linear filtering convention. */
  const float x0 = (px + cubic_h0(fx) + 0.5f) * info.inv_width;
  const float x1 = (px + cubic_h1(fx) + 0.5f) * info.inv_width;
  const float y0 = (py + cubic_h0(fy) + 0.5f) * info.inv_height;
  const float y1 = (py + cubic_h1(fy) + 0.5f) * info.inv_height;

  return cubic_g0(fy) * (g0x * ccl_gpu_image_object_read_2D<T>(tex, x0, y0) +
                         g1x * ccl_gpu_image_object_read_2D<T>(tex, x1, y0)) +
         cubic_g1(fy) * (g0x * ccl_gpu_image_object_read_2D<T>(tex, x0, y1) +
                         g1x * ccl_gpu_image_object_read_2D<T>(tex, x1, y1));
}

ccl_device float4 kernel_image_interp(KernelGlobals kg,
                                      ccl_private ShaderData *sd,
                                      const int image_texture_id,
                                      dual2 uv)
{
  if (image_texture_id == KERNEL_IMAGE_NONE) {
    return IMAGE_MISSING_RGBA;
  }
  const ccl_global KernelImageTexture &tex = kernel_data_fetch(image_textures, image_texture_id);
  const ccl_global KernelImageInfo *info;

  float2 sample_uv;

  if (tex.tile_descriptor_offset != KERNEL_TILE_LOAD_NONE) {
    /* Wrapping. */
    if (!kernel_image_tile_wrap(ExtensionType(tex.extension), uv.val)) {
      return zero_float4();
    }

    /* Tile mapping */
    float2 xy = zero_float2();
    const KernelTileDescriptor tile_descriptor = kernel_image_tile_map(
        kg, sd, tex, image_texture_id, uv, xy);

    if (!kernel_tile_descriptor_loaded(tile_descriptor)) {
      return (tile_descriptor == KERNEL_TILE_LOAD_FAILED) ? IMAGE_MISSING_RGBA : tex.average_color;
    }

    info = &kernel_data_fetch(image_info, kernel_tile_descriptor_image_info_id(tile_descriptor));

    /* Convert to normalized space again. */
    sample_uv = make_float2(xy.x * info->inv_width, xy.y * info->inv_height);
  }
  else {
    /* Full image sampling. */
    if (tex.image_info_id == KERNEL_IMAGE_NONE) {
      return IMAGE_MISSING_RGBA;
    }

    info = &kernel_data_fetch(image_info, tex.image_info_id);
    sample_uv = uv.val;
  }

  /* float4, byte4, ushort4 and half4 */
  const int texture_type = info->data_type;
  if (texture_type == IMAGE_DATA_TYPE_FLOAT4 || texture_type == IMAGE_DATA_TYPE_BYTE4 ||
      texture_type == IMAGE_DATA_TYPE_HALF4 || texture_type == IMAGE_DATA_TYPE_USHORT4)
  {
    if (info->interpolation == INTERPOLATION_CUBIC || info->interpolation == INTERPOLATION_SMART) {
      return kernel_image_interp_bicubic<float4>(*info, sample_uv);
    }
    else {
      ccl_gpu_image_object_2D tex = (ccl_gpu_image_object_2D)info->data;
      return ccl_gpu_image_object_read_2D<float4>(tex, sample_uv.x, sample_uv.y);
    }
  }
  /* float, byte and half */
  else {
    float f;

    if (info->interpolation == INTERPOLATION_CUBIC || info->interpolation == INTERPOLATION_SMART) {
      f = kernel_image_interp_bicubic<float>(*info, sample_uv);
    }
    else {
      ccl_gpu_image_object_2D tex = (ccl_gpu_image_object_2D)info->data;
      f = ccl_gpu_image_object_read_2D<float>(tex, sample_uv.x, sample_uv.y);
    }

    return make_float4(f, f, f, 1.0f);
  }
}

ccl_device_forceinline float4 kernel_image_interp_with_udim(KernelGlobals kg,
                                                            ccl_private ShaderData *sd,
                                                            const int udim_id,
                                                            dual2 uv)
{
  const int image_texture_id = kernel_image_udim_map(kg, udim_id, uv.val);
  if (image_texture_id == KERNEL_IMAGE_NONE) {
    return IMAGE_MISSING_RGBA;
  }

  return kernel_image_interp(kg, sd, image_texture_id, uv);
}

CCL_NAMESPACE_END
