/*
    This file is part of darktable,
    Copyright (C) 2024 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/>.
*/
typedef enum dt_iop_colorequal_channel_t
{
  HUE = 0,
  SATURATION = 1,
  BRIGHTNESS = 2,
  NUM_CHANNELS = 3,
  GRAD_SWITCH = 4,
  SATURATION_GRAD = SATURATION + GRAD_SWITCH,
  BRIGHTNESS_GRAD = BRIGHTNESS + GRAD_SWITCH
} dt_iop_colorequal_channel_t;

#define SAT_EFFECT 2.0f
#define BRIGHT_EFFECT 8.0f

#include "common.h"
#include "colorspace.h"
#include "color_conversion.h"

#define SATSIZE 4096.0f

static inline float _get_satweight(const float sat, global float *weights)
{
  const float isat = (float)SATSIZE * (1.0f + clamp(sat, -1.0f, 1.0f - (1.0f / SATSIZE)));
  const float base = floor(isat);
  const int i = (int)base;
  return weights[i] + (isat - base) * (weights[i+1] - weights[i]);
}

static inline float gamut_map_HSB(const float4 HSB, global float *gamut_LUT, const float L_white)
{
  const float4 JCH = dt_UCS_HSB_to_JCH(HSB);
  const float max_colorfulness = lookup_gamut(gamut_LUT, JCH.z);
  const float max_chroma = 15.932993652962535f * dtcl_pow(JCH.x * L_white, 0.6523997524738018f) * dtcl_pow(max_colorfulness, 0.6007557017508491f) / L_white;
  const float4 JCH_gamut_boundary = { JCH.x, max_chroma, JCH.z, 0.0f };
  const float4 HSB_gamut_boundary = dt_UCS_JCH_to_HSB(JCH_gamut_boundary);

  // Soft-clip the current pixel saturation at constant brightness
  return soft_clip(HSB.y, 0.8f * HSB_gamut_boundary.y, HSB_gamut_boundary.y);
}

__kernel void init_covariance(global float4 *covariance,
                              global float2 *uv,
                              const int width,
                              const int height)
{
  const int col = get_global_id(0);
  const int row = get_global_id(1);
  if(col >= width || row >= height) return;

  const int k = mad24(row, width, col);

  covariance[k].x = uv[k].x * uv[k].x;
  covariance[k].y = uv[k].x * uv[k].y;
  covariance[k].z = uv[k].x * uv[k].y;
  covariance[k].w = uv[k].y * uv[k].y;
}

__kernel void finish_covariance(global float4 *covariance,
                                global float2 *uv,
                                const int width,
                                const int height)
{
  const int col = get_global_id(0);
  const int row = get_global_id(1);
  if(col >= width || row >= height) return;

  const int k = mad24(row, width, col);

  covariance[k].x -= uv[k].x * uv[k].x;
  covariance[k].y -= uv[k].x * uv[k].y;
  covariance[k].z -= uv[k].x * uv[k].y;
  covariance[k].w -= uv[k].y * uv[k].y;
}

__kernel void prepare_prefilter(global float2 *uv,
                                global float4 *covariance,
                                global float4 *a,
                                global float2 *b,
                                const float eps,
                                const int width,
                                const int height)
{
  const int col = get_global_id(0);
  const int row = get_global_id(1);
  if(col >= width || row >= height) return;

  const int k = mad24(row, width, col);

  const float4 sigma = { covariance[k].x + eps, covariance[k].y, covariance[k].z, covariance[k].w + eps};
  const float det = sigma.x * sigma.w - sigma.y * sigma.z;

  if(fabs(det) > 4.0f * FLT_EPSILON)
  {
    const float4 sigma_inv = { sigma.w / det, -sigma.y / det, -sigma.z / det, sigma.x / det };
    a[k].x = covariance[k].x * sigma_inv.x + covariance[k].y * sigma_inv.y;
    a[k].y = covariance[k].x * sigma_inv.z + covariance[k].y * sigma_inv.w;
    a[k].z = covariance[k].z * sigma_inv.x + covariance[k].w * sigma_inv.y;
    a[k].w = covariance[k].z * sigma_inv.z + covariance[k].w * sigma_inv.w;
  }
  else
    a[k].x = a[k].y = a[k].z = a[k].w = 0.0f;

  b[k].x = uv[k].x - a[k].x * uv[k].x - a[k].y * uv[k].y;
  b[k].y = uv[k].y - a[k].z * uv[k].x - a[k].w * uv[k].y;
}

