// Copyright (c) 2020 Chris Richardson
// FEniCS Project
// SPDX-License-Identifier:    MIT

#include "polyset.h"
#include "cell.h"
#include "indexing.h"
#include <Eigen/Dense>
#include <array>
#include <cmath>

using namespace basix;

namespace
{
// Compute coefficients in the Jacobi Polynomial recurrence relation
constexpr std::array<double, 3> jrc(int a, int n)
{
  double an = (a + 2 * n + 1) * (a + 2 * n + 2)
              / static_cast<double>(2 * (n + 1) * (a + n + 1));
  double bn = a * a * (a + 2 * n + 1)
              / static_cast<double>(2 * (n + 1) * (a + n + 1) * (a + 2 * n));
  double cn = n * (a + n) * (a + 2 * n + 2)
              / static_cast<double>((n + 1) * (a + n + 1) * (a + 2 * n));
  return {an, bn, cn};
}
//-----------------------------------------------------------------------------
// Compute the complete set of derivatives from 0 to nderiv, for all the
// polynomials up to order n on a line segment. The polynomials used are
// Legendre Polynomials, with the recurrence relation given by
// n P(n) = (2n - 1) x P_{n-1} - (n - 1) P_{n-2} in the interval [-1, 1]. The
// range is rescaled here to [0, 1].
std::vector<Eigen::ArrayXXd>
tabulate_polyset_line_derivs(int degree, int nderiv, const Eigen::ArrayXXd& x)
{
  assert(x.cols() == 1);
  const Eigen::ArrayXXd X = x * 2.0 - 1.0;

  const int m = (degree + 1);

  std::vector<Eigen::ArrayXXd> dresult(nderiv + 1);
  for (int k = 0; k < nderiv + 1; ++k)
  {
    // Get reference to this derivative
    Eigen::ArrayXXd result(x.rows(), m);

    if (k == 0)
      result.col(0).fill(1.0);
    else
      result.col(0).setZero();

    for (int p = 1; p < degree + 1; ++p)
    {
      const double a = 1.0 - 1.0 / static_cast<double>(p);
      result.col(p) = X * result.col(p - 1) * (a + 1.0);
      if (k > 0)
        result.col(p) += 2 * k * dresult[k - 1].col(p - 1) * (a + 1.0);
      if (p > 1)
        result.col(p) -= result.col(p - 2) * a;
    }

    dresult[k] = result;
  }

  // Normalise
  for (int k = 0; k < nderiv + 1; ++k)
  {
    for (int p = 0; p < degree + 1; ++p)
      dresult[k].col(p) *= std::sqrt(p + 0.5);
  }

  return dresult;
}
//-----------------------------------------------------------------------------
// Compute the complete set of derivatives from 0 to nderiv, for all the
// polynomials up to order n on a triangle in [0, 1][0, 1].
// The polynomials P_{pq} are built up in sequence, firstly along q = 0, which
// is a line segment, as in tabulate_polyset_interval_derivs above, but with a
// change of variables. The polynomials are then extended in the q direction,
// using the relation given in Sherwin and Karniadakis 1995
// (https://doi.org/10.1016/0045-7825(94)00745-9)
std::vector<Eigen::ArrayXXd>
tabulate_polyset_triangle_derivs(int n, int nderiv, const Eigen::ArrayXXd& pts)

{
  assert(pts.cols() == 2);

  Eigen::ArrayXXd x = pts * 2.0 - 1.0;

  const int m = (n + 1) * (n + 2) / 2;
  const int md = (nderiv + 1) * (nderiv + 2) / 2;
  std::vector<Eigen::ArrayXXd> dresult(md);

  // f3 = ((1-y)/2)^2
  const Eigen::ArrayXd f3 = (1.0 - x.col(1)).square() * 0.25;

  // Iterate over derivatives in increasing order, since higher derivatives
  // depend on earlier calculations
  Eigen::ArrayXXd result(pts.rows(), m);
  for (int k = 0; k < nderiv + 1; ++k)
  {
    for (int kx = 0; kx < k + 1; ++kx)
    {
      const int ky = k - kx;

      if (kx == 0 and ky == 0)
        result.col(0).fill(1.0);
      else
        result.col(0).setZero();

      for (int p = 1; p < n + 1; ++p)
      {
        const double a
            = static_cast<double>(2 * p - 1) / static_cast<double>(p);
        result.col(idx(p, 0))
            = (x.col(0) + 0.5 * x.col(1) + 0.5) * result.col(idx(p - 1, 0)) * a;
        if (kx > 0)
        {
          result.col(idx(p, 0))
              += 2 * kx * a * dresult[idx(kx - 1, ky)].col(idx(p - 1, 0));
        }

        if (ky > 0)
        {
          result.col(idx(p, 0))
              += ky * a * dresult[idx(kx, ky - 1)].col(idx(p - 1, 0));
        }

        if (p > 1)
        {
          // y^2 terms
          result.col(idx(p, 0)) -= f3 * result.col(idx(p - 2, 0)) * (a - 1.0);

          if (ky > 0)
          {
            result.col(idx(p, 0))
                -= ky * (x.col(1) - 1.0)
                   * dresult[idx(kx, ky - 1)].col(idx(p - 2, 0)) * (a - 1.0);
          }

          if (ky > 1)
          {
            result.col(idx(p, 0))
                -= ky * (ky - 1) * dresult[idx(kx, ky - 2)].col(idx(p - 2, 0))
                   * (a - 1.0);
          }
        }
      }

      for (int p = 0; p < n; ++p)
      {
        result.col(idx(p, 1))
            = result.col(idx(p, 0)) * (x.col(1) * (1.5 + p) + 0.5 + p);
        if (ky > 0)
        {
          result.col(idx(p, 1))
              += 2 * ky * (1.5 + p) * dresult[idx(kx, ky - 1)].col(idx(p, 0));
        }

        for (int q = 1; q < n - p; ++q)
        {
          const auto [a1, a2, a3] = jrc(2 * p + 1, q);
          result.col(idx(p, q + 1))
              = result.col(idx(p, q)) * (x.col(1) * a1 + a2)
                - result.col(idx(p, q - 1)) * a3;
          if (ky > 0)
          {
            result.col(idx(p, q + 1))
                += 2 * ky * a1 * dresult[idx(kx, ky - 1)].col(idx(p, q));
          }
        }
      }

      // Store this derivative
      dresult[idx(kx, ky)] = result;
    }
  }

  // Normalisation
  for (std::size_t j = 0; j < dresult.size(); ++j)
  {
    for (int p = 0; p < n + 1; ++p)
      for (int q = 0; q < n - p + 1; ++q)
        dresult[j].col(idx(p, q)) *= std::sqrt((p + 0.5) * (p + q + 1));
  }

  return dresult;
}
//-----------------------------------------------------------------------------
std::vector<Eigen::ArrayXXd>
tabulate_polyset_tetrahedron_derivs(int n, int nderiv,
                                    const Eigen::ArrayXXd& pts)
{
  assert(pts.cols() == 3);

  Eigen::ArrayXXd x = pts * 2.0 - 1.0;

  const int m = (n + 1) * (n + 2) * (n + 3) / 6;
  const int md = (nderiv + 1) * (nderiv + 2) * (nderiv + 3) / 6;
  std::vector<Eigen::ArrayXXd> dresult(md);

  const Eigen::ArrayXd f2 = (x.col(1) + x.col(2)).square() * 0.25;
  const Eigen::ArrayXd f3 = (1.0 + x.col(1) * 2.0 + x.col(2)) * 0.5;
  const Eigen::ArrayXd f4 = (1.0 - x.col(2)) * 0.5;
  const Eigen::ArrayXd f5 = f4 * f4;

  // Traverse derivatives in increasing order
  Eigen::ArrayXXd result(pts.rows(), m);
  for (int k = 0; k < nderiv + 1; ++k)
  {
    for (int j = 0; j < k + 1; ++j)
    {
      for (int kx = 0; kx < j + 1; ++kx)
      {
        const int ky = j - kx;
        const int kz = k - j;
        if (kx == 0 and ky == 0 and kz == 0)
          result.col(0).fill(1.0);
        else
          result.col(0).setZero();

        for (int p = 1; p < n + 1; ++p)
        {
          double a = static_cast<double>(2 * p - 1) / static_cast<double>(p);
          result.col(idx(p, 0, 0))
              = (x.col(0) + 0.5 * (x.col(1) + x.col(2)) + 1.0)
                * result.col(idx(p - 1, 0, 0)) * a;
          if (kx > 0)
          {
            result.col(idx(p, 0, 0))
                += 2 * kx * a
                   * dresult[idx(kx - 1, ky, kz)].col(idx(p - 1, 0, 0));
          }

          if (ky > 0)
          {
            result.col(idx(p, 0, 0))
                += ky * a * dresult[idx(kx, ky - 1, kz)].col(idx(p - 1, 0, 0));
          }

          if (kz > 0)
          {
            result.col(idx(p, 0, 0))
                += kz * a * dresult[idx(kx, ky, kz - 1)].col(idx(p - 1, 0, 0));
          }

          if (p > 1)
          {
            result.col(idx(p, 0, 0))
                -= f2 * result.col(idx(p - 2, 0, 0)) * (a - 1.0);
            if (ky > 0)
            {
              result.col(idx(p, 0, 0))
                  -= ky * (x.col(1) + x.col(2))
                     * dresult[idx(kx, ky - 1, kz)].col(idx(p - 2, 0, 0))
                     * (a - 1.0);
            }

            if (ky > 1)
            {
              result.col(idx(p, 0, 0))
                  -= ky * (ky - 1)
                     * dresult[idx(kx, ky - 2, kz)].col(idx(p - 2, 0, 0))
                     * (a - 1.0);
            }

            if (kz > 0)
            {
              result.col(idx(p, 0, 0))
                  -= kz * (x.col(1) + x.col(2))
                     * dresult[idx(kx, ky, kz - 1)].col(idx(p - 2, 0, 0))
                     * (a - 1.0);
            }

            if (kz > 1)
            {
              result.col(idx(p, 0, 0))
                  -= kz * (kz - 1)
                     * dresult[idx(kx, ky, kz - 2)].col(idx(p - 2, 0, 0))
                     * (a - 1.0);
            }

            if (ky > 0 and kz > 0)
            {
              result.col(idx(p, 0, 0))
                  -= 2.0 * ky * kz
                     * dresult[idx(kx, ky - 1, kz - 1)].col(idx(p - 2, 0, 0))
                     * (a - 1.0);
            }
          }
        }

        for (int p = 0; p < n; ++p)
        {
          result.col(idx(p, 1, 0))
              = result.col(idx(p, 0, 0))
                * ((1.0 + x.col(1)) * p
                   + (2.0 + x.col(1) * 3.0 + x.col(2)) * 0.5);
          if (ky > 0)
          {
            result.col(idx(p, 1, 0))
                += 2 * ky * dresult[idx(kx, ky - 1, kz)].col(idx(p, 0, 0))
                   * (1.5 + p);
          }

          if (kz > 0)
          {
            result.col(idx(p, 1, 0))
                += kz * dresult[idx(kx, ky, kz - 1)].col(idx(p, 0, 0));
          }

          for (int q = 1; q < n - p; ++q)
          {
            auto [aq, bq, cq] = jrc(2 * p + 1, q);
            result.col(idx(p, q + 1, 0))
                = result.col(idx(p, q, 0)) * (f3 * aq + f4 * bq)
                  - result.col(idx(p, q - 1, 0)) * f5 * cq;

            if (ky > 0)
            {
              result.col(idx(p, q + 1, 0))
                  += 2 * ky * dresult[idx(kx, ky - 1, kz)].col(idx(p, q, 0))
                     * aq;
            }

            if (kz > 0)
            {
              result.col(idx(p, q + 1, 0))
                  += kz * dresult[idx(kx, ky, kz - 1)].col(idx(p, q, 0))
                         * (aq - bq)
                     + kz * (1.0 - x.col(2))
                           * dresult[idx(kx, ky, kz - 1)].col(idx(p, q - 1, 0))
                           * cq;
            }

            if (kz > 1)
            {
              // Quadratic term in z
              result.col(idx(p, q + 1, 0))
                  -= kz * (kz - 1)
                     * dresult[idx(kx, ky, kz - 2)].col(idx(p, q - 1, 0)) * cq;
            }
          }
        }

        for (int p = 0; p < n; ++p)
        {
          for (int q = 0; q < n - p; ++q)
          {
            result.col(idx(p, q, 1))
                = result.col(idx(p, q, 0))
                  * ((1.0 + p + q) + x.col(2) * (2.0 + p + q));
            if (kz > 0)
            {
              result.col(idx(p, q, 1))
                  += 2 * kz * (2.0 + p + q)
                     * dresult[idx(kx, ky, kz - 1)].col(idx(p, q, 0));
            }
          }
        }

        for (int p = 0; p < n - 1; ++p)
        {
          for (int q = 0; q < n - p - 1; ++q)
          {
            for (int r = 1; r < n - p - q; ++r)
            {
              auto [ar, br, cr] = jrc(2 * p + 2 * q + 2, r);
              result.col(idx(p, q, r + 1))
                  = result.col(idx(p, q, r)) * (x.col(2) * ar + br)
                    - result.col(idx(p, q, r - 1)) * cr;
              if (kz > 0)
              {
                result.col(idx(p, q, r + 1))
                    += 2 * kz * ar
                       * dresult[idx(kx, ky, kz - 1)].col(idx(p, q, r));
              }
            }
          }
        }

        // Store this derivative
        dresult[idx(kx, ky, kz)] = result;
      }
    }
  }

  for (Eigen::ArrayXXd& result : dresult)
  {
    for (int p = 0; p < n + 1; ++p)
    {
      for (int q = 0; q < n - p + 1; ++q)
      {
        for (int r = 0; r < n - p - q + 1; ++r)
        {
          result.col(idx(p, q, r))
              *= std::sqrt((p + 0.5) * (p + q + 1.0) * (p + q + r + 1.5));
        }
      }
    }
  }

  return dresult;
}
//-----------------------------------------------------------------------------
std::vector<Eigen::ArrayXXd>
tabulate_polyset_pyramid_derivs(int n, int nderiv, const Eigen::ArrayXXd& pts)
{
  assert(pts.cols() == 3);

  Eigen::ArrayXXd x = pts * 2.0 - 1.0;

  const int m = (n + 1) * (n + 2) * (2 * n + 3) / 6;
  const int md = (nderiv + 1) * (nderiv + 2) * (nderiv + 3) / 6;
  std::vector<Eigen::ArrayXXd> dresult(md);

  // Indexing for pyramidal basis functions
  auto pyr_idx = [&n](int p, int q, int r) -> int {
    const int rv = n - r + 1;
    const int r0 = r * (n + 1) * (n - r + 2) + (2 * r - 1) * (r - 1) * r / 6;
    return r0 + p * rv + q;
  };

  const Eigen::ArrayXd f2 = (1.0 - x.col(2)).square() * 0.25;

  // Traverse derivatives in increasing order
  Eigen::ArrayXXd result(pts.rows(), m);
  for (int k = 0; k < nderiv + 1; ++k)
  {
    for (int j = 0; j < k + 1; ++j)
    {
      for (int kx = 0; kx < j + 1; ++kx)
      {
        const int ky = j - kx;
        const int kz = k - j;
        result.setZero();

        const int pyramidal_index = pyr_idx(0, 0, 0);
        assert(pyramidal_index < m);
        if (kx == 0 and ky == 0 and kz == 0)
          result.col(pyramidal_index).fill(1.0);
        else
          result.col(pyramidal_index).setZero();

        // r = 0
        for (int p = 0; p < n + 1; ++p)
        {
          if (p > 0)
          {
            const double a
                = static_cast<double>(p - 1) / static_cast<double>(p);
            result.col(pyr_idx(p, 0, 0)) = (0.5 + x.col(0) + x.col(2) * 0.5)
                                           * result.col(pyr_idx(p - 1, 0, 0))
                                           * (a + 1.0);
            if (kx > 0)
            {
              result.col(pyr_idx(p, 0, 0))
                  += 2.0 * kx
                     * dresult[idx(kx - 1, ky, kz)].col(pyr_idx(p - 1, 0, 0))
                     * (a + 1.0);
            }

            if (kz > 0)
            {
              result.col(pyr_idx(p, 0, 0))
                  += kz * dresult[idx(kx, ky, kz - 1)].col(pyr_idx(p - 1, 0, 0))
                     * (a + 1.0);
            }

            if (p > 1)
            {
              result.col(pyr_idx(p, 0, 0))
                  -= f2 * result.col(pyr_idx(p - 2, 0, 0)) * a;

              if (kz > 0)
              {
                result.col(pyr_idx(p, 0, 0))
                    += kz * (1.0 - x.col(2))
                       * dresult[idx(kx, ky, kz - 1)].col(pyr_idx(p - 2, 0, 0))
                       * a;
              }

              if (kz > 1)
              {
                // quadratic term in z
                result.col(pyr_idx(p, 0, 0))
                    -= kz * (kz - 1)
                       * dresult[idx(kx, ky, kz - 2)].col(pyr_idx(p - 2, 0, 0))
                       * a;
              }
            }
          }

          for (int q = 1; q < n + 1; ++q)
          {
            const double a
                = static_cast<double>(q - 1) / static_cast<double>(q);
            result.col(pyr_idx(p, q, 0)) = (0.5 + x.col(1) + x.col(2) * 0.5)
                                           * result.col(pyr_idx(p, q - 1, 0))
                                           * (a + 1.0);
            if (ky > 0)
            {
              result.col(pyr_idx(p, q, 0))
                  += 2.0 * ky
                     * dresult[idx(kx, ky - 1, kz)].col(pyr_idx(p, q - 1, 0))
                     * (a + 1.0);
            }

            if (kz > 0)
            {
              result.col(pyr_idx(p, q, 0))
                  += kz * dresult[idx(kx, ky, kz - 1)].col(pyr_idx(p, q - 1, 0))
                     * (a + 1.0);
            }

            if (q > 1)
            {
              result.col(pyr_idx(p, q, 0))
                  -= f2 * result.col(pyr_idx(p, q - 2, 0)) * a;

              if (kz > 0)
              {
                result.col(pyr_idx(p, q, 0))
                    += kz * (1.0 - x.col(2))
                       * dresult[idx(kx, ky, kz - 1)].col(pyr_idx(p, q - 2, 0))
                       * a;
              }

              if (kz > 1)
              {
                result.col(pyr_idx(p, q, 0))
                    -= kz * (kz - 1)
                       * dresult[idx(kx, ky, kz - 2)].col(pyr_idx(p, q - 2, 0))
                       * a;
              }
            }
          }
        }

        // Extend into r > 0
        for (int p = 0; p < n; ++p)
        {
          for (int q = 0; q < n; ++q)
          {
            result.col(pyr_idx(p, q, 1))
                = result.col(pyr_idx(p, q, 0))
                  * ((1.0 + p + q) + x.col(2) * (2.0 + p + q));
            if (kz > 0)
            {
              result.col(pyr_idx(p, q, 1))
                  += 2 * kz * dresult[idx(kx, ky, kz - 1)].col(pyr_idx(p, q, 0))
                     * (2.0 + p + q);
            }
          }
        }

        for (int r = 1; r < n + 1; ++r)
        {
          for (int p = 0; p < n - r; ++p)
          {
            for (int q = 0; q < n - r; ++q)
            {
              auto [ar, br, cr] = jrc(2 * p + 2 * q + 2, r);
              result.col(pyr_idx(p, q, r + 1))
                  = result.col(pyr_idx(p, q, r)) * (x.col(2) * ar + br)
                    - result.col(pyr_idx(p, q, r - 1)) * cr;
              if (kz > 0)
              {
                result.col(pyr_idx(p, q, r + 1))
                    += ar * 2 * kz
                       * dresult[idx(kx, ky, kz - 1)].col(pyr_idx(p, q, r));
              }
            }
          }
        }

        dresult[idx(kx, ky, kz)] = result;
      }
    }
  }

  for (Eigen::ArrayXXd& result : dresult)
  {
    for (int r = 0; r < n + 1; ++r)
    {
      for (int p = 0; p < n - r + 1; ++p)

      {
        for (int q = 0; q < n - r + 1; ++q)
        {
          result.col(pyr_idx(p, q, r))
              *= std::sqrt((q + 0.5) * (p + 0.5) * (p + q + r + 1.5));
        }
      }
    }
  }

  return dresult;
}
//-----------------------------------------------------------------------------
std::vector<Eigen::ArrayXXd>
tabulate_polyset_quad_derivs(int n, int nderiv, const Eigen::ArrayXXd& pts)
{
  assert(pts.cols() == 2);
  const int m = (n + 1) * (n + 1);
  const int md = (nderiv + 1) * (nderiv + 2) / 2;

  std::vector<Eigen::ArrayXXd> dresult(md);
  std::vector<Eigen::ArrayXXd> px
      = tabulate_polyset_line_derivs(n, nderiv, pts.col(0));
  std::vector<Eigen::ArrayXXd> py
      = tabulate_polyset_line_derivs(n, nderiv, pts.col(1));

  Eigen::ArrayXXd result(pts.rows(), m);
  for (int kx = 0; kx < nderiv + 1; ++kx)
  {
    for (int ky = 0; ky < nderiv + 1 - kx; ++ky)
    {
      int c = 0;
      for (int i = 0; i < px[kx].cols(); ++i)
        for (int j = 0; j < py[ky].cols(); ++j)
          result.col(c++) = px[kx].col(i) * py[ky].col(j);
      dresult[idx(kx, ky)] = result;
    }
  }

  return dresult;
}
//-----------------------------------------------------------------------------
std::vector<Eigen::ArrayXXd>
tabulate_polyset_hex_derivs(int n, int nderiv, const Eigen::ArrayXXd& pts)
{
  assert(pts.cols() == 3);
  const int m = (n + 1) * (n + 1) * (n + 1);
  const int md = (nderiv + 1) * (nderiv + 2) * (nderiv + 3) / 6;

  std::vector<Eigen::ArrayXXd> px
      = tabulate_polyset_line_derivs(n, nderiv, pts.col(0));
  std::vector<Eigen::ArrayXXd> py
      = tabulate_polyset_line_derivs(n, nderiv, pts.col(1));
  std::vector<Eigen::ArrayXXd> pz
      = tabulate_polyset_line_derivs(n, nderiv, pts.col(2));

  std::vector<Eigen::ArrayXXd> dresult(md);
  Eigen::ArrayXXd result(pts.rows(), m);
  for (int kx = 0; kx < nderiv + 1; ++kx)
  {
    for (int ky = 0; ky < nderiv + 1 - kx; ++ky)
    {
      for (int kz = 0; kz < nderiv + 1 - kx - ky; ++kz)
      {
        int c = 0;
        for (int i = 0; i < px[kx].cols(); ++i)
          for (int j = 0; j < py[ky].cols(); ++j)
            for (int k = 0; k < pz[kz].cols(); ++k)
              result.col(c++) = px[kx].col(i) * py[ky].col(j) * pz[kz].col(k);

        dresult[idx(kx, ky, kz)] = result;
      }
    }
  }

  return dresult;
}
//-----------------------------------------------------------------------------
std::vector<Eigen::ArrayXXd>
tabulate_polyset_prism_derivs(int n, int nderiv, const Eigen::ArrayXXd& pts)
{
  assert(pts.cols() == 3);
  const int m = (n + 1) * (n + 1) * (n + 2) / 2;
  const int md = (nderiv + 1) * (nderiv + 2) * (nderiv + 3) / 6;

  std::vector<Eigen::ArrayXXd> pxy
      = tabulate_polyset_triangle_derivs(n, nderiv, pts.leftCols(2));
  std::vector<Eigen::ArrayXXd> pz
      = tabulate_polyset_line_derivs(n, nderiv, pts.col(2));

  std::vector<Eigen::ArrayXXd> dresult(md);
  Eigen::ArrayXXd result(pts.rows(), m);
  for (int kx = 0; kx < nderiv + 1; ++kx)
  {
    for (int ky = 0; ky < nderiv + 1 - kx; ++ky)
    {
      for (int kz = 0; kz < nderiv + 1 - kx - ky; ++kz)
      {
        int c = 0;
        for (int i = 0; i < pxy[idx(kx, ky)].cols(); ++i)
          for (int k = 0; k < pz[kz].cols(); ++k)
            result.col(c++) = pxy[idx(kx, ky)].col(i) * pz[kz].col(k);

        dresult[idx(kx, ky, kz)] = result;
      }
    }
  }

  return dresult;
}
} // namespace
//-----------------------------------------------------------------------------
std::vector<Eigen::ArrayXXd> polyset::tabulate(cell::type celltype, int n,
                                               int nderiv,
                                               const Eigen::ArrayXXd& pts)
{
  switch (celltype)
  {
  case cell::type::interval:
    return tabulate_polyset_line_derivs(n, nderiv, pts);
  case cell::type::triangle:
    return tabulate_polyset_triangle_derivs(n, nderiv, pts);
  case cell::type::tetrahedron:
    return tabulate_polyset_tetrahedron_derivs(n, nderiv, pts);
  case cell::type::quadrilateral:
    return tabulate_polyset_quad_derivs(n, nderiv, pts);
  case cell::type::prism:
    return tabulate_polyset_prism_derivs(n, nderiv, pts);
  case cell::type::pyramid:
    return tabulate_polyset_pyramid_derivs(n, nderiv, pts);
  case cell::type::hexahedron:
    return tabulate_polyset_hex_derivs(n, nderiv, pts);
  default:
    throw std::runtime_error("Polynomial set: Unsupported cell type");
  }
}
//-----------------------------------------------------------------------------
int polyset::dim(cell::type celltype, int n)
{
  switch (celltype)
  {
  case cell::type::triangle:
    return (n + 1) * (n + 2) / 2;
  case cell::type::tetrahedron:
    return (n + 1) * (n + 2) * (n + 3) / 6;
  case cell::type::prism:
    return (n + 1) * (n + 1) * (n + 2) / 2;
  case cell::type::pyramid:
    return (n + 1) * (n + 2) * (2 * n + 3) / 6;
  case cell::type::interval:
    return (n + 1);
  case cell::type::quadrilateral:
    return (n + 1) * (n + 1);
  case cell::type::hexahedron:
    return (n + 1) * (n + 1) * (n + 1);
  default:
    return 1;
  }
}
//-----------------------------------------------------------------------------


