// Copyright (c) 2010-2018 The Regents of the University of Michigan
// This file is from the freud project, released under the BSD 3-Clause License.

#ifndef LINKCELL_H
#define LINKCELL_H

#include <memory>
#include <vector>

#include "Box.h"
#include "NeighborList.h"
#include "Index1D.h"

/*! \file LinkCell.h
    \brief Build a cell list from a set of points.
*/

namespace freud { namespace locality {

/*! \internal
    \brief Signifies the end of the linked list
*/
const unsigned int LINK_CELL_TERMINATOR = 0xffffffff;

//! Iterates over particles in a link cell list generated by LinkCell
/*! The link-cell structure is not trivial to iterate over. This helper class
 *  makes that easier both in C++ and provides a Python compatible interface
 *  for direct usage there. An IteratorLinkCell is given the bare essentials
 *  it needs to iterate over a given cell, the cell list, the number of
 *  particles, number of cells and the cell to iterate over. Call next() to
 *  get the index of the next particle in the cell, atEnd() will return true
 *  if you are at the end. In C++, next() will crash the code if you attempt
 *  to iterate past the end (no bounds checking for performance). When called
 *  from Python, a different version of next() is used that will throw
 *  StopIteration at the end.
 *
 *  A loop over all of the particles in a cell can be accomplished with the
 *   following code in C++.
 * \code
 * LinkCell::iteratorcell it = lc.itercell(cell);
 * for (unsigned int i = it.next(); !it.atEnd(); i=it.next())
 *     {
 *     // do something with particle i
 *     }
 * \endcode
 *
 * \note Behavior is undefined if an IteratorLinkCell is accessed after the
 * parent LinkCell is destroyed.
*/
class IteratorLinkCell
    {
    public:
        IteratorLinkCell():
            m_cell_list(NULL), m_Np(0), m_Nc(0),
            m_cur_idx(LINK_CELL_TERMINATOR), m_cell(0)
            {
            }

        IteratorLinkCell(const std::shared_ptr<unsigned int>& cell_list,
                         unsigned int Np,
                         unsigned int Nc,
                         unsigned int cell)
                         : m_cell_list(cell_list.get()), m_Np(Np), m_Nc(Nc)
            {
            assert(cell < Nc);
            assert(Np > 0);
            assert(Nc > 0);
            m_cell = cell;
            m_cur_idx = m_Np + cell;
            }

        //! Copy the position of rhs into this object
        void copy(const IteratorLinkCell &rhs)
        {
            m_cell_list = rhs.m_cell_list;
            m_Np = rhs.m_Np;
            m_Nc = rhs.m_Nc;
            m_cur_idx = rhs.m_cur_idx;
            m_cell = rhs.m_cell;
        }

        //! Test if the iteration over the cell is complete
        bool atEnd()
            {
            return (m_cur_idx == LINK_CELL_TERMINATOR);
            }

        //! Get the next particle index in the list
        unsigned int next()
            {
            m_cur_idx = m_cell_list[m_cur_idx];
            return m_cur_idx;
            }

        //! Get the first particle index in the list
        unsigned int begin()
            {
            m_cur_idx = m_Np + m_cell;
            m_cur_idx = m_cell_list[m_cur_idx];
            return m_cur_idx;
            }

    private:
        const unsigned int *m_cell_list;  //!< The cell list
        unsigned int m_Np;                //!< Number of particles in the cell list
        unsigned int m_Nc;                //!< Number of cells in the cell list
        unsigned int m_cur_idx;           //!< Current index
        unsigned int m_cell;              //!< Cell being considered
    };

//! Iterates over sets of shells in a cell list
/*! This class provides a convenient way to iterate over distinct
    shells in a cell list structure. For a range of N, these are the
    faces, edges, and corners of a cube of edge length 2*N + 1 cells
    large. While IteratorLinkCell provides a way to iterate over
    neighbors given a cell, IteratorCellShell provides a way to find
    which cell offsets should be applied to find all the neighbors of
    a particular reference shell within a distance some number of
    cells away.

\code
// Grab neighbor cell offsets within the 3x3x3 typical search distance
for(IteratorCellShell iter(0); iter != IteratorCellShell(2); ++iter)
{
    // still need to apply modulo operation for dimensions of the cell list
    const vec3<int> offset(*iter);
}
\endcode
 */
class IteratorCellShell
    {
    public:
        IteratorCellShell(unsigned int range=0, bool is2D=false):
            m_is2D(is2D)
            {
            reset(range);
            }

