/*

Copyright (C) 2014 David Spies

This file is part of Octave.

Octave 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.

Octave 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 Octave; see the file COPYING.  If not, see
<http://www.gnu.org/licenses/>.

*/
#if !defined (octave_nz_iterators_h)
#define octave_nz_iterators_h 1

#include "interp-idx.h"
#include "oct-inttypes.h"
#include "Array.h"
#include "DiagArray2.h"
#include "PermMatrix.h"
#include "Sparse.h"
#include "direction.h"

// This file contains generic column-major iterators over
// the nonzero elements of any array or matrix.  If you have a matrix mat
// of type M, you can construct the proper iterator type using
// M::iter_type iter(mat) and iter will iterate efficiently (forwards
// or backwards) over the nonzero elements of mat.
//
// The parameter T is the element-type except for PermMatrix where
// the element type is always bool
//
// Use a dir_handler to indicate which direction you intend to iterate.
// (see step-dir.h).  begin() resets the iterator to the beginning or
// end of the matrix (for dir_handler<FORWARD> and <BACKWARD> respectively).
// finished(dirc) indicates whether the iterators has finished traversing
// the nonzero elements.  step(dirc) steps from one element to the next.
//
// You can, for instance, use a for-loop as follows:
//
// typedef M::iter_type iter_t;
// typedef M::element_type T;
//
// iter_t iter(mat);
// dir_handler<1> dirc;
//
// for(iter.begin (dirc); !iter.finished (dirc); iter.step (dirc))
// {
//   octave_idx_type row = iter.row();
//   octave_idx_type col = iter.col();
//   double doub_index = iter.interp_index ();
//   T elem = iter.data();
//   // ... Do something with these
// }
//
// Note that array_iter for indexing over full matrices also includes
// a iter.flat_index () method which returns an octave_idx_type.
//
// The other iterators to not have a flat_index() method because they
// risk overflowing octave_idx_type.  It is recommended you take care
// to implement your function in a way that accounts for this problem.
//
// FIXME: I'd like to add in these
// default no-parameter versions of
// begin() and step() to each of the
// classes.  But the C++ compiler complains
// because apparently I'm not allowed to overload
// templated methods with non-templated ones.  Any
// ideas for work-arounds?
//
//#define INCLUDE_DEFAULT_STEPS \
//  void begin (void) \
//    { \
//      dir_handler<FORWARD> dirc; \
//      begin (dirc); \
//    } \
//  void step (void) \
//    { \
//      dir_handler<FORWARD> dirc; \
//      step (dirc); \
//    } \
//  bool finished (void) const \
//    { \
//      dir_handler<FORWARD> dirc; \
//      return finished (dirc); \
//    }

// A generic method for checking if some element of a matrix with
// element type T is zero.
template<typename T>
bool
is_zero (T t)
{
  return t == static_cast<T> (0);
}

// An iterator over full arrays.  When the number of dimensions exceeds
// 2, calls to iter.col() may exceed mat.cols() up to mat.dims().numel(1)
//
// This mimics the behavior of the "find" method (both in Octave and Matlab)
// on many-dimensional matrices.

template<typename T>
class array_iterator
{
private:
  const Array<T>& mat;

  //Actual total number of columns = mat.dims().numel(1)
  //can be different from length of row dimension
  const octave_idx_type totcols;
  const octave_idx_type numels;

  octave_idx_type coli;
  octave_idx_type rowj;
  octave_idx_type my_idx;

  template<direction dir>
  void
  step_once (dir_handler<dir> dirc)
  {
    my_idx += dir;
    rowj += dir;
    if (dirc.is_ended (rowj, mat.rows ()))
      {
        rowj = dirc.begin (mat.rows ());
        coli += dir;
      }
  }

  template<direction dir>
  void
  move_to_nz (dir_handler<dir> dirc)
  {
    while (!finished (dirc) && is_zero (data ()))
      {
        step_once (dirc);
      }
  }


public:
  array_iterator (const Array<T>& arg_mat)
      : mat (arg_mat), totcols (arg_mat.dims ().numel (1)), numels (
          totcols * arg_mat.rows ())
  {
    dir_handler<FORWARD> dirc;
    begin (dirc);
  }

  template<direction dir>
  void
  begin (dir_handler<dir> dirc)
  {
    coli = dirc.begin (totcols);
    rowj = dirc.begin (mat.rows ());
    my_idx = dirc.begin (mat.numel ());
    move_to_nz (dirc);
  }