__kernel void apply_prefilter(global float2 *uv,
                              global float *saturation,
                              global float4 *a,
                              global float2 *b,
                              global float *weights,
                              const float sat_shift,
                              const int width,
                              const int height)
{
  const int col = get_global_id(0);
  const int row = get_global_id(1);
  if(col >= width || row >= height) return;

  const int k = mad24(row, width, col);

  const float2 UV = uv[k];
  const float2 cv = (float2)(a[k].x * UV.x + a[k].y * UV.y + b[k].x,
                             a[k].z * UV.x + a[k].w * UV.y + b[k].y);

  const float satweight = _get_satweight(saturation[k] - sat_shift, weights);
  uv[k].x = mix(UV.x, cv.x, satweight);
  uv[k].y = mix(UV.y, cv.y, satweight);
}

__kernel void prepare_correlations(global float2 *corrections,
                                   global float *b_corrections,
                                   global float2 *uv,
                                   global float4 *correlations,
                                   const int width,
                                   const int height)
{
  const int col = get_global_id(0);
  const int row = get_global_id(1);
  if(col >= width || row >= height) return;

  const int k = mad24(row, width, col);

  correlations[k].x = uv[k].x * corrections[k].y;
  correlations[k].y = uv[k].y * corrections[k].y;
  correlations[k].z = uv[k].x * b_corrections[k];
  correlations[k].w = uv[k].y * b_corrections[k];
}

// also write covariance
__kernel void finish_correlations(global float2 *corrections,
                                  global float *b_corrections,
                                  global float2 *uv,
                                  global float4 *correlations,
                                  global float4 *covariance,
                                  const int width,
                                  const int height)
{
  const int col = get_global_id(0);
  const int row = get_global_id(1);
  if(col >= width || row >= height) return;

  const int k = mad24(row, width, col);

  covariance[k].x -= uv[k].x * uv[k].x;
  covariance[k].y -= uv[k].x * uv[k].y;
  covariance[k].z -= uv[k].x * uv[k].y;
  covariance[k].w -= uv[k].y * uv[k].y;

  correlations[k].x -= uv[k].x * corrections[k].y;
  correlations[k].y -= uv[k].y * corrections[k].y;
  correlations[k].z -= uv[k].x * b_corrections[k];
  correlations[k].w -= uv[k].y * b_corrections[k];
}

__kernel void final_guide(global float4 *covariance,
                          global float4 *correlations,
                          global float2 *corrections,
                          global float *b_corrections,
                          global float2 *uv,
                          global float4 *a,
                          global float2 *b,
                          const float eps,
                          const int width,
                          const int height)
{
  const int col = get_global_id(0);
  const int row = get_global_id(1);
  if(col >= width || row >= height) return;

  const int k = mad24(row, width, col);

  const float4 sigma = { covariance[k].x + eps, covariance[k].y, covariance[k].z, covariance[k].w + eps };
  const float det = sigma.x * sigma.w - sigma.y * sigma.z;

  if(fabs(det) > 4.0f * FLT_EPSILON)
  {
    const float4 sigma_inv = { sigma.w / det, -sigma.y / det, -sigma.z / det, sigma.x / det };
    a[k].x = correlations[k].x * sigma_inv.x + correlations[k].y * sigma_inv.y;
    a[k].y = correlations[k].x * sigma_inv.z + correlations[k].y * sigma_inv.w;
    a[k].z = correlations[k].z * sigma_inv.x + correlations[k].w * sigma_inv.y;
    a[k].w = correlations[k].z * sigma_inv.z + correlations[k].w * sigma_inv.w;
  }
  else
    a[k].x = a[k].y = a[k].z = a[k].w = 0.0f;

  b[k].x = corrections[k].y - a[k].x * uv[k].x - a[k].y * uv[k].y;
  b[k].y = b_corrections[k] - a[k].z * uv[k].x - a[k].w * uv[k].y;
}

