/*
    This file is part of darktable,
    Copyright (C) 2019-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/>.
*/

#include <math.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <time.h>

#include "common/darktable.h"
#include "common/imagebuf.h"
#include "develop/imageop_math.h"


/* DOCUMENTATION
 *
 * Choleski decomposition is a fast way to solve linear systems of equations
 * described by a positive definite hermitian (= square and symmetrical) matrix.
 * It is a special case of LU decompositions enabling special optimizations.
 * For matrices not matching this requirement,
 * you need to use the Gauss-Jordan elimination in iop/gaussian_elemination.h,
 * which is about twice as slow but more general.
 *
 * To solve A x = y, for x, with A a positive definite hermitian real matrix :
 *
 *  1) find L such that A = L × L' (Choleski decomposition)
 *  2) Second step : solve L × b = y for b (triangular descent)
 *  3) Third step : solve L' × x = b for x (triangular ascent)
 *
 * L is a lower diagonal matrix such that :
 *
 *     [ l11  0    0   ]
 * L = [ l12  l22  0   ] (if n = 3)
 *     [ l13  l23  l33 ]
 *
 * and L' is its transpose :
 *
 *      [ l11 l12 l13 ]
 * L' = [ 0   l22 l23 ] (if n = 3)
 *      [ 0   0   l33 ]
 *
 * We use the Cholesky-Banachiewicz algorithm for the decomposition
 * because it operates row by row, which is more suitable with C
 * memory layout.
 *
 * The return codes are in sync with /iop/gaussian_elimination.h, in
 * order, maybe one day, to have a general linear solving lib that
 * could for example iterate over methods until one succeed.
 *
 * We didn't bother to parallelize the code because it's already
 * fast enough. Double precision is required to limit rounding error
 * accumulation in the summing loops: forming A'A squares the condition
 * number, and the last Cholesky pivot is then a difference of nearly
 * equal quantities that float cannot resolve.
 *
 * Vectorization leads to slow-downs here since we access matrices
 * row-wise and column-wise, in a non-contiguous fashion.
 *
 * References :
 *  "Analyse numérique pour ingénieurs", 4e edition, André Fortin,
 *  Presses Internationales de Polytechnique Montréal, 2011.
 *
 *  Cholesky method,
 *  https://algowiki-project.org/en/Cholesky_method#The_.5Bmath.5DLL.5ET.5B.2Fmath.5D_decomposition
 *  https://en.wikipedia.org/wiki/Cholesky_decomposition
 *  https://rosettacode.org/wiki/Cholesky_decomposition#C
 *
 */


__DT_CLONE_TARGETS__
static inline gboolean _choleski_decompose(const double *const restrict A,
                                           double *const restrict L,
                                           size_t n,
                                           const gboolean verbose)
{
  // A is input n×n matrix, decompose it into L such that A = L × L'
  // slow and safe variant : we check values for negatives in sqrt and
  // divisions by 0.

  if(A[0] <= 0.0) return FALSE; // failure : non positive definite matrix

  gboolean valid = TRUE;

  for(size_t i = 0; i < n; i++)
    for(size_t j = 0; j < (i + 1); j++)
    {
      double sum = 0.0;

      for(size_t k = 0; k < j; k++)
        sum += L[i * n + k] * L[j * n + k];

      if(i == j)
      {
        const double temp = A[i * n + i] - sum;

        if(temp < 0.0)
        {
          valid = FALSE;
          L[i * n + j] = NAN;
        }
        else
          L[i * n + j] = sqrt(temp);
      }
      else
      {
        const double temp = L[j * n + j];

        if(temp == 0.0)
        {
          valid = FALSE;
          L[i * n + j] = NAN;
        }
        else
          L[i * n + j] = (A[i * n + j] - sum) / temp;
      }
    }

  if(!valid && verbose)
    dt_print(DT_DEBUG_ALWAYS, "Cholesky decomposition returned NaNs");
  return valid;
}


__DT_CLONE_TARGETS__
static inline gboolean _triangular_descent(const double *const restrict L,
                                           const double *const restrict y,
                                           double *const restrict b,
                                           const size_t n,
                                           const gboolean verbose)
{
  // solve L × b = y for b
  // use the lower triangular part of L from top to bottom

  gboolean valid = TRUE;

  for(size_t i = 0; i < n; ++i)
  {
    double sum = y[i];
    for(size_t j = 0; j < i; ++j)
      sum -= L[i * n + j] * b[j];

    const double temp = L[i * n + i];

    if(temp != 0.0)
      b[i] = sum / temp;
    else
    {
      b[i] = NAN;
      valid = FALSE;
    }
  }
  if(!valid && verbose)
    dt_print(DT_DEBUG_ALWAYS, "Cholesky LU triangular descent returned NaNs");
  return valid;
}