        void operator++()
            {
            // this bool indicates that we have wrapped over in whichever
            // direction we are looking and should move to the next
            // row/plane
            bool wrapped(false);

            switch(m_stage)
                {
                // +y wedge: iterate over x and (possibly) z
                // zs = list(range(-N + 1, N)) if threeD else [0]
                // for r in itertools.product(range(-N, N), [N], zs):
                //     yield r
                case 0:
                    ++m_current_x;
                    wrapped = m_current_x >= m_range;
                    m_current_x -= 2*wrapped*m_range;
                    if(!m_is2D)
                        {
                        m_current_z += wrapped;
                        wrapped = m_current_z >= m_range;
                        m_current_z += wrapped*(1 - 2*m_range);
                        }
                    if(wrapped)
                        {
                        ++m_stage;
                        m_current_x = m_range;
                        }
                    break;
                    // +x wedge: iterate over y and (possibly) z
                    // for r in itertools.product([N], range(N, -N, -1), zs):
                    //     yield r
                case 1:
                    --m_current_y;
                    wrapped = m_current_y <= -m_range;
                    m_current_y += 2*wrapped*m_range;
                    if(!m_is2D)
                        {
                        m_current_z += wrapped;
                        wrapped = m_current_z >= m_range;
                        m_current_z += wrapped*(1 - 2*m_range);
                        }
                    if(wrapped)
                        {
                        ++m_stage;
                        m_current_y = -m_range;
                        }
                    break;
                    // -y wedge: iterate over x and (possibly) z
                    // for r in itertools.product(range(N, -N, -1), [-N], zs):
                    //     yield r
                case 2:
                    --m_current_x;
                    wrapped = m_current_x <= -m_range;
                    m_current_x += 2*wrapped*m_range;
                    if(!m_is2D)
                        {
                        m_current_z += wrapped;
                        wrapped = m_current_z >= m_range;
                        m_current_z += wrapped*(1 - 2*m_range);
                        }
                    if(wrapped)
                        {
                        ++m_stage;
                        m_current_x = -m_range;
                        }
                    break;
                    // -x wedge: iterate over y and (possibly) z
                    // for r in itertools.product([-N], range(-N, N), zs):
                    //     yield r
                case 3:
                    ++m_current_y;
                    wrapped = m_current_y >= m_range;
                    m_current_y -= 2*wrapped*m_range;
                    if(!m_is2D)
                        {
                        m_current_z += wrapped;
                        wrapped = m_current_z >= m_range;
                        m_current_z += wrapped*(1 - 2*m_range);
                        }
                    if(wrapped)
                        {
                        if(m_is2D) // we're done for this range
                            reset(m_range + 1);
                        else
                            {
                            ++m_stage;
                            m_current_x = -m_range;
                            m_current_y = -m_range;
                            m_current_z = -m_range;
                            }
                        }
                    break;
                    // -z face and +z face: iterate over x and y
                    // grid = list(range(-N, N + 1))
                    // if threeD:
                    //     # make front and back in z
                    //     for (x, y) in itertools.product(grid, grid):
                    //         yield (x, y, N)
                    //         if N > 0:
                    //             yield (x, y, -N)
                    // elif N == 0:
                    //     yield (0, 0, 0)
                case 4:
                case 5:
                default:
                    ++m_current_x;
                    wrapped = m_current_x > m_range;
                    m_current_x -= wrapped*(2*m_range + 1);
                    m_current_y += wrapped;
                    wrapped = m_current_y > m_range;
                    m_current_y -= wrapped*(2*m_range + 1);
                    if(wrapped)
                        {
                        // 2D cases have already moved to the next stage by
                        // this point, only deal with 3D
                        ++m_stage;
                        m_current_z = m_range;

                        // if we're done, move on to the next range
                        if(m_stage > 5)
                            reset(m_range + 1);
                        }
                    break;
                }
            }

        vec3<int> operator*()
            {
            return vec3<int>(m_current_x, m_current_y, m_current_z);
            }

        bool operator==(const IteratorCellShell &other)
            {
            return m_range == other.m_range && m_current_x == other.m_current_x && m_current_y == other.m_current_y && m_current_z == other.m_current_z && m_stage == other.m_stage && m_is2D == other.m_is2D;
            }

        bool operator!=(const IteratorCellShell &other)
            {
            return !(*this == other);
            }