//
// Copyright (c) 2009, Markus Rickert
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// * Redistributions of source code must retain the above copyright notice,
//   this list of conditions and the following disclaimer.
// * Redistributions in binary form must reproduce the above copyright notice,
//   this list of conditions and the following disclaimer in the documentation
//   and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.
//

#include <iostream>
#include <boost/lexical_cast.hpp>
#include <rl/math/Polynomial.h>

int
main(int argc, char** argv)
{
	if (argc < 2)
	{
		std::cout << "Usage: rlPolynomialRootsDemo C0 ... CN" << std::endl;
		return EXIT_FAILURE;
	}

	std::vector<rl::math::Real> c(argc - 1);

	for (std::size_t i = 0; i < c.size(); ++i)
	{
		c[i] = boost::lexical_cast<rl::math::Real>(argv[i + 1]);
		std::cout << (i > 0 ? " + " : "") << c[i] << " * x^" << i;
	}

	std::cout << " = 0" << std::endl;

	std::vector<rl::math::Real> roots = rl::math::Polynomial<rl::math::Real>::realRoots(c);

	std::cout << roots.size() << " solution" << (roots.size() != 1 ? "(s)" : "") << std::endl;

	for (std::size_t i = 0; i < roots.size(); ++i)
	{
		std::cout << "x[" << i << "] = " << roots[i] << std::endl;
	}

	return EXIT_SUCCESS;
}