__kernel void apply_guided(global float2 *uv,
                           global float *saturation,
                           global float *scharr,
                           global float4 *a,
                           global float2 *b,
                           global float2 *corrections,
                           global float *b_corrections,
                           global float *weights,
                           const float sat_shift,
                           const float bright_shift,
                           const int width,
                           const int height)
{
  const int col = get_global_id(0);
  const int row = get_global_id(1);
  if(col >= width || row >= height) return;

  const int k = mad24(row, width, col);

  const float2 CV = { a[k].x * uv[k].x + a[k].y * uv[k].y + b[k].x,
                      a[k].z * uv[k].x + a[k].w * uv[k].y + b[k].y };

  corrections[k].y = 1.0f + (CV.x - 1.0f) * _get_satweight(saturation[k] - sat_shift, weights);
  const float gradient_weight = 1.0f - clipf(scharr[k]);
  b_corrections[k] = CV.y * gradient_weight * _get_satweight(saturation[k] - bright_shift, weights);
}

__kernel void sample_input(__read_only image2d_t dev_in,
                           global float *saturation,
                           global float *lum,
                           global float2 *uv,
                           global float4 *pix_out,
                           global float4 *mat,
                           const int width,
                           const int height)
{
  const int col = get_global_id(0);
  const int row = get_global_id(1);
  if(col >= width || row >= height) return;

  const int k = mad24(row, width, col);

  const float4 pix_in = Areadpixel(dev_in, col, row);
  // calc saturation from input data
  const float dmin = fmin(pix_in.x, fmin(pix_in.y, pix_in.z));
  const float dmax = fmax(pix_in.x, fmax(pix_in.y, pix_in.z));
  const float delta = dmax - dmin;
  saturation[k] = (dmax > NORM_MIN && delta > NORM_MIN) ? delta / dmax : 0.0f;

  const float4 M[3] = { mat[0], mat[1], mat[2] };
  const float4 XYZ_D65 = matrix_dot(pix_in, M);
  const float4 xyY = dt_D65_XYZ_to_xyY(XYZ_D65);

  lum[k] = Y_to_dt_UCS_L_star(xyY.z);
  uv[k] = xyY_to_dt_UCS_UV(xyY);
  pix_out[k].w = pix_in.w;
}

__kernel void write_output(__write_only image2d_t dev_out,
                            global float4 *pix_out,
                            global float2 *corrections,
                            global float *b_corrections,
                            global float4 *mat,
                            global float *gamut_LUT,
                            const float white,
                            const int width,
                            const int height)
{
  const int col = get_global_id(0);
  const int row = get_global_id(1);
  if(col >= width || row >= height) return;

  const float4 M[3] = { mat[0], mat[1], mat[2] };

  const int k = mad24(row, width, col);

  pix_out[k].x += corrections[k].x;
  pix_out[k].y = fmax(0.0f, pix_out[k].y * (1.0f + SAT_EFFECT * (corrections[k].y - 1.0f)));
  pix_out[k].z = fmax(0.0f, pix_out[k].z * (1.0f + BRIGHT_EFFECT * b_corrections[k]));

  pix_out[k].y = gamut_map_HSB(pix_out[k], gamut_LUT, white);
  const float4 XYZ_D65 = dt_UCS_HSB_to_XYZ(pix_out[k], white);
  const float4 pout = matrix_dot(XYZ_D65, M);
  write_imagef(dev_out, (int2)(col, row), pout);
}