    private:
        void reset(unsigned int range)
            {
            m_range = range;
            m_stage = 0;
            m_current_x = -m_range;
            m_current_y = m_range;
            if(m_is2D)
                {
                m_current_z = 0;
                }
            else
                {
                m_current_z = -m_range + 1;
                }

            if(range == 0)
                {
                m_current_z = 0;
                // skip to the last stage
                m_stage = 5;
                }
            }
        int m_range;      //!< Find cells this many cells away
        int m_current_x;  //!< Current position in x
        int m_current_y;  //!< Current position in y
        int m_current_z;  //!< Current position in z
        char m_stage;     //!< stage of the computation (which face is being iterated over)
        bool m_is2D;      //!< true if the cell list is 2D
};

bool compareFirstNeighborPairs(const std::vector<std::tuple<size_t, size_t, float> > &left,
                               const std::vector<std::tuple<size_t, size_t, float> > &right);

//! Computes a cell id for each particle and a link cell data structure for iterating through it
/*! For simplicity in only needing a small number of arrays, the link cell
    algorithm is used to generate and store the cell list data for particles.

    Cells are given a nominal minimum width \a cell_width. Each dimension of
    the box is split into an integer number of cells no smaller than
    \a cell_width wide in that dimension. The actual number of cells along
    each dimension is stored in an Index3D which is also used to compute the
    cell index from (i,j,k).

    The cell coordinate (i,j,k) itself is computed like so:
    \code
    i = floorf((x + Lx/2) / w) % Nw
    \endcode
    and so on for j, k (y, z). Call getCellCoord() to do this computation for
    an arbitrary point.

    <b>Data structures:</b><br>
    The internal data structure used in LinkCell is a linked list of particle
    indices. See IteratorLinkCell for information on how to iterate through these.

    <b>2D:</b><br>
    LinkCell properly handles 2D boxes. When a 2D box is handed to LinkCell,
    it creates an m x n x 1 cell list and neighbor cells are only listed in
    the plane. As with everything else in freud, 2D points must be passed in
    as 3 component vectors x,y,0. Failing to set 0 in the third component will
    lead to undefined behavior.
*/
class LinkCell
    {
    public:
        //! iterator to iterate over particles in the cell
        typedef IteratorLinkCell iteratorcell;

        //! Constructor
        LinkCell(const box::Box& box, float cell_width);

        //! Null Constructor for triclinic behavior
        LinkCell();

        //! Update cell_width
        void setCellWidth(float cell_width);

        //! Update box used in linkCell
        void updateBox(const box::Box& box);

        //! Compute LinkCell dimensions
        const vec3<unsigned int> computeDimensions(const box::Box& box, float cell_width) const;

        //! Get the simulation box
        const box::Box& getBox() const
            {
            return m_box;
            }

        //! Get the cell indexer
        const Index3D& getCellIndexer() const
            {
            return m_cell_index;
            }

        //! Get the number of cells
        unsigned int getNumCells() const
            {
            return m_cell_index.getNumElements();
            }

        //! Get the cell width
        float getCellWidth() const
            {
            return m_cell_width;
            }

        //! Compute the cell id for a given position
        unsigned int getCell(const vec3<float>& p) const
            {
            vec3<unsigned int> c = getCellCoord(p);
            return m_cell_index(c.x, c.y, c.z);
            }

        //! Compute cell coordinates for a given position
        vec3<unsigned int> getCellCoord(const vec3<float> p) const
            {
            vec3<float> alpha = m_box.makeFraction(p);
            vec3<unsigned int> c;
            c.x = floorf(alpha.x * float(m_cell_index.getW()));
            c.x %= m_cell_index.getW();
            c.y = floorf(alpha.y * float(m_cell_index.getH()));
            c.y %= m_cell_index.getH();
            c.z = floorf(alpha.z * float(m_cell_index.getD()));
            c.z %= m_cell_index.getD();
            return c;
            }

        //! Iterate over particles in a cell
        iteratorcell itercell(unsigned int cell) const
            {
            assert(m_cell_list.get() != NULL);
            return iteratorcell(m_cell_list, m_Np, getNumCells(), cell);
            }

        //! Get a list of neighbors to a cell
        const std::vector<unsigned int>& getCellNeighbors(unsigned int cell) const
            {
            return m_cell_neighbors[cell];
            }

        //! Compute the cell list
        void computeCellList(box::Box& box, const vec3<float> *points, unsigned int Np);

        //! Compute the neighbor list using the cell list
        void compute(box::Box& box, const vec3<float> *ref_points,
                     unsigned int Nref, const vec3<float> *points=0, unsigned int Np=0,
                     bool exclude_ii=true);

        NeighborList *getNeighborList()
        {
            return &m_neighbor_list;
        }

    private:

        //! Rounding helper function.
        static unsigned int roundDown(unsigned int v, unsigned int m);

        box::Box m_box;                //!< Simulation box where the particles belong
        Index3D m_cell_index;          //!< Indexer to compute cell indices
        unsigned int m_Np;             //!< Number of particles last placed into the cell list
        unsigned int m_Nc;             //!< Number of cells last used
        float m_cell_width;            //!< Minimum necessary cell width cutoff
        vec3<unsigned int> m_celldim;  //!< Cell dimensions

        std::shared_ptr<unsigned int> m_cell_list;  //!< The cell list last computed

        std::vector< std::vector<unsigned int> > m_cell_neighbors;  //!< List of cell neighbors to each cell

        NeighborList m_neighbor_list;  //!< Stored neighbor list

        //! Helper function to compute cell neighbors
        void computeCellNeighbors();
    };

}; }; // end namespace freud::locality

#endif // LINKCELL_H