// Copyright (c) 2013, Manuel Blum
// All rights reserved.

// Define this symbol to enable runtime tests for allocations
//#define EIGEN_RUNTIME_NO_MALLOC

#include <Eigen/Dense>
#include <iostream>
#include <fstream>
#include <cstdio>
#include <string>
#include <ctime>
#include "nn.h"

using namespace std;

typedef vector<int> vi;
typedef vector<vi> vvi;
inline void swap(int &val)
{
    val = (val << 24) | ((val << 8) & 0x00ff0000) | ((val >> 8) & 0x0000ff00) | (val >> 24);
}
class batch_manager
{
    int num_samples;
    int current_count;
    int batch_size;

  public:
    batch_manager(int num_samples, int batch_size)
    {
        this->num_samples = num_samples;
        this->batch_size = batch_size;
        current_count = 0;
    }
    pair<int, int> next_batch()
    {
        pair<int, int> res = make_pair(current_count, current_count + batch_size);
        if (current_count + batch_size > num_samples)
        {
            //cout << "Current count: " << current_count << endl;
            if (current_count == num_samples)
            {
                current_count = 0;
                return make_pair(0, batch_size);
            }
            res.second = num_samples;
            current_count = 0;

            return res;
        }
        current_count += res.second - res.first;
        return res;
    }
};
matrix_t read_mnist_images(std::string filename)
{

    matrix_t X;
    std::ifstream fs(filename.c_str(), std::ios::binary);
    if (fs)
    {
        int magic_number, num_images, num_rows, num_columns;
        fs.read((char *)&magic_number, sizeof(magic_number));
        fs.read((char *)&num_images, sizeof(num_images));
        fs.read((char *)&num_rows, sizeof(num_rows));
        fs.read((char *)&num_columns, sizeof(num_columns));
        if (magic_number != 2051)
        {
            swap(magic_number);
            swap(num_images);
            swap(num_rows);
            swap(num_columns);
        }

        X = matrix_t::Zero(num_images, num_rows * num_columns);

        for (size_t i = 0; i < num_images; ++i)
        {
            for (size_t j = 0; j < num_rows * num_columns; ++j)
            {
                unsigned char temp = 0;
                fs.read((char *)&temp, sizeof(temp));
                X(i, j) = (double)temp;
            }
        }
        fs.close();
    }
    else
    {
        std::cout << "error reading file: " << filename << std::endl;
        exit(1);
    }
    return X;
}

matrix_t read_mnist_labels(std::string filename)
{
    matrix_t Y;
    std::ifstream fs(filename.c_str(), std::ios::binary);
    if (fs)
    {
        int magic_number, num_images, num_rows, num_columns;
        fs.read((char *)&magic_number, sizeof(magic_number));
        fs.read((char *)&num_images, sizeof(num_images));
        if (magic_number != 2049)
        {
            swap(magic_number);
            swap(num_images);
        }

        Y = matrix_t::Zero(num_images, 10);

        for (size_t i = 0; i < num_images; ++i)
        {
            unsigned char temp = 0;
            fs.read((char *)&temp, sizeof(temp));
            Y(i, (int)temp) = 1.0;
        }
        fs.close();
    }
    else
    {
        std::cout << "error reading file: " << filename << std::endl;
        exit(1);
    }
    return Y;
}

double measure_accuracy(matrix_t Y, matrix_t Y_test)
{
    int correct = 0;
    for (int i = 0; i < Y.rows(); i++)
    {
        int j = 0;
        for (; j < Y.cols(); j++)
        {
            int out = Y_test(i, j) > 0.7 ? 1 : 0;
            //printf("%3i   %4.4f", (int)Y(i, j), Y_test(i, j));

            if (out != Y(i, j))
                break;
        }
        //cout << endl;
        if (j == Y.cols())
        {
            correct++;
        }
        else
        {
            //cout << Y_test.row(i) << endl;
        }
    }
    //cout << "Correct answer " << correct << " among " << Y.rows() << " test" << endl;
    double accuracy = (double)correct / (double)Y.rows() * 100;
    return accuracy;
    //cout << "Accuracy: " <<  accuracy << "%" << endl;
}

void fill_type_info(matrix_t &Y, vi &type_info)
{
    int num_sample = Y.rows();
    int num_type = Y.cols();
    for (int i = 0; i < num_sample; i++)
    {
        for (int j = 0; j < num_type; j++)
        {
            if (Y(i, j) == 1)
            {
                type_info[j]++;
                break;
            }
        }
    }
}

void fill_position(matrix_t &Y, vvi &position)
{
    int num_sample = Y.rows();
    int num_type = Y.cols();
    for (int i = 0; i < num_sample; i++)
    {
        for (int j = 0; j < num_type; j++)
        {
            if (Y(i, j) == 1)
            {
                position[j].push_back(i);
                break;
            }
        }
    }
}

int distribution_in_batch(vi &type_distribution, vi &type_info, int proposed_batch_size, int num_sample)
{
    int batch_size = 0;
    int num_type = type_info.size();
    for (int i = 0; i < num_type; i++)
    {
        type_distribution[i] = (int)ceil((double)(type_info[i] * proposed_batch_size) / (double)num_sample);
        batch_size += type_distribution[i];
    }
    return batch_size;
}

