/*

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_find_h)
#define octave_find_h 1

#include "Array.h"

#include "interp-idx.h"
#include "min-with-nnz.h"
#include "nz-iterators.h"
#include "direction.h"

namespace find
{
  // This struct mostly matches the signature of ffind_result (see find.cc)
  // except without a get_list function.  Internal calls to "find" from
  // liboctave should call the "find" method which in turn constructs one
  // of these as a value to be filled.
  // It should not be used anywhere else

  struct find_result
  {
    Array<octave_idx_type> res;
    template<typename ELT_T>
    void
    add (octave_idx_type place, const array_iterator<ELT_T>& iter)
    {
      res.xelem (place) = iter.flat_idx ();
    }

    find_result (void) { }
    find_result (const dim_vector& nnz) : res (nnz) { }
  };

  // This is the generic "find" method for all matrix types.  All calls to
  // "find" get routed here one way or another as this method contains logic
  // for determining the proper output dimensions and handling forward/reversed,
  // n_to_find etc.
  // It would be unwise to try and duplicate this logic elsewhere as it's fairly
  // nuanced and unintuitive (but Matlab compatible).

  template<typename R, direction dir, typename MT>
  R
  find_to_R (dir_handler<dir> dirc, const MT& v, octave_idx_type n_to_find)
  {
    typedef typename MT::iter_type iter_t;

    octave_idx_type numres;
    if (n_to_find == -1)
      numres = v.nnz ();
    else
      numres = min_with_nnz (v, n_to_find);

    const dim_vector& dv = v.dims ();
    octave_idx_type col_len = dv.elem (0);
    octave_idx_type cols = dv.numel (1);

    // Fixup return dimensions, for Matlab compatibility.
    // find (zeros (0,0)) -> zeros (0,0)
    // find (zeros (1,0)) -> zeros (1,0)
    // find (zeros (0,1)) -> zeros (0,1)
    // find (zeros (0,X)) -> zeros (0,1)
    // find (zeros (1,1)) -> zeros (0,0) !!!! WHY?
    // find (zeros (0,1,0)) -> zeros (0,0)
    // find (zeros (0,1,0,1)) -> zeros (0,0) etc

    dim_vector res_dims;

    if ((cols == 1 && col_len == 1 && numres == 0)
        || (col_len == 0 && cols == 0))
      res_dims = dim_vector (0, 0);
    else if (col_len == 1)
      res_dims = dim_vector (1, numres);
    else
      res_dims = dim_vector (numres, 1);

    iter_t iterator (v);
    R resvec (res_dims);
    octave_idx_type count = dirc.begin (numres);

    for (iterator.begin (dirc);
        !iterator.finished (dirc) && !dirc.is_ended (count, numres);
        iterator.step (dirc), count += dir)
      {
        resvec.add (count, iterator);
      }

    return resvec;
  }

  // The method call used for internal liboctave calls to find.  It constructs
  // an find_result and fills that rather than a res_type_of (see "find.cc").
  // Note that this method only supports full arrays because octave_idx_type
  // will generally overflow if used to represent the nonzero indices of sparse
  // arrays. Instead consider using the sparse-iterator class from
  // nz-iterators.h to avoid this problem

  template<typename ELT_T>
  Array<octave_idx_type>
  find (const Array<ELT_T>& v, octave_idx_type n_to_find = -1, direction dir =
            FORWARD)
  {
    find_result resvec;

    switch (dir)
      {
      case FORWARD:
        dir_handler<FORWARD> dirc_forward;
        resvec = find_to_R<find_result> (dirc_forward, v, n_to_find);
        break;
      case BACKWARD:
        dir_handler<BACKWARD> dirc_back;
        resvec = find_to_R<find_result> (dirc_back, v, n_to_find);
        break;
      default:
        liboctave_fatal(
                "find: unknown direction: %d; must be FORWARD or BACKWARD",
                dir);
      }

    return resvec.res;
  }
}

#endif