__kernel void write_visual (__write_only image2d_t dev_out,
                            global float4 *pixout,
                            global float2 *corrections,
                            global float *b_corrections,
                            global float *saturation,
                            global float *scharr,
                            global float *weights,
                            const float bright_shift,
                            const float sat_shift,
                            const float white,
                            const int mode,
                            const int width,
                            const int height)
{
  const int col = get_global_id(0);
  const int row = get_global_id(1);
  if(col >= width || row >= height) return;

  const int k = mad24(row, width, col);
  const float val = dtcl_sqrt(fmax(0.0f,  pixout[k].z * white));
  float corr = 0.0f;
  switch(mode)
  {
    case BRIGHTNESS:
        corr = BRIGHT_EFFECT * b_corrections[k];
        break;
    case SATURATION:
        corr = SAT_EFFECT * (corrections[k].y - 1.0f);
        break;
    case BRIGHTNESS_GRAD:
        corr = _get_satweight(saturation[k] - bright_shift, weights) - 0.5f;
        break;
    case SATURATION_GRAD:
        corr = _get_satweight(saturation[k] - sat_shift, weights) - 0.5f;
        break;
    default:  // HUE
        corr = 0.2f * corrections[k].x;
  }

  const int neg = corr < 0.0f;
  corr = fabs(corr);
  corr = corr < 2e-3f ? 0.0f : corr;

  float4 pout = (float4)(fmax(0.0f, neg ? val - corr : val),
                         fmax(0.0f, val - corr),
                         fmax(0.0f, neg ? val : val - corr),
                         0.0f);


  if(mode == BRIGHTNESS)
  {
    if(scharr[k] > 0.1f)
    {
      pout.x = pout.z = 0.0f;
      pout.y = scharr[k];
    }
  }
  write_imagef(dev_out, (int2)(col, row), pout);
}

__kernel void draw_weight(__write_only image2d_t dev_out,
                          global float *weights,
                          const float bright_shift,
                          const float sat_shift,
                          const int mode,
                          const int width,
                          const int height)
{
  const int col = get_global_id(0);

  if(col >= width) return;

  const float eps = 0.5f / (float)height;
  const float shift = (mode == SATURATION_GRAD) ? sat_shift : bright_shift;
  const float4 mark = (float4)( 0.0f, 1.0f, 0.0f, 0.0f);
  for(int i = 0; i < 16; i++)
  {
    const float pos = (float)(16*col +i) / (float)(16*width);
    const float weight = _get_satweight(pos - shift, weights);
    if(weight > eps && weight < 1.0f - eps)
    {
      const int row = (int)((1.0f - weight) * (float)(height-1));
      write_imagef(dev_out, (int2)(col, row), mark);
    }
  }
}

__kernel void process_data(global float2 *uv,
                           global float *Lscharr,
                           global float *saturation,
                           global float2 *corrections,
                           global float *b_corrections,
                           global float4 *pixout,
                           global float *LUT_saturation,
                           global float *LUT_hue,
                           global float *LUT_brightness,
                           const float white,
                           const float gradient_amp,
                           const int guiding,
                           const int width,
                           const int height)
{
  const int col = get_global_id(0);
  const int row = get_global_id(1);
  if(col >= width || row >= height) return;

  const int k = mad24(row, width, col);

  const float4 JCH = dt_UCS_LUV_to_JCH(Lscharr[k], white, uv[k]);
  const float4 HSB = dt_UCS_JCH_to_HSB(JCH);

  const float hue = HSB.x;
  const float sat = HSB.y;
  pixout[k].x = hue;
  pixout[k].y = sat;
  pixout[k].z = HSB.z;

  if(guiding)
  {
    const int kk = mad24(clamp(row, 1, height - 2), width, clamp(col, 1, width - 2));

    const float kscharr = fmax(0.0f, scharr_gradient(saturation, kk, width) - 0.02f);
    Lscharr[k] = gradient_amp * kscharr * kscharr;
  }

  if(JCH.y > NORM_MIN)
  {
    corrections[k].x = lookup_gamut(LUT_hue, hue);
    corrections[k].y = lookup_gamut(LUT_saturation, hue);
    b_corrections[k] = sat * (lookup_gamut(LUT_brightness, hue) - 1.0f);
  }
  else
  {
    corrections[k].x = 0.0f;
    corrections[k].y = 1.0f;
    b_corrections[k] = 0.0f;
  }
}