void distribute_into_batch(matrix_t &X, matrix_t &Y, matrix_t &Xm, matrix_t &Ym, vi &type_distribution, vvi &position, int batch_size, vi &type_point, int batch_no)
{
    int cnt = 0;
    for (int i = 0; i < type_distribution.size(); i++)
    {
        for (int j = 0; j < type_distribution[i]; j++)
        {
            if (batch_no * batch_size + cnt == Xm.rows())
            {
                Xm.conservativeResize(Xm.rows() + batch_size - cnt, Xm.cols());
                Ym.conservativeResize(Ym.rows() + batch_size - cnt, Ym.cols());
            }
            Xm.row(batch_no * batch_size + cnt) = X.row(position[i][type_point[i]]);
            Ym.row(batch_no * batch_size + cnt) = Y.row(position[i][type_point[i]]);
            cnt++;
            type_point[i]++;
            if (type_point[i] >= position[i].size())
            {
                type_point[i] = 0;
            }
        }
    }
}
void distribute_into_matrix(matrix_t &X, matrix_t &Y, matrix_t &Xm, matrix_t &Ym, vi &type_info, vi &type_distribution, vvi &position, int batch_size)
{
    vi type_point(type_info.size(), 0);
    int num_batch = (int)ceil((double)Y.rows() / (double)batch_size);

    for (int i = 0; i < num_batch; i++)
    {
        distribute_into_batch(X, Y, Xm, Ym, type_distribution, position, batch_size, type_point, i);
    }
}
pair<matrix_t, matrix_t> make_uniform_dataset(matrix_t &X, matrix_t &Y, int &batch_size)
{
    int num_sample = Y.rows();
    int num_type = Y.cols();

    vi type_info(num_type, 0);
    fill_type_info(Y, type_info);

    vi type_distribution(num_type, 0);
    batch_size = distribution_in_batch(type_distribution, type_info, batch_size, num_sample);

    vvi position(num_type);
    fill_position(Y, position);

    matrix_t Xm(num_sample, X.cols());
    matrix_t Ym(num_sample, num_type);

    distribute_into_matrix(X, Y, Xm, Ym, type_info, type_distribution, position, batch_size);

    return make_pair(Xm, Ym);
}
pair<matrix_t, matrix_t> make_worst_batch(matrix_t &X, matrix_t &Y)
{
    matrix_t Xm(X.rows(), X.cols());
    matrix_t Ym(Y.rows(), Y.cols());

    int num_sample = Y.rows();
    int num_type = Y.cols();

    vi type_info(num_type, 0);
    fill_type_info(Y, type_info);

    vvi position(num_type);
    fill_position(Y, position);
    int cnt = 0;
    for (int i = 0; i < num_type; i++)
    {
        for (int j = 0; j < position[i].size(); j++)
        {
            Xm.row(cnt) = X.row(position[i][j]);
            Ym.row(cnt) = Y.row(position[i][j]);
            cnt++;
        }
    }

    return make_pair(Xm, Ym);
}
int main(int argc, const char *argv[])
{

    if (argc != 2)
    {
        std::cout << "please provide path to mnist data ..." << std::endl;
        std::cout << "you can download the dataset at http://yann.lecun.com/exdb/mnist/" << std::endl;
        std::cout << std::endl
                  << "usage: " << argv[0] << " path_to_data" << std::endl
                  << std::endl;
        return 1;
    }

    std::string path = argv[1];

    std::cout << "reading data" << std::endl;

    // matrix_t X_train = read_mnist_images(path + "/train-images.idx3-ubyte");
    // matrix_t Y_train = read_mnist_labels(path + "/train-labels.idx1-ubyte");
    // matrix_t X_test = read_mnist_images(path + "/t10k-images.idx3-ubyte");
    // matrix_t Y_test = read_mnist_labels(path + "/t10k-labels.idx1-ubyte");

    matrix_t X_train = read_mnist_images(path + "/train-images.idx3-ubyte");
    matrix_t Y_train = read_mnist_labels(path + "/train-labels.idx1-ubyte");
    matrix_t X_test = read_mnist_images(path + "/t10k-images.idx3-ubyte");
    matrix_t Y_test = read_mnist_labels(path + "/t10k-labels.idx1-ubyte");

    //cout << "Number of training sample: " << X_train.rows() << endl;
    int max_steps = 1500;
    double lambda = 0.001;

    // specify network topology
    Eigen::VectorXi topo(3);
    topo << X_train.cols(), 300, Y_test.cols();
    std::cout << "topology: " << topo.transpose() << std::endl;

    // initialize a neural network with given topology
    std::cout << "initializing network" << std::endl;
    NeuralNet nn(topo);

    int batch_size = 100;
    // pair<matrix_t, matrix_t> xy = make_uniform_dataset(X_train, Y_train, batch_size);

    // X_train = xy.first;
    // Y_train = xy.second;

    // pair<matrix_t, matrix_t> xy = make_worst_batch(X_train, Y_train);

    // X_train = xy.first;
    // Y_train = xy.second;

    std::cout<< "scaling the data" << std::endl;
    nn.autoscale(X_train, Y_train);

    int num_attribute = X_train.cols();
    int num_type = Y_train.cols();
    batch_manager batch(X_train.rows(), batch_size);

    std::cout << "starting training" << std::endl;
    std::cout << "iter        error" << std::endl;

    double err;
    clock_t begin = clock();
    for (int i = 0; i < max_steps; ++i)
    {
        pair<int, int> start_end = batch.next_batch();
        const int batch_size = start_end.second - start_end.first;

        matrix_t Xm = X_train.block(start_end.first, 0, batch_size, num_attribute);
        matrix_t Ym = Y_train.block(start_end.first, 0, batch_size, num_type);

        err = nn.loss(Xm, Ym, lambda);
        nn.rprop();
        printf("%4i   %10.7f\n", i, err);
    }
    clock_t end = clock();
    double elapsed_secs = double(end - begin) / CLOCKS_PER_SEC;
    cout << "Training time: " << elapsed_secs << endl;

    // test accuracy
    nn.forward_pass(X_test);
    matrix_t prediction = nn.get_activation();
    int correct = 0;
    int k;
    for (size_t i = 0; i < Y_test.rows(); ++i)
    {
        prediction.row(i).maxCoeff(&k);
        correct += Y_test(i, k);
    }

    std::cout << "test accuracy: " << correct * 1.0 / Y_test.rows() * 100 << "%" << std::endl;
    // double accuracy = measure_accuracy(Y_test, prediction);
    // std::cout << "Accuracy: " << accuracy << std::endl;
    nn.write("mnist.nn");

    return 0;
}


// Copyright 2018-2019 Hans Dembinski and Henry Schreiner
//
// Distributed under the Boost Software License, version 1.0.
// (See accompanying file LICENSE_1_0.txt
// or copy at http://www.boost.org/LICENSE_1_0.txt)
//
// Based on boost/histogram/accumulators/weighted_mean.hpp
//
// Changes:
//  * Internal values are public for access from Python
//  * A special constructor added for construction from Python

#pragma once

#include <boost/core/nvp.hpp>
#include <boost/histogram/weight.hpp>

namespace accumulators {

/**
  Calculates mean and variance of weighted sample.

  Uses West's incremental algorithm to improve numerical stability
  of mean and variance computation.
*/
template <typename ValueType>
struct weighted_mean {
    using value_type      = ValueType;
    using const_reference = const value_type&;

    weighted_mean() = default;

    weighted_mean(const value_type& wsum,
                  const value_type& wsum2,
                  const value_type& mean,
                  const value_type& variance)
        : sum_of_weights(wsum)
        , sum_of_weights_squared(wsum2)
        , value(mean)
        , sum_of_weighted_deltas_squared(
              variance * (sum_of_weights - sum_of_weights_squared / sum_of_weights)) {}

    weighted_mean(const value_type& wsum,
                  const value_type& wsum2,
                  const value_type& mean,
                  const value_type& sum_of_weighted_deltas_squared,
                  bool /* tag to trigger Python internal constructor */)
        : sum_of_weights(wsum)
        , sum_of_weights_squared(wsum2)
        , value(mean)
        , sum_of_weighted_deltas_squared(sum_of_weighted_deltas_squared) {}

    void operator()(const value_type& x) { operator()(boost::histogram::weight(1), x); }

    void operator()(const boost::histogram::weight_type<value_type>& w,
                    const value_type& x) {
        sum_of_weights += w.value;
        sum_of_weights_squared += w.value * w.value;
        const auto delta = x - value;
        value += w.value * delta / sum_of_weights;
        sum_of_weighted_deltas_squared += w.value * delta * (x - value);
    }

    weighted_mean& operator+=(const weighted_mean& rhs) {
        if(sum_of_weights != 0 || rhs.sum_of_weights != 0) {
            const auto tmp = value * sum_of_weights + rhs.value * rhs.sum_of_weights;
            sum_of_weights += rhs.sum_of_weights;
            sum_of_weights_squared += rhs.sum_of_weights_squared;
            value = tmp / sum_of_weights;
        }
        sum_of_weighted_deltas_squared += rhs.sum_of_weighted_deltas_squared;
        return *this;
    }

    weighted_mean& operator*=(const value_type& s) {
        value *= s;
        sum_of_weighted_deltas_squared *= s * s;
        return *this;
    }

    bool operator==(const weighted_mean& rhs) const noexcept {
        return sum_of_weights == rhs.sum_of_weights
               && sum_of_weights_squared == rhs.sum_of_weights_squared
               && value == rhs.value
               && sum_of_weighted_deltas_squared == rhs.sum_of_weighted_deltas_squared;
    }

    bool operator!=(const weighted_mean rhs) const noexcept { return !operator==(rhs); }

    value_type variance() const {
        return sum_of_weighted_deltas_squared
               / (sum_of_weights - sum_of_weights_squared / sum_of_weights);
    }

    template <class Archive>
    void serialize(Archive& ar, unsigned /* version */) {
        ar& boost::make_nvp("sum_of_weights", sum_of_weights);
        ar& boost::make_nvp("sum_of_weights_squared", sum_of_weights_squared);
        ar& boost::make_nvp("value", value);
        ar& boost::make_nvp("sum_of_weighted_deltas_squared",
                            sum_of_weighted_deltas_squared);
    }

    value_type sum_of_weights{};
    value_type sum_of_weights_squared{};
    value_type value{};
    value_type sum_of_weighted_deltas_squared{};
};

} // namespace accumulators


#include <iostream>
#include <chrono>
#include <memory>
#include <cmath>
#include <Eigen/Dense>
#include <Eigen/Dense>
#include "minimize.h"

using Matrix = minimize::Matrix;
using Vector = minimize::Vector;

// Define interpolation procedure for minimization algorithm
double minimize::interpolate(double x2, double f2, double d2, double x3, double f3, double d3, double f0, double INT, double RHO)
{

  // choose subinterval
  // move point 3 to point 4
  double x4 = x3;
  double f4 = f3;
  double d4 = d3;

  //double tolerance = 1e-32;
  double tolerance = 1e-64;

  if ( f4 > f0 )
    {
      double denom = f4-f2-d2*(x4-x2);
      if ( std::abs(denom) < tolerance )
        // bisect
        x3 = (x2+x4)/2;
      else
        // quadratic interpolation
        x3 = x2-(0.5*d2*std::pow(x4-x2,2))/(denom);
    }
  else
    {
      // cubic interpolation
      double A = 6*(f2-f4)/(x4-x2)+3*(d4+d2);
      double B = 3*(f4-f2)-(2*d2+d4)*(x4-x2);
      double radical = B*B-A*d2*std::pow(x4-x2,2);

      if ( ( radical < 0 ) || ( std::abs(A) < tolerance ) )
        x3 = (x2+x4)/2;
      else
        x3 = x2+( std::sqrt(radical) - B)/A;
    }

  // don't accept too close
  if ( x4-INT*(x4-x2) < x3 )
    x3 = x4-INT*(x4-x2);

  if ( x2+INT*(x4-x2) > x3 )
    x3 = x2+INT*(x4-x2);

  return x3;
};

// Define cubic extrapolation routine for minimization algorithm
double minimize::cubic_extrap(double x1, double x2, double f1, double f2, double d1, double d2, double EXT, double INT)
{
  // make cubic extrapolation
  double A = 6*(f1-f2)+3*(d2+d1)*(x2-x1);
  double B = 3*(f2-f1)-(2*d1+d2)*(x2-x1);

  double x3;
  //double tolerance = 1e-32;
  double tolerance = 1e-64;
  double radical = B*B-A*d1*(x2-x1);

  if ( radical < 0.0 )
    x3 = x2*EXT;
  else if ( B + std::sqrt(radical) < tolerance )
    x3 = x2*EXT;
  else
    {
      x3 = x1-d1*std::pow(x2-x1,2)/( B + std::sqrt(radical) );

      if ( ( x3 < 0 ) || ( x3 > x2*EXT ) )
        x3 = x2*EXT;
      else if ( x3 < x2+INT*(x2-x1) )
        x3 = x2+INT*(x2-x1);
    }

  return x3;
};

//
//  ORIGINAL CODE BY CARL EDWARD RASMUSSEN
//  http://learning.eng.cam.ac.uk/carl/code/minimize/
//
//  % Minimize a differentiable multivariate function.
//  %
//  % Usage: [X, fX, i] = minimize(X, f, length, P1, P2, P3, ... )
//  %
//  % where the starting point is given by "X" (D by 1), and the function named in
//  % the string "f", must return a function value and a vector of partial
//  % derivatives of f wrt X, the "length" gives the length of the run: if it is
//  % positive, it gives the maximum number of line searches, if negative its
//  % absolute gives the maximum allowed number of function evaluations. You can
//  % (optionally) give "length" a second component, which will indicate the
//  % reduction in function value to be expected in the first line-search (defaults
//  % to 1.0). The parameters P1, P2, P3, ... are passed on to the function f.
//  %
//  % The function returns when either its length is up, or if no further progress
//  % can be made (ie, we are at a (local) minimum, or so close that due to
//  % numerical problems, we cannot get any closer). NOTE: If the function
//  % terminates within a few iterations, it could be an indication that the
//  % function values and derivatives are not consistent (ie, there may be a bug in
//  % the implementation of your "f" function). The function returns the found
//  % solution "X", a vector of function values "fX" indicating the progress made
//  % and "i" the number of iterations (line searches or function evaluations,
//  % depending on the sign of "length") used.
//  %
//  % The Polack-Ribiere flavour of conjugate gradients is used to compute search
//  % directions, and a line search using quadratic and cubic polynomial
//  % approximations and the Wolfe-Powell stopping criteria is used together with
//  % the slope ratio method for guessing initial step sizes. Additionally a bunch
//  % of checks are made to make sure that exploration is taking place and that
//  % extrapolation will not be unboundedly large.
//  %
//  % See also: checkgrad
//  %
//  % Copyright (C) 2001 - 2006 by Carl Edward Rasmussen (2006-09-08).
//
//  INT = 0.1;    % don't reevaluate within 0.1 of the limit of the current bracket
//  EXT = 3.0;                  % extrapolate maximum 3 times the current step-size
//  MAX = 20;                         % max 20 function evaluations per line search
//  RATIO = 10;                                       % maximum allowed slope ratio
//  SIG = 0.1; RHO = SIG/2; % SIG and RHO are the constants controlling the Wolfe-
//  % Powell conditions. SIG is the maximum allowed absolute ratio between
//  % previous and new slopes (derivatives in the search direction), thus setting
//  % SIG to low (positive) values forces higher precision in the line-searches.
//  % RHO is the minimum allowed fraction of the expected (from the slope at the
//  % initial point in the linesearch). Constants must satisfy 0 < RHO < SIG < 1.
//  % Tuning of SIG (depending on the nature of the function to be optimized) may
//  % speed up the minimization; it is probably not worth playing much with RHO.
//
//  % The code falls naturally into 3 parts, after the initial line search is
//  % started in the direction of steepest descent. 1) we first enter a while loop
//  % which uses point 1 (p1) and (p2) to compute an extrapolation (p3), until we
//  % have extrapolated far enough (Wolfe-Powell conditions). 2) if necessary, we
//  % enter the second loop which takes p2, p3 and p4 chooses the subinterval
//  % containing a (local) minimum, and interpolates it, unil an acceptable point
//  % is found (Wolfe-Powell conditions). Note, that points are always maintained
//  % in order p0 <= p1 <= p2 < p3 < p4. 3) compute a new search direction using
//  % conjugate gradients (Polack-Ribiere flavour), or revert to steepest if there
//  % was a problem in the previous line-search. Return the best value so far, if
//  % two consecutive line-searches fail, or whenever we run out of function
//  % evaluations or line-searches. During extrapolation, the "f" function may fail
//  % either with an error or returning Nan or Inf, and minimize should handle this
//  % gracefully.
//

