/*
    GNU Octave level-set package.
    Copyright (C) 2013-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/>.
*/

#include "FastMarching.hpp"

#include <algorithm>
#include <cmath>
#include <memory>
#include <stdexcept>

namespace fastMarching
{

/* Tolerance for safety checks.  */
static const realT ASSERT_TOL = 1.0e-9;

/* ************************************************************************** */
/* Update equation class.  */

/**
 * Destruct and free all memory held in the entries map.
 */
UpdateEquation::~UpdateEquation ()
{
  for (auto e : entries)
    delete e.second;
}

/**
 * Add a new entry together with g.
 * @param d Dimension along which the point neighbours the centre.
 * @param h Grid spacing to the neighbour.
 * @param u Value of the distance function at the neighbour.
 * @param g Value of the extended function at the neighbour.
 */
void
UpdateEquation::add (dimensionT d, realT h, realT u, realT g)
{
  std::unique_ptr<EqnEntry> newEntry(new EqnEntry (h, u, g));

  auto iter = entries.find (d);
  if (iter == entries.end ())
    {
      entries.insert (std::make_pair (d, newEntry.release ()));
      return;
    }

  if (iter->second->getEffDistance () > newEntry->getEffDistance ())
    {
      delete iter->second;
      iter->second = newEntry.release ();
    }
}

/**
 * Calculate the new distance value solving the quadratic equation.
 * @param f Speed value at the current cell.
 * @return Distance value at the centre cell.
 */
realT
UpdateEquation::solve (realT f) const
{
  sortedEntries.clear ();

  /* First, sort the entries by distance.  */
  for (const auto& e : entries)
    sortedEntries.push_back (e.second);
  std::sort (sortedEntries.begin (), sortedEntries.end (), &EqnEntry::compare);

  /*
  The equation we solve is:

    (X - N_1)^2 / h_1^2 + (X - N_2)^2 / h_2^2 + ... = f^2,

  where N_i are the smaller, alive neighbours in each direction.  During
  the loop, we build up the coefficients accordingly.
  */
      
  realT a, b, c;
  realT d;

  while (true)
    {
      assert (!sortedEntries.empty ());

      a = 0.0;
      b = 0.0;
      c = -std::pow (f, 2);

      for (const auto cur : sortedEntries)
        {
          const realT invH2 = std::pow (cur->h, -2);
          a += invH2;
          b -= 2.0 * cur->u * invH2;
          c += std::pow (cur->u, 2) * invH2;
        }

      d = std::pow (b, 2) - 4.0 * a * c;
      if (d >= 0.0)
        break;

      issueWarning ("level-set:fast-marching:too-far-alive",
                    "found too far alive point in fast marching, d = %f", d);
      sortedEntries.pop_back ();
    }

  /* Now the equation is solvable, actually do it and take the larger
     of the solutions.  */
  u = (-b + sqrt (d)) / (2.0 * a);

  return u;
}

/**
 * Calculate the value of the extended function at the centre.  This
 * can only be used if all g values were given for all entries, and
 * if solve() was already called.
 * @return The value of the extended function at the centre.
 * @throws std::logic_error if the equation was built without all g's.
 * @throws std::logic_error if solve() was not yet called.
 */
realT
UpdateEquation::getExtendedFunction () const
{
  if (gMissing)
    throw std::logic_error ("not all values of g given in update equation");
  if (sortedEntries.empty ())
    throw std::logic_error ("update equation is not yet solved");

  /*
  For the extended function g, the equation to solve is:

    (G - G_1) (U - N_1) / h_1^2 + (G - G_2) (U - N_2) / h_2^2 + ... = 0,

  which is even linear in G since we know U already.  We solve this equation
  in the form "a G = b".
  */

  realT a = 0.0;
  realT b = 0.0;
  for (const auto cur : sortedEntries)
    {
      const realT invH2 = std::pow (cur->h, -2);
      const realT grad = u - cur->u;
      a += grad * invH2;
      b += cur->g * grad * invH2;
    }

  assert (a > 0.0);
  return b / a;
}

/* ************************************************************************** */
/* Single cell entry.  */

/**
 * Set state to alive.
 * @see setNarrowBand
 */
void
Entry::setAlive ()
{
  /* Note that this routine is only used to set narrow_band -> alive, since
     far_away -> narrow_band must be done with setNarrowBand.  */
  assert (state == NARROW_BAND);
  state = ALIVE;

  assert (heapEntry);
  heapEntry = nullptr;
}

/**
 * Set state to narrow band, adding also the heap entry.
 * @param e Heap entry.
 * @see setState
 */
void
Entry::setNarrowBand (const heapT::Entry& e)
{
  assert (state == FAR_AWAY);
  state = NARROW_BAND;

  assert (!heapEntry);
  heapEntry = e;
}

/**
 * Update distance value from the neighbours.  They are found
 * using our coordinates and the given grid.  This may raise assertion
 * failures under certain conditions.
 * @param grid The grid to use to access neighbours.
 */
void
Entry::update (const Grid& grid)
{
  assert (state == NARROW_BAND);

  /*
  Go over all neighbours and remember in each dimension the smaller of
  the neighbour values.  We only use alive values as suggested in the
  technical report of J. A. Baerentzen.
  */

  UpdateEquation eqn;
  const auto buildEqn
    = [&eqn, &grid] (const IndexTuple& neighbour, dimensionT d, indexT)
    {
      const Entry* e = grid.get (neighbour);
      if (e && e->state == ALIVE && e->u < LARGE_DIST)
        eqn.add (d, 1.0, e->u, e->g);
    };
  grid.iterateNeighbours (coord, buildEqn);

  /* Solve the equation.  */
  const realT newU = eqn.solve (f);
  g = eqn.getExtendedFunction ();

  /* Update.  Usually the distance should be smaller, but in certain
     circumstances we want to accomodate also for increases correcting
     prior inaccuracies (and inaccuracies in the initialisation possibly).  */
  bool bubbleDown = false;
  if (newU > u)
    {
      if (newU > u + ASSERT_TOL)
        issueWarning ("level-set:fast-marching:increased-distance",
                      "increased distance in fast marching %f to %f", u, newU);
      bubbleDown = true;
    }
  u = newU;

  /* Bubble up the heap.  */
  assert (heapEntry);
  if (bubbleDown)
    heapEntry->bubbleDown ();
  else
    heapEntry->bubbleUp ();
}

/* ************************************************************************** */
/* Full grid.  */

/**
 * Calculate index into entries for given coordinates.
 * @param c Coordinates in indices.
 * @return Single index into entries or -1 if c is out of range.
 */
indexT
Grid::lineariseIndex (const IndexTuple& c) const
{
  assert (size.size () == c.size ());

  indexT res = 0;
  for (dimensionT dim = 0; dim < size.size (); ++dim)
    {
      if (c[dim] < 0 || c[dim] >= size[dim])
        return -1;
      res *= size[dim];
      res += c[dim];
    }

  return res;
}

/**
 * Perform a single update step on the currently best next narrow-band
 * entry.  It is removed and set alive, and its neighbours are then
 * added into the narrow-band if not yet there and updated.
 */
void
Grid::update ()
{
  assert (!narrowBand.empty ());

  Entry* top = **narrowBand.top ();
  narrowBand.pop ();

  top->setAlive ();
  recalculateNeighbours (*top);
}

/**
 * Recalculate neighbours of a given entry.  This is used both for
 * update and initialisation of narrow band.
 * @param center Recalculate this entry's neighbours.
 */
void
Grid::recalculateNeighbours (const Entry& center)
{
  IndexTuple neighbour = center.getCoordinate ();
  const auto updateIt
    = [this] (const IndexTuple& neighbour, dimensionT, indexT)
    {
      Entry* e = get (neighbour);
      if (!e)
        return;

      switch (e->getState ())
        {
          case Entry::FAR_AWAY:
            /* Note:  The heap push below inserts first at the very
               bottom of the heap and needs no bubbling up, since the
               distance is currently LARGE_DIST.  It will be bubbled
               only in the update below.  */
            {
              const auto heapEntry = narrowBand.push (e);
              e->setNarrowBand (heapEntry);
            }
            /* Fall through.  */
          case Entry::NARROW_BAND:
            e->update (*this);
            break;

          default:
            /* Nothing to do.  */
            break;
        }
    };
  iterateNeighbours (center.getCoordinate (), updateIt);
}

/**
 * Construct the grid with a given size.  The entries will be empty
 * and will still have to be filled in.
 * @param s Size to use.
 */
Grid::Grid (const IndexTuple& s)
  : size(s), entries(), narrowBand(&Entry::compareForHeap)
{
  indexT fullSize = 1;
  for (auto i : s)
    fullSize *= i;

  entries.resize (fullSize);
  for (auto& p : entries)
    p = nullptr;

  assert (narrowBand.empty ());
}

/**
 * Destructor, freeing all memory.
 */
Grid::~Grid ()
{
  /* Narrow band heap only holds pointers, memory not owned there!  */

  /* Free memory in entries.  */
  for (auto p : entries)
    if (p)
      delete p;
  entries.clear ();
}

/**
 * Initialise a given grid cell by the given entry.  The coordinate
 * where to put the entry is taken from the object itself, and an
 * error is raised if that entry is already filled in.  The entries
 * state must be either "alive" or "far away", narrow band is
 * automatically initialised later on.
 * @param e The new entry.  Ownership is taken over.
 * @throws std::invalid_argument if something is wrong with e.
 */
void
Grid::setInitial (Entry* e)
{
  if (!e)
    throw std::invalid_argument ("entry is NULL");
  if (e->getState () != Entry::ALIVE && e->getState () != Entry::FAR_AWAY)
    throw std::invalid_argument ("only alive and far-away entries can be"
                                 " set initially");

  const IndexTuple& indTuple = e->getCoordinate ();
  const indexT ind = lineariseIndex (indTuple);

  if (ind < 0)
    throw std::invalid_argument ("invalid coordinate of entry");
  if (entries[ind])
    throw std::invalid_argument ("duplicate entry set on grid");

  entries[ind] = e;
}

/**
 * Perform the calculation, updating all narrow-band entries until
 * all of them are alive.  It is necessary to have narrow-band already
 * filled in, meaning that far-away entries not reached by the propagation
 * from any narrow-band will not be updated!
 */
void
Grid::perform ()
{
  for (const auto center : entries)
    if (center && center->getState () == Entry::ALIVE)
      recalculateNeighbours (*center);

  while (!narrowBand.empty ())
    update ();
}

/* ************************************************************************** */

} // Namespace fastMarching.