  octave_idx_type
  col (void) const
  {
    return coli;
  }
  octave_idx_type
  row (void) const
  {
    return rowj;
  }
  double
  interp_idx (void) const
  {
    return to_interp_idx (my_idx);
  }
  octave_idx_type
  flat_idx (void) const
  {
    return my_idx;
  }
  T
  data (void) const
  {
    return mat.elem (my_idx);
  }

  template<direction dir>
  void
  step (dir_handler<dir> dirc)
  {
    step_once (dirc);
    move_to_nz (dirc);
  }
  template<direction dir>
  bool
  finished (dir_handler<dir> dirc) const
  {
    return dirc.is_ended (my_idx, numels);
  }
};

template<typename T>
class sparse_iterator
{
private:
  const Sparse<T>& mat;
  octave_idx_type coli;
  octave_idx_type my_idx;

  template<direction dir>
  void
  adjust_col (dir_handler<dir> dirc)
  {
    while (!finished (dirc)
        && dirc.is_ended (my_idx, mat.cidx (coli), mat.cidx (coli + 1)))
      coli += dir;
  }

public:
  sparse_iterator (const Sparse<T>& arg_mat) :
      mat (arg_mat)
  {
    dir_handler<FORWARD> dirc;
    begin (dirc);
  }

  template<direction dir>
  void
  begin (dir_handler<dir> dirc)
  {
    coli = dirc.begin (mat.cols ());
    my_idx = dirc.begin (mat.nnz ());
    adjust_col (dirc);
  }

  double
  interp_idx (void) const
  {
    return to_interp_idx (row (), col (), mat.dims ());
  }
  octave_idx_type
  col (void) const
  {
    return coli;
  }
  octave_idx_type
  row (void) const
  {
    return mat.ridx (my_idx);
  }
  T
  data (void) const
  {
    return mat.data (my_idx);
  }
  template<direction dir>
  void
  step (dir_handler<dir> dirc)
  {
    my_idx += dir;
    adjust_col (dirc);
  }
  template<direction dir>
  bool
  finished (dir_handler<dir> dirc) const
  {
    return dirc.is_ended (coli, mat.cols ());
  }
};

template<typename T>
class diag_iterator
{
private:
  const DiagArray2<T>& mat;
  octave_idx_type my_idx;

  template <direction dir>
  void
  move_to_nz (dir_handler<dir> dirc)
  {
    while (!finished (dirc) && is_zero (data ()))
      {
        my_idx += dir;
      }
  }

public:
  diag_iterator (const DiagArray2<T>& arg_mat) :
      mat (arg_mat)
  {
    dir_handler<FORWARD> dirc;
    begin (dirc);
  }

  template<direction dir>
  void
  begin (dir_handler<dir> dirc)
  {
    my_idx = dirc.begin (mat.diag_length ());
    move_to_nz (dirc);
  }

  double
  interp_idx (void) const
  {
    return to_interp_idx (row (), col (), mat.dims ());
  }
  octave_idx_type
  col (void) const
  {
    return my_idx;
  }
  octave_idx_type
  row (void) const
  {
    return my_idx;
  }
  T
  data (void) const
  {
    return mat.dgelem (my_idx);
  }
  template<direction dir>
  void
  step (dir_handler<dir> dirc)
  {
    my_idx += dir;
    move_to_nz (dirc);
  }
  template<direction dir>
  bool
  finished (dir_handler<dir> dirc) const
  {
    return dirc.is_ended (my_idx, mat.diag_length ());
  }
};

class perm_iterator
{
private:
  const PermMatrix& mat;
  octave_idx_type my_idx;

public:
  perm_iterator (const PermMatrix& arg_mat) :
      mat (arg_mat)
  {
    dir_handler<FORWARD> dirc;
    begin (dirc);
  }

  template<direction dir>
  void
  begin (dir_handler<dir> dirc)
  {
    my_idx = dirc.begin (mat.cols ());
  }

  octave_idx_type
  interp_idx (void) const
  {
    return to_interp_idx (row (), col (), mat.dims ());
  }
  octave_idx_type
  col (void) const
  {
    return my_idx;
  }
  octave_idx_type
  row (void) const
  {
    return mat.perm_elem (my_idx);
  }
  bool
  data (void) const
  {
    return true;
  }
  template<direction dir>
  void
  step (dir_handler<dir>)
  {
    my_idx += dir;
  }
  template<direction dir>
  bool
  finished (dir_handler<dir> dirc) const
  {
    return dirc.is_ended (my_idx, mat.rows ());
  }
};

#endif