// Conjugate gradient minimization algorithm
void minimize::cg_minimize(Vector & X, minimize::GradientObj * target, Vector & D, int length, double SIG, double EXT, double INT, int MAX)
{
  // specify optimization hyperparameters
  //double RATIO = 10.0;
  double RATIO = 100.0;
  double RHO = SIG/2;

  // determine problem dimension
  int N = static_cast<int>(X.size());

  // initialize values
  int i = 0;
  bool ls_failed = false;
  Vector df0(N);
  double f0;
  (*target).computeValueAndGradient(X, f0, df0);

  // initial search direction (steepest) and slope
  Vector s = -df0;
  double d0 = -s.transpose()*s;

  // initial step is 1/(|s|+1)
  double x3 = 1/(1-d0);

  // declare placeholders for storing optimal values
  Vector X0(N);
  double F0;
  Vector dF0(N);

  // declare variables in main loop
  int M;
  bool continue_extrap;
  double x1, f1, d1, x2, f2, d2, f3, d3;
  Vector df3(N);

  // "realmin" = smallest positive normalized floating-point number in IEEE double precision format
  double realmin = 2.2251e-308;

  // MAIN LOOP
  bool request_break = false;
  while ( ( i < length ) && ( !request_break ) )
    {
      i++;

      // make a copy of current values
      X0 = X;
      F0 = f0;
      dF0 = df0;

      // Display current parameter values
      //std::cout << " X  =  " << X.transpose().array().exp().matrix() << std::endl;

      // initialize iteration count
      M = MAX;

      // EXTRAPOLATE
      continue_extrap = true;
      while ( continue_extrap )
        {
          x2 = 0.0;
          f2 = f0;
          d2 = d0;
          M = M - 1;
          (*target).computeValueAndGradient(X+x3*s, f3, df3);

          // keep best values
          if ( f3 < F0 )
            {
              X0 = X+x3*s;
              F0 = f3;
              dF0 = df3;
            }

          // new slope
          d3 = df3.transpose()*s;

          // are we done extrapolating?
          if ( ( ( d3 > SIG*d0 ) || ( f3 > f0+x3*RHO*d0 ) ) || ( M == 0 ) )
              continue_extrap = false;

          // move point 2 to point 1
          x1 = x2;
          f1 = f2;
          d1 = d2;
          // move point 3 to point 2
          x2 = x3;
          f2 = f3;
          d2 = d3;

          // cubic extrapolation
          x3 = cubic_extrap(x1, x2, f1, f2, d1, d2, EXT, INT);

        } // END EXTRAPOLATE

      // INTERPOLATE
      while ( ( ( std::abs(d3) > -SIG*d0 ) || ( f3 > f0+x3*RHO*d0) )  &&  ( M > 0 ) )
        {

          x3 = interpolate(x2, f2, d2, x3, f3, d3, f0, INT, RHO);

          (*target).computeValueAndGradient(X+x3*s, f3, df3);

          // keep best values
          if ( f3 < F0 )
            {
              X0 = X+x3*s;
              F0 = f3;
              dF0 = df3;
            }

          // decrement line-search count
          M = M - 1;

          // new slope
          d3 = df3.transpose()*s;

        } // END INTERPOLATE

      //  START COMPUTE NEW SEARCH DIRECTION
      if ( ( std::abs(d3) < -SIG*d0 ) && ( f3 < f0+x3*RHO*d0 ) )
        {
          // if line search succeeded
          // update variables
          X = X+x3*s;
          f0 = f3;

          // Polack-Ribiere CG direction
          s = ( (df3.transpose()*df3 - df0.transpose()*df3)(0) / (df0.transpose()*df0)(0) )*s - df3;

          // swap derivatives
          df0 = df3;
          d3 = d0;
          d0 = df0.transpose()*s;

          // new slope must be negative
          if ( d0 > 0 )
            {
              // otherwise use steepest direction
              s = -df0;
              d0 = -s.transpose()*s;
            }

          // slope ratio but max RATIO
          if ( RATIO <  d3/(d0-realmin) )
            {
              x3 = x3 * RATIO;
              //std::cout << "\n[*] RATIO parameter enforced\n";
            }
          else
              x3 = x3 * d3/(d0-realmin);

          // this line search did not fail
          ls_failed = false;
        }

      else
        {
          // restore best point so far
          X = X0;
          f0 = F0;
          df0 = dF0;

          // line search failed twice in a row
          // or we ran out of time, so we give up
          if ( ( ls_failed ) || ( i > length ) )
              request_break = true;

          // DEBUGGING INFO TO SEE WHY OPTIMIZATION EXITS EARLY
          /*
          if ( ls_failed )
            {
              std::cout << "\n[*] Line Search Failed  ( i = " << i << " )\n";
              std::cout << "\n The following conditions failed:   [ SIG = " << SIG << " , RHO = " << RHO << " ]\n";
              if ( std::abs(d3) >= -SIG*d0 )
                {
                  double lhs = std::abs(d3);
                  double rhs = -SIG*d0;
                  std::cout << "abs(" << d3 << ")   <   -SIG * " << d0 << " [   i.e. " << lhs << " < " << rhs << " ]\n";
                }
              if ( f3 >= f0+x3*RHO*d0 )
                {
                  double lhs = f3;
                  double rhs = f0+x3*RHO*d0;
                  std::cout <<  f3 << "   <   " << f0 << " + " << x3 << " * RHO * " << d0 << "   [ i.e. " << lhs << " < " << rhs << " ]\n";
                }
            }
          if ( i > length )
            std::cout << "\n[*] Exceeded 'length' value\n";
          */

          // try steepest
          s = -df0;
          d0 = -s.transpose()*s;
          x3 = 1/(1-d0);

          // this line search failed
          ls_failed = true;

        } // END COMPUTE NEW SEARCH DIRECTION

    }  // END MAIN LOOP

  //std::cout << "\nMinimized Function Value (???) :\n";
  //std::cout << f0 << std::endl;

};


// Software License for MTL
//
// Copyright (c) 2007 The Trustees of Indiana University.
//               2008 Dresden University of Technology and the Trustees of Indiana University.
//               2010 SimuNova UG (haftungsbeschränkt), www.simunova.com.
// All rights reserved.
// Authors: Peter Gottschling and Andrew Lumsdaine
//
// This file is part of the Matrix Template Library
//
// See also license.mtl.txt in the distribution.

#ifndef MATH_POWER_INCLUDE
#define MATH_POWER_INCLUDE

#include <concepts>
#include <boost/numeric/linear_algebra/concepts.hpp>
#include <boost/numeric/linear_algebra/identity.hpp>
#include <stdexcept>

namespace math {

    template <typename Op, std::Semiregular Element, Integral Exponent>
        requires std::Callable2<Op, Element, Element>
              && std::Convertible<std::Callable2<Op, Element, Element>::result_type, Element>
    inline Element power(const Element& a, Exponent n, Op op)
    {
	std::cout << "[Magma] ";
	if (n < 1) throw std::range_error("power [magma]: n must be > 0");

	Element value= a;
	for (; n > 1; --n)
	    value= op(value, a);
	return value;
    }

#if 0
    template <typename Op, std::Semiregular Element, Integral Exponent>
        requires SemiGroup<Op, Element>
              && std::Callable2<Op, Element, Element>
              && std::Convertible<std::Callable2<Op, Element, Element>::result_type, Element>
    inline Element multiply_and_square_horner(const Element& a, Exponent n, Op op)
    {
	if (n < 1) throw std::range_error("mult&square Horner: n must be > 0");

        // Set mask to highest bit
        Exponent mask= 1 << (8 * sizeof(mask) - 1);

        // If this is a negative number right shift can insert 1s instead of 0s -> infinite loop
        // Therefore we take the 2nd-highest bit
        if (mask < 0)
	    mask= 1 << (8 * sizeof(mask) - 2);

        // Find highest 1 bit
        while(!bool(mask & n)) mask>>= 1;

        Element value= a;
        for (mask>>= 1; mask; mask>>= 1) {
	    value= op(value, value);
	    if (n & mask)
		value= op(value, a);
        }
        return value;
    }

    template <typename Op, std::Semiregular Element, Integral Exponent>
        requires SemiGroup<Op, Element>
              && std::Callable2<Op, Element, Element>
              && std::Convertible<std::Callable2<Op, Element, Element>::result_type, Element>
    inline Element power(const Element& a, Exponent n, Op op)
    {
	return multiply_and_square_horner(a, n, op);
    }
#endif

#if 1
    // With Horner scheme we can avoid recursion
    // This one is more intuitive (I believe)
    template <typename Op, std::Semiregular Element, Integral Exponent>
        requires SemiGroup<Op, Element>
              && std::Callable2<Op, Element, Element>
              && std::Convertible<std::Callable2<Op, Element, Element>::result_type, Element>
    inline Element power(const Element& a, Exponent n, Op op)
    {
	std::cout << "[SemiGroup] ";
	if (n < 1) throw std::range_error("power [SemiGroup]: n must be > 0");

	Exponent half(n / 2);
        // If half is 0 then n must be 1 and the result is a
        if (half == 0)
	    return a;

        // Compute power of downward rounded exponent and "square" the result
        Element value= power(a, half, op);
        value= op(value, value);

        // If n is odd another operation with a is needed
        if (n & 1)
	    value= op(value, a);
        return value;
    }
#endif

    template <typename Op, std::Semiregular Element, Integral Exponent>
        requires Monoid<Op, Element>
              && std::Callable2<Op, Element, Element>
              && std::Convertible<std::Callable2<Op, Element, Element>::result_type, Element>
    inline Element multiply_and_square(const Element& a, Exponent n, Op op)
    {
	// Same as the simpler form except that the first multiplication is made before
	// the loop and one squaring is saved this way
	if (n < 0) throw std::range_error("mult&square: n must be >= 0");

	using math::identity;
	Element value= bool(n & 1) ? Element(a) : Element(identity(op, a)), square= a;

	for (n>>= 1; n > 0; n>>= 1) {
	    square= op(square, square);
	    if (n & 1)
		value= op(value, square);
	}
	return value;
    }

    template <typename Op, std::Semiregular Element, Integral Exponent>
        requires Monoid<Op, Element>
              && std::Callable2<Op, Element, Element>
              && std::Convertible<std::Callable2<Op, Element, Element>::result_type, Element>
    inline Element power(const Element& a, Exponent n, Op op)
    {
	std::cout << "[Monoid] ";
	return multiply_and_square(a, n, op);
    }

