/*
    This file is part of darktable,
    copyright (c) 2011-2026 darktable developers.

    darktable is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    darktable is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with darktable.  If not, see <http://www.gnu.org/licenses/>.
*/

#include "common.h"

/*
    To speed up processing we use an algorithm proposed by B. Goossens, H.Q. Luong, J. Aelterman, A. Pizurica,  and W. Philips,
    "A GPU-Accelerated Real-Time NLMeans Algorithm for Denoising Color Video Sequences", in Proc. ACIVS (2), 2010, pp.46-57.
*/

static inline float ddirac(const int2 q)
{
  return ((q.x || q.y) ? 1.0f : 0.0f);
}

kernel void
denoiseprofile_precondition(read_only image2d_t in,
                            write_only image2d_t out,
                            const int width,
                            const int height,
                            const float4 a,
                            const float4 sigma2)
{
  const int x = get_global_id(0);
  const int y = get_global_id(1);

  if(x >= width || y >= height) return;

  float4 pixel = readpixel(in, x, y);
  const float alpha = pixel.w;

  float4 t = fmax(pixel / a, 0.f);
  float4 d = fmax((float4)0.0f, t + (float4)0.375f + sigma2);
  float4 s = 2.0f * dtcl_sqrt(d);

  s.w = alpha;
  write_imagef (out, (int2)(x, y), s);
}


kernel void
denoiseprofile_precondition_v2(read_only image2d_t in,
                               write_only image2d_t out,
                               const int width,
                               const int height,
                               const float4 a,
                               const float4 p,
                               const float4 b,
                               const float4 wb)
{
  const int x = get_global_id(0);
  const int y = get_global_id(1);

  if(x >= width || y >= height) return;

  float4 pixel = readpixel(in, x, y);
  const float alpha = pixel.w;

  float4 t = fmax(2.0f * dtcl_pow(fmax((float4)0.0f, pixel / wb + b), 1.0f - p / 2.0f) / ((-p + 2.0f) * dtcl_sqrt(a)), 0.f);

  t.w = alpha;

  write_imagef (out, (int2)(x, y), t);
}

kernel void denoiseprofile_precondition_Y0U0V0(read_only image2d_t in,
                                               write_only image2d_t out,
                                               const int width,
                                               const int height,
                                               const float4 a,
                                               const float4 p,
                                               const float4 b,
                                               global float *toY0U0V0)
{
  const int x = get_global_id(0);
  const int y = get_global_id(1);

  if(x >= width || y >= height) return;

  float4 pixel = readpixel(in, x, y);
  const float alpha = pixel.w;

  const float4 t = fmax(2.0f * dtcl_pow(fmax((float4)0.0f, pixel + b), 1.0f - p / 2.0f) / ((-p + 2.0f) * dtcl_sqrt(a)), 0.f);

  float4 outpx = (float4)0.0f;
  outpx.x += toY0U0V0[0] * t.x;
  outpx.x += toY0U0V0[1] * t.y;
  outpx.x += toY0U0V0[2] * t.z;
  outpx.y += toY0U0V0[3] * t.x;
  outpx.y += toY0U0V0[4] * t.y;
  outpx.y += toY0U0V0[5] * t.z;
  outpx.z += toY0U0V0[6] * t.x;
  outpx.z += toY0U0V0[7] * t.y;
  outpx.z += toY0U0V0[8] * t.z;

  outpx.w = alpha;

  write_imagef (out, (int2)(x, y), outpx);
}

kernel void
denoiseprofile_dist(read_only image2d_t in, global float* U4, const int width, const int height,
             const int2 q, const int border)
{
  // U4 covers the image plus a border on each side, read from the image by mirroring
  const int pwidth = width + 2 * border;
  const int pheight = height + 2 * border;
  const int x = get_global_id(0);
  const int y = get_global_id(1);

  if(x >= pwidth || y >= pheight) return;

  const int xi = x - border;
  const int yi = y - border;
  float4 p1 = readpixel(in, fold_mirror(xi, width - 1), fold_mirror(yi, height - 1));
  float4 p2 = readpixel(in, fold_mirror(xi + q.x, width - 1), fold_mirror(yi + q.y, height - 1));
  float4 tmp = (p1 - p2)*(p1 - p2);
  float dist = tmp.x + tmp.y + tmp.z;

  U4[mad24(y, pwidth, x)] = dist;
}

