// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Utility functions for optimizing multi-dimensional nonlinear functions.

#ifndef LIB_JXL_OPTIMIZE_H_
#define LIB_JXL_OPTIMIZE_H_

#include <cmath>
#include <cstdio>

#include "lib/jxl/base/status.h"

namespace jxl {
namespace optimize {

// An array type of numeric values that supports math operations with operator-,
// operator+, etc.
template <typename T, size_t N>
class Array {
 public:
  Array() = default;
  explicit Array(T v) {
    for (size_t i = 0; i < N; i++) v_[i] = v;
  }

  size_t size() const { return N; }

  T& operator[](size_t index) {
    JXL_DASSERT(index < N);
    return v_[index];
  }
  T operator[](size_t index) const {
    JXL_DASSERT(index < N);
    return v_[index];
  }

 private:
  // The values used by this Array.
  T v_[N];
};

template <typename T, size_t N>
Array<T, N> operator+(const Array<T, N>& x, const Array<T, N>& y) {
  Array<T, N> z;
  for (size_t i = 0; i < N; ++i) {
    z[i] = x[i] + y[i];
  }
  return z;
}

template <typename T, size_t N>
Array<T, N> operator-(const Array<T, N>& x, const Array<T, N>& y) {
  Array<T, N> z;
  for (size_t i = 0; i < N; ++i) {
    z[i] = x[i] - y[i];
  }
  return z;
}

template <typename T, size_t N>
Array<T, N> operator*(T v, const Array<T, N>& x) {
  Array<T, N> y;
  for (size_t i = 0; i < N; ++i) {
    y[i] = v * x[i];
  }
  return y;
}

template <typename T, size_t N>
T operator*(const Array<T, N>& x, const Array<T, N>& y) {
  T r = 0.0;
  for (size_t i = 0; i < N; ++i) {
    r += x[i] * y[i];
  }
  return r;
}

// Implementation of the Scaled Conjugate Gradient method described in the
// following paper:
//   Moller, M. "A Scaled Conjugate Gradient Algorithm for Fast Supervised
//   Learning", Neural Networks, Vol. 6. pp. 525-533, 1993
//   http://sci2s.ugr.es/keel/pdf/algorithm/articulo/moller1990.pdf
//
// The Function template parameter is a class that has the following method:
//
//   // Returns the value of the function at point w and sets *df to be the
//   // negative gradient vector of the function at point w.
//   double Compute(const optimize::Array<T, N>& w,
//                  optimize::Array<T, N>* df) const;
//
// Returns a vector w, such that |df(w)| < grad_norm_threshold.
template <typename T, size_t N, typename Function>
Array<T, N> OptimizeWithScaledConjugateGradientMethod(
    const Function& f, const Array<T, N>& w0, const T grad_norm_threshold,
    size_t max_iters) {
  const size_t n = w0.size();
  const T rsq_threshold = grad_norm_threshold * grad_norm_threshold;
  const T sigma0 = static_cast<T>(0.0001);
  const T l_min = static_cast<T>(1.0e-15);
  const T l_max = static_cast<T>(1.0e15);

  Array<T, N> w = w0;
  Array<T, N> wp;
  Array<T, N> r;
  Array<T, N> rt;
  Array<T, N> e;
  Array<T, N> p;
  T psq;
  T fp;
  T D;
  T d;
  T m;
  T a;
  T b;
  T s;
  T t;

  T fw = f.Compute(w, &r);
  T rsq = r * r;
  e = r;
  p = r;
  T l = static_cast<T>(1.0);
  bool success = true;
  size_t n_success = 0;
  size_t k = 0;

  while (k++ < max_iters) {
    if (success) {
      m = -(p * r);
      if (m >= 0) {
        p = r;
        m = -(p * r);
      }
      psq = p * p;
      s = sigma0 / std::sqrt(psq);
      f.Compute(w + (s * p), &rt);
      t = (p * (r - rt)) / s;
    }

    d = t + l * psq;
    if (d <= 0) {
      d = l * psq;
      l = l - t / psq;
    }

    a = -m / d;
    wp = w + a * p;
    fp = f.Compute(wp, &rt);

    D = 2.0 * (fp - fw) / (a * m);
    if (D >= 0.0) {
      success = true;
      n_success++;
      w = wp;
    } else {
      success = false;
    }

    if (success) {
      e = r;
      r = rt;
      rsq = r * r;
      fw = fp;
      if (rsq <= rsq_threshold) {
        break;
      }
    }

    if (D < 0.25) {
      l = std::min(4.0 * l, l_max);
    } else if (D > 0.75) {
      l = std::max(0.25 * l, l_min);
    }

    if ((n_success % n) == 0) {
      p = r;
      l = 1.0;
    } else if (success) {
      b = ((e - r) * r) / m;
      p = b * p + r;
    }
  }

  return w;
}

}  // namespace optimize
}  // namespace jxl

#endif  // LIB_JXL_OPTIMIZE_H_