__DT_CLONE_TARGETS__
static inline gboolean _triangular_ascent(const double *const restrict L,
                                          const double *const restrict b,
                                          double *const restrict x,
                                          const size_t n,
                                          const gboolean verbose)
{
  // solve L' × x = b for x
  // use the lower triangular part of L transposed from bottom to top

  gboolean valid = TRUE;

  for(int i = (n - 1); i > -1 ; --i)
  {
    double sum = b[i];
    for(int j = (n - 1); j > i; --j)
      sum -= L[j * n + i] * x[j];

    const double temp = L[i * n + i];
    if(temp != 0.0)
      x[i] = sum / temp;
    else
    {
      x[i] = NAN;
      valid = FALSE;
    }
  }

  if(!valid && verbose)
    dt_print(DT_DEBUG_ALWAYS, "Cholesky LU triangular ascent returned NaNs");
  return valid;
}


__DT_CLONE_TARGETS__
static inline gboolean _solve_hermitian(const double *const restrict A,
                                        double *const restrict y,
                                        const size_t n,
                                        const gboolean verbose)
{
  // Solve A x = y where A an hermitian positive definite matrix n × n
  // x and y are n vectors. Output the result in y

  double *const restrict x = dt_alloc_align_double(n);
  double *const restrict L = dt_alloc_align_double(n * n);

  if(!x || !L)
  {
    dt_free_align(x);
    dt_free_align(L);
    dt_control_log(_("Choleski decomposition failed to allocate memory,"
                     " check your RAM settings"));
    dt_print(DT_DEBUG_PIPE, "Choleski decomposition failed to allocate memory");
    return FALSE;
  }

  gboolean valid = FALSE;

  // LU decomposition
  valid = _choleski_decompose(A, L, n, verbose);

  // Triangular descent
  if(valid)
    valid = _triangular_descent(L, y, x, n, verbose);

  // Triangular ascent
  if(valid)
    valid = _triangular_ascent(L, x, y, n, verbose);

  dt_free_align(x);
  dt_free_align(L);

  return valid;
}


__DT_CLONE_TARGETS__
static inline void _transpose_dot_matrix(float *const restrict A, // input
                                         double *const restrict A_square, // output
                                         const size_t m,
                                         const size_t n)
{
  // Construct the square symmetrical definite positive matrix A' A,
  // BUT only compute the lower triangle part for performance

  for(size_t i = 0; i < n; ++i)
    for(size_t j = 0; j < (i + 1); ++j)
    {
      double sum = 0.0;
      for(size_t k = 0; k < m; ++k)
        sum += (double)A[k * n + i] * (double)A[k * n + j];

      A_square[i * n + j] = sum;
    }
}


__DT_CLONE_TARGETS__
static inline void _transpose_dot_vector(float *const restrict A, // input
                                         float *const restrict y, // input
                                         double *const restrict y_square, // output
                                         const size_t m,
                                         const size_t n)
{
  // Construct the vector A' y
  for(size_t i = 0; i < n; ++i)
  {
    double sum = 0.0;
    for(size_t k = 0; k < m; ++k)
      sum += (double)A[k * n + i] * (double)y[k];

    y_square[i] = sum;
  }
}


__DT_CLONE_TARGETS__
static inline gboolean pseudo_solve(float *const restrict A,
                                    float *const restrict y,
                                    const size_t m,
                                    const size_t n,
                                    const gboolean verbose)
{
  // Solve the linear problem A x = y with the over-constrained
  // rectanguler matrice A of dimension m × n (m >= n) by the least
  // squares method

  if(m < n || n < 2 || m < 2)
  {
    dt_print(DT_DEBUG_ALWAYS, "Pseudo solve: cannot cast %zu × %zu matrice", m, n);
    return FALSE;
  }

  double *const restrict A_square = dt_alloc_align_double(n * n);
  double *const restrict y_square = dt_alloc_align_double(n);

  if(!A_square || !y_square)
  {
    dt_free_align(A_square);
    dt_free_align(y_square);
    dt_print(DT_DEBUG_PIPE, "Choleski decomposition failed to allocate memory");
    dt_control_log(_("Choleski decomposition failed to allocate memory,"
                     " check your RAM settings"));
    return FALSE;
  }

  DT_OMP_PRAGMA(parallel sections num_threads(darktable.num_openmp_threads))
  {
    DT_OMP_PRAGMA(section)
    {
      // Prepare the least squares matrix = A' A
      _transpose_dot_matrix(A, A_square, m, n);
    }
    DT_OMP_PRAGMA(section)
    {
      // Prepare the y square vector = A' y
      _transpose_dot_vector(A, y, y_square, m, n);
    }
  }

  // Solve A' A x = A' y for x
  gboolean valid = _solve_hermitian(A_square, y_square, n, verbose);
  if(valid)
  {
    for(size_t i = 0; i < n; i++)
      y[i] = (float) y_square[i];
  }

  dt_free_align(y_square);
  dt_free_align(A_square);

  return valid;
}

// clang-format off
// modelines: These editor modelines have been set for all relevant files by tools/update_modelines.py
// vim: shiftwidth=2 expandtab tabstop=2 cindent
// kate: tab-indents: off; indent-width 2; replace-tabs on; indent-mode cstyle; remove-trailing-spaces modified;
// clang-format on