    template <typename Op, std::Semiregular Element, Integral Exponent>
        requires PIMonoid<Op, Element>
              && std::Callable2<Op, Element, Element>
              && std::Convertible<std::Callable2<Op, Element, Element>::result_type, Element>
    inline Element power(const Element& a, Exponent n, Op op)
    {
	std::cout << "[PIMonoid] ";
	if (n < 0 && !is_invertible(op, a))
	    throw std::range_error("power [PIMonoid]: a must be invertible with n < 0");

	return n < 0 ? multiply_and_square(Element(inverse(op, a)), Exponent(-n), op)
	             : multiply_and_square(a, n, op);
    }

#if 1
    template <typename Op, std::Semiregular Element, Integral Exponent>
        requires Group<Op, Element>
              && std::Callable2<Op, Element, Element>
              && std::Convertible<std::Callable2<Op, Element, Element>::result_type, Element>
    inline Element power(const Element& a, Exponent n, Op op)
    {
	std::cout << "[Group] ";
	// For groups we don't need any range test

	return n < 0 ? multiply_and_square(Element(inverse(op, a)), Exponent(-n), op)
	             : multiply_and_square(a, n, op);
    }
#endif

#if 0
    template <typename Op, typename Element, typename Exponent>
        requires Group<Op, Element>
              && Integral<Exponent>
              && std::Semiregular<Element>
              && std::Callable2<Op, Element, Element>
              && std::Convertible<std::Callable2<Op, Element, Element>::result_type, Element>
              && std::Semiregular<math::Inversion<Op, Element>::result_type>
              && std::HasNegate<Exponent>
              && math::Monoid<Op, math::Inversion<Op, Element>::result_type>
              && Integral< std::HasNegate<Exponent>::result_type>
              && std::Callable2<Op, math::Inversion<Op, Element>::result_type,
				math::Inversion<Op, Element>::result_type>
              && std::Convertible<std::Callable2<Op, math::Inversion<Op, Element>::result_type,
						 math::Inversion<Op, Element>::result_type>::result_type,
				  math::Inversion<Op, Element>::result_type>
    inline Element power(const Element& a, Exponent n, Op op)
    {
	return n < 0 ? multiply_and_square(inverse(op, a), -n, op)
	             : multiply_and_square(a, n, op);
    }
#endif

} // namespace math

#endif // MATH_POWER_INCLUDE


//
//  ClosedFormSVDSVD2d.hpp
//  IPC
//
//  Created by Minchen Li on 8/31/18.
//  based on https://www.researchgate.net/publication/263580188_Closed_Form_SVD_Solutions_for_2_x_2_Matrices_-_Rev_2
//

#ifndef ClosedFormSVD2d_hpp
#define ClosedFormSVD2d_hpp

#include <Eigen/Eigen>

#include <iostream>

namespace IPC {

template <typename MatrixType>
class AutoFlipSVD : Eigen::JacobiSVD<MatrixType> {
    typedef Eigen::JacobiSVD<MatrixType> Base;

public:
    AutoFlipSVD(void) {}
    AutoFlipSVD(const MatrixType& mtr, unsigned int computationOptions = 0)
        : Base(2, 2, computationOptions)
    {
        if (MatrixType::RowsAtCompileTime == Eigen::Dynamic) {
            assert(mtr.rows() == 2);
        }
        else {
            assert(MatrixType::RowsAtCompileTime == 2);
        }

        if (MatrixType::ColsAtCompileTime == Eigen::Dynamic) {
            assert(mtr.cols() == 2);
        }
        else {
            assert(MatrixType::ColsAtCompileTime == 2);
        }

        Base::m_isInitialized = true;

        compute(mtr);
    }

public:
    AutoFlipSVD& compute(const MatrixType& A)
    {
        bool computeU = (Base::m_computeFullU || Base::m_computeThinU);
        bool computeV = (Base::m_computeFullV || Base::m_computeThinV);

        const double a = A(0, 0);
        const double b = A(1, 0);
        const double c = A(0, 1);
        const double d = A(1, 1);

        const double ad = a * d;
        const double bc = b * c;
        const double _2admbc = 2.0 * (ad - bc);
        const double sqn = A.squaredNorm();

        const double sum = sqn + _2admbc;
        const double dif = sqn - _2admbc;
        if (dif <= 0.0) {
            // avoid dividing by 0 in general formula
            const double aa = (a + d) / 2.0;
            const double bb = (b - c) / 2.0;
            const double lambda = std::sqrt(aa * aa + bb * bb);
            Base::m_singularValues.setConstant(lambda);

            if (computeU) {
                if (lambda == 0.0) {
                    // avoid dividing by 0
                    Base::m_matrixU.setIdentity();
                }
                else {
                    const double cosl = aa / lambda;
                    const double sinl = bb / lambda;
                    Base::m_matrixU << cosl, -sinl, sinl, cosl;
                }
            }

            if (computeV) {
                Base::m_matrixV.setIdentity();
            }
        }
        else if (sum <= 0.0) {
            // avoid dividing by 0 in general formula
            // symmetric matrix with a=-d
            const double aa = (a - d) / 2.0;
            const double bb = (b + c) / 2.0;
            const double lambda = std::sqrt(aa * aa + bb * bb);
            Base::m_singularValues << lambda, -lambda;

            if (computeU || computeV) {
                if (bb == 0.0) {
                    // avoid dividing by 0 and sqrt(<0)
                    if (computeU) {
                        Base::m_matrixU.setIdentity();
                    }

                    if (computeV) {
                        Base::m_matrixV.setIdentity();
                    }
                }
                else {
                    const double a_div_lambda_half = aa / lambda / 2.0;
                    bool neg_b = (bb < 0.0);
                    const double cos2 = 0.5 + a_div_lambda_half;
                    const double cos = ((cos2 <= 0.0) ? 0.0 : std::sqrt(cos2));
                    const double sin2 = 0.5 - a_div_lambda_half;
                    const double sin = ((sin2 <= 0.0) ? 0.0 : (neg_b ? -std::sqrt(sin2) : std::sqrt(sin2)));
                    if (computeU) {
                        Base::m_matrixU << cos, -sin, sin, cos;
                    }

                    if (computeV) {
                        Base::m_matrixV << cos, -sin, sin, cos;
                    }
                }
            }
        }
        else {
            const double sqrt_sum = std::sqrt(sum); // safe
            const double sqrt_dif = std::sqrt(dif); // safe

            Base::m_singularValues[0] = (sqrt_sum + sqrt_dif) / 2.0;
            Base::m_singularValues[1] = ((_2admbc < 0.0) ? (-std::abs(sqrt_sum - sqrt_dif) / 2.0) : (std::abs(sqrt_sum - sqrt_dif) / 2.0));

            if (computeU || computeV) {
                const double a2 = a * a;
                const double b2 = b * b;
                const double c2 = c * c;
                const double d2 = d * d;

                const double denom = sqrt_sum * sqrt_dif * 2.0;

                const double a2md2 = a2 - d2;
                const double b2mc2 = b2 - c2;

                const double ab = a * b;
                const double cd = c * d;
                const bool neg_ab_p_cd = ((ab + cd) < 0);

                if (computeU) {
                    const double a2md2_m_b2mc2_div_ = (a2md2 - b2mc2) / denom; // safe

                    // avoid sqrt(<0)
                    const double cosl2 = 0.5 + a2md2_m_b2mc2_div_;
                    const double cosl = ((cosl2 <= 0.0) ? 0.0 : std::sqrt(cosl2));
                    const double sinl2 = 0.5 - a2md2_m_b2mc2_div_;
                    const double sinl = ((sinl2 <= 0.0) ? 0.0 : (neg_ab_p_cd ? -std::sqrt(sinl2) : std::sqrt(sinl2)));

                    Base::m_matrixU << cosl, -sinl, sinl, cosl;
                }

                if (computeV) {
                    const double ac = a * c;
                    const double bd = b * d;
                    const bool neg_ac_p_bd = ((ac + bd) < 0);
                    const double a2md2_p_b2mc2_div_ = (a2md2 + b2mc2) / denom; // safe

                    // avoid sqrt(<0)
                    const double cosr2 = 0.5 + a2md2_p_b2mc2_div_;
                    const double cosr = ((cosr2 <= 0.0) ? 0.0 : std::sqrt(cosr2));
                    const double sinr2 = 0.5 - a2md2_p_b2mc2_div_;
                    const double sinr = ((sinr2 <= 0.0) ? 0.0 : (neg_ac_p_bd ? -std::sqrt(sinr2) : std::sqrt(sinr2)));

                    const bool s = neg_ab_p_cd ^ neg_ac_p_bd;
                    const bool neg_apsd = ((a + (s ? -d : d)) < 0.0);
                    if (neg_apsd) {
                        Base::m_matrixV << -cosr, sinr, -sinr, -cosr;
                    }
                    else {
                        Base::m_matrixV << cosr, -sinr, sinr, cosr;
                    }
                }
            }
        }

        return *this;
    }
    AutoFlipSVD& compute(const MatrixType& mtr, unsigned int computationOptions)
    {
        if (MatrixType::RowsAtCompileTime == Eigen::Dynamic) {
            assert(mtr.rows() == 2);
        }
        else {
            assert(MatrixType::RowsAtCompileTime == 2);
        }

        if (MatrixType::ColsAtCompileTime == Eigen::Dynamic) {
            assert(mtr.cols() == 2);
        }
        else {
            assert(MatrixType::ColsAtCompileTime == 2);
        }

        allocate(computationOptions);
        Base::m_isInitialized = true;

        compute(mtr);

        return *this;
    }

protected:
    void allocate(unsigned int computationOptions)
    {
        if (Base::m_isAllocated && 2 == Base::m_rows && 2 == Base::m_cols && computationOptions == Base::m_computationOptions) {
            return;
        }

        Base::m_rows = 2;
        Base::m_cols = 2;
        Base::m_isInitialized = false;
        Base::m_isAllocated = true;
        Base::m_computationOptions = computationOptions;
        Base::m_computeFullU = (computationOptions & Eigen::ComputeFullU) != 0;
        Base::m_computeThinU = (computationOptions & Eigen::ComputeThinU) != 0;
        Base::m_computeFullV = (computationOptions & Eigen::ComputeFullV) != 0;
        Base::m_computeThinV = (computationOptions & Eigen::ComputeThinV) != 0;
        eigen_assert(!(Base::m_computeFullU && Base::m_computeThinU) && "JacobiSVD: you can't ask for both full and thin U");
        eigen_assert(!(Base::m_computeFullV && Base::m_computeThinV) && "JacobiSVD: you can't ask for both full and thin V");
        eigen_assert(EIGEN_IMPLIES(Base::m_computeThinU || Base::m_computeThinV,
                         MatrixType::ColsAtCompileTime == Eigen::Dynamic)
            && "JacobiSVD: thin U and V are only available when your matrix has a dynamic number of columns.");

        Base::m_diagSize = 2;
        Base::m_singularValues.resize(2);
        Base::m_matrixU.resize(2, 2);
        Base::m_matrixV.resize(2, 2);
        Base::m_workMatrix.resize(2, 2);
    }

public:
    const typename Eigen::JacobiSVD<MatrixType>::SingularValuesType& singularValues(void) const
    {
        return Base::m_singularValues;
    }
    const MatrixType& matrixU(void) const
    {
        return Base::m_matrixU;
    }
    const MatrixType& matrixV(void) const
    {
        return Base::m_matrixV;
    }
};

} // namespace IPC

#endif /* ClosedFormSVD2d_hpp */


#include "stdafx.h"
#include "KinectDistortionModel.h"
#include <Eigen\Dense>

using namespace DirectX;
using namespace Eigen;

namespace k4u
{
  float ApplyMediaFoundationRadialDistortion(const DirectX::XMFLOAT3& coeffs, float r)
  {
    return 1 +
      coeffs.x * pow(r, 2.f) +
      coeffs.y * pow(r, 4.f) +
      coeffs.z * pow(r, 6.f);
  }

  float ApplyKinectRadialDistortion(const DirectX::XMFLOAT3& coeffsA, const DirectX::XMFLOAT3& coeffsB, float r)
  {
    auto a = ApplyMediaFoundationRadialDistortion(coeffsA, r);
    auto b = ApplyMediaFoundationRadialDistortion(coeffsB, r);
    auto bi = b == 0.f ? 1.f : 1.f / b;
    return a * bi;
  }