kernel void
denoiseprofile_horiz(global float* U4_in, global float* U4_out, const int width, const int height,
              const int2 q, const int P, local float *buffer)
{
  const int lid = get_local_id(0);
  const int lsz = get_local_size(0);
  const int x = get_global_id(0);
  const int y = get_global_id(1);
  const int gidx = mad24(min(y, height-1), width, min(x, width-1));


  if(y < height)
  {
    /* fill center part of buffer */
    buffer[P + lid] = U4_in[gidx];

    /* left wing of buffer */
    for(int n=0; n <= P/lsz; n++)
    {
      const int l = mad24(n, lsz, lid + 1);
      if(l > P) continue;
      int xx = mad24((int)get_group_id(0), lsz, -l);
      xx = max(xx, 0);
      buffer[P - l] = U4_in[mad24(y, width, xx)];
    }

    /* right wing of buffer */
    for(int n=0; n <= P/lsz; n++)
    {
      const int r = mad24(n, lsz, lsz - lid);
      if(r > P) continue;
      int xx = mad24((int)get_group_id(0), lsz, lsz - 1 + r);
      xx = min(xx, width-1);
      buffer[P + lsz - 1 + r] = U4_in[mad24(y, width, xx)];
    }
  }

  barrier(CLK_LOCAL_MEM_FENCE);

  if(x >= width || y >= height) return;

  buffer += lid + P;

  float distacc = 0.0f;
  for(int pi = -P; pi <= P; pi++)
  {
    distacc += buffer[pi];
  }

  U4_out[gidx] = distacc;
}


kernel void
denoiseprofile_vert(global float* U4_in, global float* U4_out, const int width, const int height,
              const int2 q, const int P, const float norm, local float *buffer,
              const float central_pixel_weight, global float* U4_single_pixel)
{
  const int lid = get_local_id(1);
  const int lsz = get_local_size(1);
  const int x = get_global_id(0);
  const int y = get_global_id(1);
  const int gidx = mad24(min(y, height-1), width, min(x, width-1));

  if(x < width)
  {
    /* fill center part of buffer */
    buffer[P + lid] = U4_in[gidx];

    /* left wing of buffer */
    for(int n=0; n <= P/lsz; n++)
    {
      const int l = mad24(n, lsz, lid + 1);
      if(l > P) continue;
      int yy = mad24((int)get_group_id(1), lsz, -l);
      yy = max(yy, 0);
      buffer[P - l] = U4_in[mad24(yy, width, x)];
    }

    /* right wing of buffer */
    for(int n=0; n <= P/lsz; n++)
    {
      const int r = mad24(n, lsz, lsz - lid);
      if(r > P) continue;
      int yy = mad24((int)get_group_id(1), lsz, lsz - 1 + r);
      yy = min(yy, height-1);
      buffer[P + lsz - 1 + r] = U4_in[mad24(yy, width, x)];
    }
  }

  barrier(CLK_LOCAL_MEM_FENCE);

  if(x >= width || y >= height) return;

  buffer += lid + P;

  float distacc = 0.0f;
  for(int pj = -P; pj <= P; pj++)
  {
    distacc += buffer[pj];
  }

  distacc += U4_single_pixel[gidx] * (2 * P + 1) * (2 * P + 1) * central_pixel_weight;
  distacc /= (1.0f + central_pixel_weight);

  distacc = fast_mexp2f(fmax(0.0f, distacc*norm - 2.0f));

  U4_out[gidx] = distacc;
}


kernel void
denoiseprofile_accu(read_only image2d_t in, global float4* U2, global float* U4,
             const int width, const int height, const int2 q, const int border)
{
  const int x = get_global_id(0);
  const int y = get_global_id(1);
  const int gidx = mad24(y, width, x);

  if(x >= width || y >= height) return;

  // U4 holds the patch distances for offset q over the mirror-extended image; the
  // distance for offset -q at (x, y) is the one for offset q at (x, y) - q
  const int pwidth = width + 2 * border;
  float4 u1_pq = readpixel(in, fold_mirror(x + q.x, width - 1), fold_mirror(y + q.y, height - 1));
  float4 u1_mq = readpixel(in, fold_mirror(x - q.x, width - 1), fold_mirror(y - q.y, height - 1));

  float  u4    = U4[mad24(y + border, pwidth, x + border)];
  float  u4_mq = U4[mad24(y + border - q.y, pwidth, x + border - q.x)];

  float u4_mq_dd = u4_mq * ddirac(q);

  float4 accu = (u4 * u1_pq) + (u4_mq_dd * u1_mq);
  accu.w = u4 + u4_mq_dd;

  U2[gidx] += accu;
}


