/*
    GNU Octave level-set package.
    Copyright (C) 2015  Daniel Kraft <d@domob.eu>

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

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

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

/* Implement an approximation to |grad phi| using an appropriate
   upwind scheme.  The scheme implemented is basically (6.17)
   from Sethian's book "Level Set Methods and Fast Marching Methods",
   Cambridge University Press, second edition, 1999.

   The discretisation could probably be discretised also in Octave code
   with proper vector operations, but this would to harder-to-read code.
   The direct, point-wise C++ formulation is closer to the formulas
   that are derived in the book.  */

#include <octave/oct.h>
#include <octave/lo-ieee.h>

#include <algorithm>
#include <cmath>
#include <cstddef>
#include <stdexcept>

/** Type used to index n-dimensional Octave arrays.  */
typedef Array<octave_idx_type> GridIndex;
/** Type used to denote number of dimensions.  */
typedef size_t dimensionT;

/**
 * Calculate the gradient norm in a given grid point.
 * @param idx The current point's index.
 * @param phi The function whose gradient is requested.
 * @param F The sign gives the upwind direction.
 * @param size The total domain size.
 * @param D The number of dimensions.
 * @return The approximate gradient norm in the point, not yet divided by h.
 */
static double
gradientNormInPoint (const GridIndex& idx, const NDArray& phi, double F,
                     const dim_vector& size, dimensionT D)
{
  double res = 0.0;
  for (dimensionT d = 0; d < D; ++d)
    {
      GridIndex neighbour(idx);
      double forward = 0.0;
      double backward = 0.0;

      if (idx(d) > 0)
        {
          neighbour(d) = idx(d) - 1;
          backward = phi(idx) - phi(neighbour);
        }
      if (idx(d) + 1 < size(d))
        {
          neighbour(d) = idx(d) + 1;
          forward = phi(neighbour) - phi(idx);
        }

      if (F < 0.0)
        std::swap (forward, backward);

      if (backward > 0.0)
        res += std::pow (backward, 2);
      if (forward < 0.0)
        res += std::pow (forward, 2);
    }

  return std::sqrt (res);
}

/**
 * Recursively (since the number of dimensions and thus required loops
 * is not known) iterate the whole array, calling gradientNormInPoint
 * at each grid point.
 * @param idx The current index.
 * @param depth The current depth in the recursion.
 * @param phi The function whose gradient is requested.
 * @param F The sign defines the upwind direction.
 * @param size Size of the array.
 * @param D Number of dimensions.
 * @param grad Return the gradient norm here.
 */
static void
iterate (GridIndex& idx, dimensionT depth, const NDArray& phi, const NDArray& F,
         const dim_vector& size, dimensionT D, NDArray& grad)
{
  if (depth == D)
    {
      grad(idx) = gradientNormInPoint (idx, phi, F(idx), size, D);
      return;
    }

  for (idx(depth) = 0; idx(depth) < size(depth); ++idx(depth))
    iterate (idx, depth + 1, phi, F, size, D, grad);
}

/* Octave routine interfacing to fast marching code.  */
DEFUN_DLD (__levelset_upwindGrad, args, nargout,
  "GNORM = upwindGrad (PHI, F, H)\n\n"
  "Internal routine to calculate the upwind-discretised gradient norm\n"
  "of PHI.  F is the speed field, and only the sign of F is used to decide\n"
  "about the upwind direction.  H is the grid size as usual.\n\n"
  "Returned is an array GNORM of the same size as PHI, with |grad phi| in\n"
  "the entries.\n")
{
  try
    {
      if (args.length () != 3 || nargout != 1)
        throw std::invalid_argument ("invalid argument counts");

      const NDArray phi = args(0).array_value ();
      const NDArray F = args(1).array_value ();
      const double h = args(2).double_value ();

      const dim_vector size = phi.dims ();
      const dimensionT D = size.length ();

      if (size != F.dims ())
        throw std::invalid_argument ("argument dimensions mismatch");

      NDArray grad(size);
      grad.fill (octave_NA);

      GridIndex idx(dim_vector (D, 1));
      iterate (idx, 0, phi, F, size, D, grad);
      grad /= h;

      octave_value_list res;
      res(0) = grad;

      return res;
    }
  catch (const std::exception& exc)
    {
      error (exc.what ());
      return octave_value_list ();
    }
}