  DirectX::XMFLOAT3 ConvertKinectRadialDistortionToMediaFoundation(const k4a_calibration_camera_t& calibration)
  {
    const int sampleSize = 100;
    Matrix<float, sampleSize, 1> y1;
    Matrix<float, sampleSize, 3> x;

    auto& params = calibration.intrinsics.parameters.param;
    XMFLOAT3 coeffsA{ params.k1, params.k2, params.k3 };
    XMFLOAT3 coeffsB{ params.k4, params.k5, params.k6 };

    auto rStep = calibration.metric_radius / (sampleSize - 1);
    auto rCurrent = 0.f;
    for (auto i = 0; i < sampleSize; i++)
    {
      auto r = rCurrent;
      y1(i) = ApplyKinectRadialDistortion(coeffsA, coeffsB, r) - 1;

      auto r2 = r * r;
      auto r4 = r2 * r2;
      auto r6 = r4 * r2;

      x(i, 0) = r2;
      x(i, 1) = r4;
      x(i, 2) = r6;
      rCurrent += rStep;
    }

    auto k = x.fullPivHouseholderQr().solve(y1);
    return { k(0), k(1), k(2) };
  }
}


//==================================================================================================
/*!
  @file

  @copyright 2016 NumScale SAS
  @copyright 2016 J.T. Lapreste

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_FUNCTION_TANH_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_TANH_HPP_INCLUDED

#if defined(DOXYGEN_ONLY)
namespace boost { namespace simd
{

 /*!

    @ingroup group-hyperbolic
    Function object implementing tanh capabilities

    Returns the hyperbolic tangent: \f$\frac{\sinh(x)}{\cosh(x)}\f$.

    @par Semantic:

    For every parameter of floating type T

    @code
    T r = tanh(x);
    @endcode

    is similar to:

    @code
    T r = sinh(x)/cosh(x);
    @endcode

  **/
  const boost::dispatch::functor<tag::tanh_> tanh = {};
} }
#endif

#include <boost/simd/function/scalar/tanh.hpp>
#include <boost/simd/function/simd/tanh.hpp>

#endif


#include "utils.h"
#include "soft_clustering.h"
#include <boost/random/discrete_distribution.hpp>
#include <boost/random/mersenne_twister.hpp>

using namespace std;

Hyperparam::Hyperparam(AdaCluster& mdl_)
    : copt::BoundedProblem<double>(mdl_.get_n_dims()),
      mdl(mdl_),
      k(mdl_.get_k()),
      n_samples(mdl_.get_n_samples()),
      n_dims(mdl_.get_n_dims()),
      kappa(mdl_.get_kappa()),
      log_pi(mdl_.get_log_pi()),
      data(mdl_.get_data()),
      log_asg(mdl_.get_log_asg()),
      mu(mdl_.get_mu()) {
  buffer.resize(k);
}

double Hyperparam::value(const copt::Vector<double>& x) {
  // x = alpha
  double obj = 0.0;
  for (size_t i = 0; i < n_samples; ++i) {
    for (size_t c = 0; c < k; ++c) {
      buffer[c] = log_pi[c];
      for (unsigned long int j = 0; j < n_dims; ++j){
        if (!isnan(data[i][j])){
          double dist = mdl.distance(data[i][j], mu[c][j], x[j], j) / kappa[j];
          double var = mdl.variance(data[i][j], x[j], j);
          double upd = dist + 0.5 * (log(2 * M_PI) + log(kappa[j]) + log(var));
          // if (isnan(buffer[c])){
          //   cout << "buffer nan is detected ";
          //   cout << m_lowerBound[j] << " " << m_upperBound[j] << " ";
          //   cout << data[i][j] << " " << mu[c][j] << " " << x[j] << " ";
          //   cout << mdl.distance(data[i][j], mu[c][j], x[j], j) / kappa[j] << " ";
          //   cout << mdl.variance(data[i][j], x[j], j) << endl;
          //   break;
          // }
          if (!isnan(upd))
            buffer[c] -= upd;
        }
      }
    }
    obj += logsumexp(buffer);
  }
  return -obj;
}

void Hyperparam::gradient(const copt::Vector<double>& x,
                          copt::Vector<double>& grad) {
  bool terminate = 0;
  double gr = 0.0;
  for (unsigned long int j = 0; j < n_dims; ++j) {
    terminate = 0;
    gr = 0.0;
    for (size_t i = 0; i < n_samples; ++i){
      for (size_t c = 0; c < k; ++c)
        if (!isnan(data[i][j])){
          double diff_dist = mdl.diff_distance(data[i][j], mu[c][j], x[j], j) / kappa[j];
          double diff_var = mdl.diff_variance(data[i][j], x[j], j);
          double var = mdl.variance(data[i][j], x[j], j);
          double pr_asg = exp(log_asg[i][c]);
          double upd = 0.0;
          if (diff_var == 0.0)
            upd = pr_asg * diff_dist;
          else
            upd = pr_asg * (diff_dist + 0.5 * diff_var / var);
          gr += upd;
          // if (isnan(gr)){
          //   cout << "grad nan is detected";
          //   cout << m_lowerBound[j] << " " << m_upperBound[j] << " ";
          //   cout << data[i][j] << " " << mu[c][j] << " " << x[j] << " ";
          //   cout << mdl.diff_distance(data[i][j], mu[c][j], x[j], j) / kappa[j] << " ";
          //   cout << mdl.diff_variance(data[i][j], x[j], j) << " ";
          //   cout << mdl.variance(data[i][j], x[j], j) << endl;
          //   break;
          // }
          if (isnan(upd)){
            terminate = 1;
            break;
          }
        }
      if (terminate)
        break;
    }
    if (terminate)
      grad[j] = 0;
    else
      grad[j] = gr;
  }
}

AdaCluster::AdaCluster(const vector<vector<double>>& data_,
                       const vector<unsigned long int>& label_,
                       unsigned long int max_round_, unsigned long int k_)
    : data(data_), label(label_) {
  max_round = max_round_;
  k = k_;
  nmis.resize(max_round);
  fill(nmis.begin(), nmis.end(), 0);
  logliks.resize(max_round);
  fill(logliks.begin(), logliks.end(), -numeric_limits<double>::max());
  n_samples = data.size();
  n_dims = data[0].size();
  log_pi.resize(k);
  fill(log_pi.begin(), log_pi.end(), -log(k));
  log_asg_sum.resize(k);
  max_log_pi.resize(k);
  kappa.resize(n_dims);
  kappa_a.resize(n_dims);
  kappa_b.resize(n_dims);
  mu.resize(k);
  mu_a.resize(k);
  mu_b.resize(k);
  for (size_t c = 0; c < k; ++c){
    mu[c].resize(n_dims);
    mu_a[c].resize(n_dims);
    mu_b[c].resize(n_dims);
  }
  log_asg.resize(n_samples);
  asg.resize(n_samples);
  for (size_t i = 0; i < n_samples; ++i) log_asg[i].resize(k);
  attr_discrete.resize(n_dims);
  attr_positive.resize(n_dims);
  attr_nonnegative.resize(n_dims);

  alpha.resize(n_dims);
  lb.resize(n_dims);
  ub.resize(n_dims);
  for (unsigned long int j = 0; j < n_dims; ++j) {
    attr_discrete[j] = is_discrete(data, j);
    attr_positive[j] = is_positive(data, j);
    attr_nonnegative[j] = attr_positive[j];
    if (!attr_positive[j]) attr_nonnegative[j] = is_nonnegative(data, j);
  }
  for (unsigned long int j = 0; j < n_dims; ++j) {
    if (attr_discrete[j]) {
      if (attr_nonnegative[j]) {
        alpha[j] = 1;
        lb[j] = 0;
        // ub[j] = numeric_limits<double>::max();
        ub[j] = 10;
      } else {
        alpha[j] = 1;
        lb[j] = 0;
        // ub[j] = numeric_limits<double>::max();
        ub[j] = 10;
      }
    } else {
      if (attr_positive[j]) {
        alpha[j] = 0;
        // lb[j] = -numeric_limits<double>::max();
        lb[j] = -10;
        ub[j] = 2;
      } else if (attr_nonnegative[j]) {
        alpha[j] = 0.5;
        lb[j] = 0;
        ub[j] = 1;
      } else {
        alpha[j] = 1;
        lb[j] = 0;
        // ub[j] = numeric_limits<double>::max();
        ub[j] = 10;
      }
    }
  }
}

double AdaCluster::distance(const double x,
                            const double y,
                            const double alpha,
                            const unsigned long int dim) {
  // return nnc_distance(x, y, alpha);
  if (attr_discrete[dim]){
    if (attr_nonnegative[dim])
      return nnd_distance(x, y, alpha);
    else
      return rc_distance(x, y, alpha);
  } else{
    if (attr_nonnegative[dim])
      return nnc_distance(x, y, alpha);
    else
      return rc_distance(x, y, alpha);
  }
}

double AdaCluster::variance(const double x,
                            const double alpha,
                            const unsigned long int dim){
  // return nnc_variance(x, alpha);
  if (attr_discrete[dim]){
    if (attr_nonnegative[dim])
      return nnd_variance(x, alpha);
    else
      return rc_variance(x, alpha);
  } else {
    if (attr_nonnegative[dim])
      return nnc_variance(x, alpha);
    else
      return rc_variance(x, alpha);
  }
}

double AdaCluster::diff_distance(const double x,
                                 const double y,
                                 const double alpha,
                                 const unsigned long int dim) {
  // return nnc_diff_distance(x, y, alpha);
  if (attr_discrete[dim]){
    if (attr_nonnegative[dim])
      return nnd_diff_distance(x, y, alpha);
    else
      return rc_diff_distance(x, y, alpha);
  } else {
    if (attr_nonnegative[dim])
      return nnc_diff_distance(x, y, alpha);
    else
      return rc_diff_distance(x, y, alpha);
  }
}

double AdaCluster::diff_variance(const double x,
                                 const double alpha,
                                 const unsigned long int dim) {
  // return nnc_diff_variance(x, alpha);
  if (attr_discrete[dim]){
    if (attr_nonnegative[dim])
      return nnd_diff_variance(x, alpha);
    else
      return rc_diff_variance(x, alpha);
  } else {
    if (attr_nonnegative[dim])
      return nnc_diff_variance(x, alpha);
    else
      return rc_diff_variance(x, alpha);
  }
}

unsigned long int AdaCluster::get_k() { return k; }

unsigned long int AdaCluster::get_max_round() { return max_round; }

unsigned long int AdaCluster::get_n_dims() { return n_dims; }

unsigned long int AdaCluster::get_n_samples() { return n_samples; }

vector<double>& AdaCluster::get_kappa() { return kappa; }

vector<double>& AdaCluster::get_log_pi() { return log_pi; }

vector<double>& AdaCluster::get_logliks() { return logliks; }

vector<double>& AdaCluster::get_nmis() { return nmis; }

copt::Vector<double>& AdaCluster::get_alpha() { return alpha; }

copt::Vector<double>& AdaCluster::get_lb() { return lb; }

copt::Vector<double>& AdaCluster::get_ub() { return ub; }

const vector<vector<double>>& AdaCluster::get_data() { return data; }

vector<vector<double>>& AdaCluster::get_log_asg() { return log_asg; }

vector<vector<double>>& AdaCluster::get_mu() { return mu; }

void AdaCluster::initialize_random() {
  unsigned long int ri;
  for (size_t c = 0; c < k; ++c) {
    ri = rand() % n_samples;
    for (unsigned long int j = 0; j < n_dims; ++j){
      mu[c][j] = data[ri][j];
      mu_a[c][j] = mu[c][j];
      mu_b[c][j] = 1.0;
    }
  }
  for (unsigned long int j = 0; j < n_dims; ++j){
    kappa[j] = 1.0;
    kappa_a[j] = 1.0;
    kappa_b[j] = 1e-9;
  }
}