kernel void
denoiseprofile_finish(read_only image2d_t in, global float4* U2, write_only image2d_t out, const int width, const int height,
                             const float4 a, const float4 sigma2)
{
  const int x = get_global_id(0);
  const int y = get_global_id(1);
  const int gidx = mad24(y, width, x);

  if(x >= width || y >= height) return;

  float4 u2   = U2[gidx];
  const float alpha = readalpha(in, x, y);

  float4 px = ((float4)u2.w > (float4)0.0f ? u2/u2.w : (float4)0.0f);

  px = (px < (float4)0.5f ? (float4)0.0f :
    0.25f*px*px + 0.25f*sqrt(1.5f)/px - 1.375f/(px*px) + 0.625f*dtcl_sqrt(1.5f)/(px*px*px) - 0.125f - sigma2);

  px *= a;
  px.w = alpha;

  write_imagef (out, (int2)(x, y), px);
}


kernel void
denoiseprofile_finish_v2(read_only image2d_t in, global float4* U2, write_only image2d_t out, const int width, const int height,
                             const float4 a, const float4 p, const float4 b, const float bias, const float4 wb)
{
  const int x = get_global_id(0);
  const int y = get_global_id(1);
  const int gidx = mad24(y, width, x);

  if(x >= width || y >= height) return;

  float4 u2   = U2[gidx];
  const float alpha = readalpha(in, x, y);

  float4 px = ((float4)u2.w > (float4)0.0f ? u2/u2.w : (float4)0.0f);

  float4 delta = px * px + (float4)bias;
  float4 denominator = 4.0f / (sqrt(a) * (2.0f - p));
  float4 z1 = (px + dtcl_sqrt(fmax((float4)0.0f, delta))) / denominator;
  px = fmax(dtcl_pow(z1, 1.0f / (1.0f - p / 2.0f)) - b, 0.f);
  px = px * wb;
  px.w = alpha;

  write_imagef (out, (int2)(x, y), px);
}



kernel void
denoiseprofile_backtransform(read_only image2d_t in, write_only image2d_t out, const int width, const int height,
                             const float4 a, const float4 sigma2)
{
  const int x = get_global_id(0);
  const int y = get_global_id(1);
  const int gidx = mad24(y, width, x);

  if(x >= width || y >= height) return;

  float4 px = readpixel(in, x, y);
  const float alpha = px.w;

  px = (px < (float4)0.5f ? (float4)0.0f :
    0.25f*px*px + 0.25f*sqrt(1.5f)/px - 1.375f/(px*px) + 0.625f*sqrt(1.5f)/(px*px*px) - 0.125f - sigma2);

  px *= a;
  px.w = alpha;

  write_imagef (out, (int2)(x, y), px);
}


kernel void
denoiseprofile_backtransform_v2(read_only image2d_t in, write_only image2d_t out, const int width, const int height,
                             const float4 a, const float4 p, const float4 b, const float bias, const float4 wb)
{
  const int x = get_global_id(0);
  const int y = get_global_id(1);
  const int gidx = mad24(y, width, x);

  if(x >= width || y >= height) return;

  float4 px = read_imagef(in, sampleri, (int2)(x, y));
  const float alpha = px.w;

  px = fmax((float4)0.0f, px);
  const float4 delta = px * px + (float4)bias;
  const float4 denominator = 4.0f / (sqrt(a) * (2.0f - p));
  const float4 z1 = (px + sqrt(fmax((float4)0.0f, delta))) / denominator;
  px = fmax(dtcl_pow(z1, 1.0f / (1.0f - p / 2.0f)) - b, 0.f);
  px = px * wb;
  px.w = alpha;

  write_imagef (out, (int2)(x, y), px);
}

kernel void
denoiseprofile_backtransform_Y0U0V0(read_only image2d_t in, write_only image2d_t out, const int width, const int height,
                             const float4 a, const float4 p, const float4 b, const float bias, const float4 wb, global float *toRGB)
{
  const int x = get_global_id(0);
  const int y = get_global_id(1);
  const int gidx = mad24(y, width, x);

  if(x >= width || y >= height) return;

  const float4 t = readpixel(in, x, y);
  const float alpha = t.w;

  float4 px = (float4)0.0f;
  px.x += toRGB[0] * t.x;
  px.x += toRGB[1] * t.y;
  px.x += toRGB[2] * t.z;
  px.y += toRGB[3] * t.x;
  px.y += toRGB[4] * t.y;
  px.y += toRGB[5] * t.z;
  px.z += toRGB[6] * t.x;
  px.z += toRGB[7] * t.y;
  px.z += toRGB[8] * t.z;

  px = fmax((float4)0.0f, px);
  const float4 delta = px * px + (float4)bias * wb;
  const float4 denominator = 4.0f / (sqrt(a) * (2.0f - p));
  const float4 z1 = (px + sqrt(fmax((float4)0.0f, delta))) / denominator;
  px = fmax(dtcl_pow(z1, 1.0f / (1.0f - p / 2.0f)) - b, 0.f);
  px.w = alpha;

  write_imagef (out, (int2)(x, y), px);
}