void AdaCluster::initialize_k_plus_plus() {
  boost::mt19937 gen;
  unsigned long int ri, k_eff;
  ri = rand() % n_samples;
  for (unsigned long int j = 0; j < n_dims; ++j) mu[0][j] = data[ri][j];
  k_eff = 1;
  vector<double> probs(n_samples, 0);
  double dist, best_dist;
  do {
    for (size_t i = 0; i < n_samples; ++i) {
      best_dist = numeric_limits<double>::max();
      for (size_t c = 0; c < k_eff; ++c) {
        dist = 0;
        for (unsigned long int j = 0; j < n_dims; ++j){
          // dist += distance(data[i][j], mu[c][j], alpha[j], j);
          dist += pow(data[i][j] - mu[c][j], 2);
        }
        if (dist < best_dist) best_dist = dist;
      }
      probs[i] = pow(best_dist, 2);
    }
    boost::random::discrete_distribution<> dist(probs.begin(), probs.end());
    ri = dist(gen);
    for (unsigned long int j = 0; j < n_dims; ++j) mu[k_eff][j] = data[ri][j];
    k_eff += 1;
  } while (k_eff != k);

  for (unsigned long int j = 0; j < n_dims; ++j)
    if (attr_nonnegative[j])
      for (size_t c = 0; c < k_eff; ++c)
        if (mu[c][j] == 0)
          mu[c][j] = 1e-9;
  for (unsigned long int j = 0; j < n_dims; ++j){
    kappa[j] = 1.0;
    kappa_a[j] = 1.0;
    kappa_b[j] = 1.0;
  }
  for (size_t c = 0; c < k_eff; ++c)
    for (unsigned long int j = 0; j < n_dims; ++j){
      mu_a[c][j] = mu[c][j];
      mu_b[c][j] = 1.0;
    }
}

void AdaCluster::fit() {
  bool updated = 0;
  double loglik = 0;
  double nmi = 0;
  double dist = 0;
  double best_dist = 0;
  unsigned long int best_asg = -1;
  Hyperparam prblm(*this);
  prblm.setLowerBound(lb);
  prblm.setUpperBound(ub);
  for (size_t r = 0; r < max_round; r++) {
    updated = 0;
    loglik = 0.0;
    fill(max_log_pi.begin(), max_log_pi.end(), -numeric_limits<double>::max());
    for (size_t i = 0; i < n_samples; ++i) {
      for (size_t c = 0; c < k; ++c) {
        log_asg[i][c] = log_pi[c];
        for (unsigned long int j = 0; j < n_dims; ++j)
          if (!isnan(data[i][j])){
            double dist = distance(data[i][j], mu[c][j], alpha[j], j) / kappa[j];
            double var = variance(data[i][j], alpha[j], j);
            log_asg[i][c] -= dist + 0.5 * (log(2 * M_PI) + log(kappa[j]) + log(var));
          }
      }
      double norm = logsumexp(log_asg[i]);
      for (size_t c = 0; c < k; ++c) log_asg[i][c] -= norm;
      loglik += norm;
      for (size_t c = 0; c < k; ++c)
        if (log_asg[i][c] > max_log_pi[c]) max_log_pi[c] = log_asg[i][c];
    }
    for (size_t c = 0; c < k; ++c) {
      log_asg_sum[c] = 0.0;
      for (size_t i = 0; i < n_samples; ++i)
        log_asg_sum[c] += exp(log_asg[i][c] - max_log_pi[c]);
      log_asg_sum[c] = log(log_asg_sum[c]) + max_log_pi[c];
      log_pi[c] = log_asg_sum[c] - log(n_samples);
    }

    for (size_t c = 0; c < k; c++)
      for (unsigned long int j = 0; j < n_dims; ++j) {
        double tmp = 0.0;
        for (size_t i = 0; i < n_samples; ++i)
          if (!isnan(data[i][j]))
            tmp += exp(log_asg[i][c])*data[i][j];
        mu[c][j] = (kappa[j]*mu_a[c][j]*mu_b[c][j] + tmp
          )/(kappa[j]*mu_b[c][j] + exp(log_asg_sum[c]));
        // vector<double> pos(n_samples, 0);
        // vector<double> neg(n_samples, 0);
        // double zero_count = 0;
        // for (size_t i = 0; i < n_samples; ++i){
        //   if (data[i][j] > 0)
        //     pos[i] = log_asg[i][c] + log(data[i][j]);
        //   else if (data[i][j] < 0)
        //     neg[i] = log_asg[i][c] + log(-data[i][j]);
        //   else
        //     zero_count += 1;
        // }
        // mu[c][j] = (kappa[j]*mu_a[c][j] + exp(logsumexp(pos)) - exp(logsumexp(neg))
        //             )/(kappa[j]*mu_b[c][j] + exp(log_asg_sum[c]));
      }

    double pr_asg = 0.0;
    double nom = 0.0;
    double denom = 0.0;
    for (unsigned long int j = 0; j < n_dims; ++j) {
      nom = 0.0;
      denom = 0.0;
      for (size_t i = 0; i < n_samples; ++i)
        for (size_t c = 0; c < k; c++) {
          pr_asg = exp(log_asg[i][c]);
          if (!isnan(data[i][j])) {
            nom += pr_asg * distance(data[i][j], mu[c][j], alpha[j], j);
            denom += pr_asg;
          }
        }
      kappa[j] =  (kappa_b[j] + nom) / (kappa_a[j] + 0.5*denom);
    }

    copt::Vector<double> alpha_prev(alpha);
    // cout << alpha.transpose() << endl;
    copt::LbfgsbSolver<Hyperparam> solver;
    solver.minimize(prblm, alpha);
    // cout << alpha.transpose() << endl;

    for (unsigned long int j = 0; j < n_dims; ++j)
      if (isnan(alpha[j]))
        alpha[j] = alpha_prev[j];

    for (size_t i = 0; i < n_samples; ++i) {
      best_dist = log_asg[i][0];
      best_asg = 0;
      for (size_t c = 1; c < k; ++c)
        if (log_asg[i][c] > best_dist) {
          best_dist = log_asg[i][c];
          best_asg = c;
        }
      if (asg[i] != best_asg) {
        asg[i] = best_asg;
        updated = 1;
      }
    }
    nmi = calc_nmi(label, asg);
    nmis[r] = nmi;
    cout << "round=" << r << " nmi=" << nmi << " loglik=" << loglik << endl;
    if (!updated) break;
  }
  cout << "Alpha=";
  for (unsigned long int j = 0; j < n_dims; ++j)
      cout << alpha[j] << " ";
  cout << endl;
  cout << "Kappa=";
  for (unsigned long int j = 0; j < n_dims; ++j)
      cout << kappa[j] << " ";
  cout << endl;
  for (size_t c = 0; c < k; ++c){
    cout << "mu[" << c << "]=";
    for (unsigned long int j = 0; j < n_dims; ++j)
        cout << mu[c][j] << " ";
    cout << endl;
  }
}

// GMM::GMM(const vector<vector<double>>& data_,
//          const vector<unsigned long int>& label_, unsigned long int max_round_,
//          unsigned long int k_)
//     : AdaCluster(data_, label_, max_round_, k_) {}

// double GMM::log_base_measure(unsigned long int sample, unsigned long int dim) {
//   return -0.5 * log(M_PI) - 0.5 * log(kappa[dim]);
// }

// double GMM::distance(const double x, double y, unsigned long int dim) {
//   return pow(x - y, 2) / 2.0;
// }

// void GMM::update_hyperparams() {}

// BSC::BSC(const vector<vector<double>>& data_,
//          const vector<unsigned long int>& label_, unsigned long int max_round_,
//          unsigned long int k_)
//     : AdaCluster(data_, label_, max_round_, k_) {
//   beta.resize(n_dims);
//   for (unsigned long int j = 0; j < n_dims; ++j) beta[j] = 2.0;
//   sum_log_data.resize(n_dims);
//   fill(sum_log_data.begin(), sum_log_data.end(), 0);
//   for (unsigned long int j = 0; j < n_dims; ++j)
//     for (size_t i = 0; i < n_samples; ++i) sum_log_data[j] += log_data[i][j];
// }

// vector<double>& BSC::get_beta() { return beta; }

// double BSC::log_base_measure(unsigned long int sample, unsigned long int dim) {
//   return -0.5 * log(M_PI) - 0.5 * log(kappa[dim]) -
//          0.5 * (2.0 - beta[dim]) * log_data[sample][dim];
// }

// double BSC::distance(const double x, double y, unsigned long int dim) {
//   return beta_div(x, y, beta[dim]);
// }

// void BSC::update_hyperparams() {
//   vector<double> x(n_dims);
//   vector<double> lb(n_dims);
//   vector<double> ub(n_dims);
//   for (unsigned long int j = 0; j < n_dims; ++j) {
//     lb[j] = -numeric_limits<double>::max();
//     ub[j] = numeric_limits<double>::max();
//     // lb[j] = -5.0;
//     // ub[j] = 5.0;
//   }
//   Saddle saddle(data, log_data, log_asg, log_pi, mu, kappa);
//   saddle.setLowerBound(lb);
//   saddle.setUpperBound(ub);
//   copt::LbfgsbSolver<double> solver;
//   int max_trial = 1;
//   double estimate = 0.0;
//   bool success = 1;
//   for (int trial = 0; trial < max_trial; trial++) {
//     for (unsigned long int j = 0; j < n_dims; ++j) x[j] = beta[j];
//     solver.minimize(saddle, x);
//     success = 1;
//     for (unsigned long int j = 0; j < n_dims; ++j)
//       if (std::isnan(x[j]) || !std::isfinite(x[j])) {
//         success = 0;
//         break;
//       }
//     if (success) break;
//   }
//   for (unsigned long int j = 0; j < n_dims; ++j) {
//     beta[j] = x[j];
//     // cout << " beta[" << j << "]=" << beta[j];
//   }
//   // cout << endl;
// }


#include "Nodal.h"
#include "Functions.h"
#include "quadrules/GaussJacobi.h"
#include "quadrules/IntervalQuadratureRule.h"
#include "util/Combinatorics.h"

#include <Eigen/LU>

#include <algorithm>
#include <cassert>
#include <iterator>

using Eigen::MatrixXd;
using Eigen::VectorXd;

namespace tndm {

std::vector<double> LegendreGaussLobattoPoints(unsigned n, unsigned a, unsigned b) {
    assert(n >= 2);

    auto gjPoints = GaussJacobi(n - 2, a + 1, b + 1).points();
    std::vector<double> glPoints;
    glPoints.reserve(n);
    glPoints.push_back(-1.0);
    std::copy(gjPoints.rbegin(), gjPoints.rend(), std::back_inserter(glPoints));
    glPoints.push_back(1.0);
    return glPoints;
}

template <std::size_t D>
MatrixXd Vandermonde(unsigned degree, std::vector<std::array<double, D>> const& points) {
    assert(binom(degree + D, D) == points.size());

    MatrixXd vandermonde(points.size(), binom(degree + D, D));

    for (std::size_t i = 0; i < points.size(); ++i) {
        std::size_t bf = 0;
        for (auto j : AllIntegerSums<D>(degree)) {
            vandermonde(i, bf++) = DubinerP(j, points[i]);
        }
    }

    return vandermonde;
}

template <std::size_t D>
LebesgueFunction<D>::LebesgueFunction(unsigned degree,
                                      std::vector<std::array<double, D>> const& nodes)
    : degree(degree), phi(nodes.size()), L(nodes.size()) {
    auto vandermonde = Vandermonde(degree, nodes);
    vInvT = vandermonde.transpose().inverse();
}

template <std::size_t D> double LebesgueFunction<D>::operator()(std::array<double, D> const& xi) {
    std::size_t bf = 0;
    for (auto j : AllIntegerSums<D>(degree)) {
        phi(bf++) = DubinerP(j, xi);
    }
    assert(bf == vInvT.cols());
    L = vInvT * phi;
    return L.lpNorm<1>();
}

template MatrixXd Vandermonde<1u>(unsigned, std::vector<std::array<double, 1u>> const&);
template MatrixXd Vandermonde<2u>(unsigned, std::vector<std::array<double, 2u>> const&);
template MatrixXd Vandermonde<3u>(unsigned, std::vector<std::array<double, 3u>> const&);

template class LebesgueFunction<1u>;
template class LebesgueFunction<2u>;
template class LebesgueFunction<3u>;

} // namespace tndm