float4
weight(const float4 c1, const float4 c2, const float inv_sigma2)
{
  const float4 sqr = (c1 - c2)*(c1 - c2);
  const float dt = (sqr.x + sqr.y + sqr.z)*inv_sigma2;
  const float var = 0.02f;
  const float off2 = 9.0f;
  const float r = fast_mexp2f(fmax(0.0f, dt*var - off2));

  return (float4)r;
}


kernel void
denoiseprofile_decompose(read_only image2d_t in, write_only image2d_t coarse, write_only image2d_t detail,
     const int width, const int height, const unsigned int scale, const float inv_sigma2, global const float *filter)
{
  const int x = get_global_id(0);
  const int y = get_global_id(1);

  if(x >= width || y >= height) return;

  const int mult = 1<<scale;

  float4 pixel = readpixel(in, x, y);
  float4 sum = (float4)(0.0f);
  float4 wgt = (float4)(0.0f);

  for(int j=0;j<5;j++) for(int i=0;i<5;i++)
  {
    int xx = mad24(mult, i - 2, x);
    int yy = mad24(mult, j - 2, y);
    int k  = mad24(j, 5, i);

    float4 px = readpixel(in, xx, yy);
    float4 w = filter[k]*weight(pixel, px, inv_sigma2);

    sum += w*px;
    wgt += w;
  }

  sum /= wgt;
  sum.w = pixel.w;

  write_imagef (detail, (int2)(x, y), pixel - sum);
  write_imagef (coarse, (int2)(x, y), sum);
}


kernel void
denoiseprofile_synthesize(read_only image2d_t coarse, read_only image2d_t detail, write_only image2d_t out,
     const int width, const int height,
     const float4 threshold, const float4 boost)
{
  const int x = get_global_id(0);
  const int y = get_global_id(1);

  if(x >= width || y >= height) return;

  float4 c = readpixel(coarse, x, y);
  float4 d = readpixel(detail, x, y);
  float4 amount = copysign(fmax((float4)(0.0f), fabs(d) - threshold), d);
  float4 sum = c + boost*amount;
  sum.w = c.w;
  write_imagef(out, (int2)(x, y), sum);
}


kernel void
denoiseprofile_reduce_first(read_only image2d_t in, const int width, const int height,
                            global float4 *accu, local float4 *buffer)
{
  const int x = get_global_id(0);
  const int y = get_global_id(1);
  const int xlsz = get_local_size(0);
  const int ylsz = get_local_size(1);
  const int xlid = get_local_id(0);
  const int ylid = get_local_id(1);

  const int l = mad24(ylid, xlsz, xlid);

  const int isinimage = (x < width && y < height);
  float4 pixel = readpixel(in, x, y);

  buffer[l] = isinimage ? pixel*pixel : (float4)0.0f;

  barrier(CLK_LOCAL_MEM_FENCE);

  const int lsz = mul24(xlsz, ylsz);

  for(int offset = lsz / 2; offset > 0; offset = offset / 2)
  {
    if(l < offset)
    {
      buffer[l] += buffer[l + offset];
    }
    barrier(CLK_LOCAL_MEM_FENCE);
  }

  const int xgid = get_group_id(0);
  const int ygid = get_group_id(1);
  const int xgsz = get_num_groups(0);

  const int m = mad24(ygid, xgsz, xgid);
  accu[m]   = buffer[0];
}


kernel void
denoiseprofile_reduce_second(const global float4* input, global float4 *result, const int length, local float4 *buffer)
{
  int x = get_global_id(0);
  float4 sum_y2 = (float4)0.0f;

  while(x < length)
  {
    sum_y2 += input[x];

    x += get_global_size(0);
  }

  int lid = get_local_id(0);
  buffer[lid] = sum_y2;

  barrier(CLK_LOCAL_MEM_FENCE);

  for(int offset = get_local_size(0) / 2; offset > 0; offset = offset / 2)
  {
    if(lid < offset)
    {
      buffer[lid] += buffer[lid + offset];
    }
    barrier(CLK_LOCAL_MEM_FENCE);
  }

  if(lid == 0)
  {
    const int gid = get_group_id(0);

    result[gid]   = buffer[0];
  }
}