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

 Infomap software package for multi-level network clustering

 Copyright (c) 2013, 2014 Daniel Edler, Martin Rosvall

 For more information, see <http://www.mapequation.org>


 This file is part of Infomap software package.

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

 Infomap software package 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 Affero General Public License for more details.

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

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


#ifndef TREEITERATORS_H_
#define TREEITERATORS_H_

#include <cstddef>
#include <iterator>
#include <deque>

namespace infomap {

#ifndef ASSERT
#include <cassert>
#define ASSERT(x) assert(x)
#endif

using std::iterator_traits;


/**
 * Child iterator.
 */
template <typename NodePointerType> // pointer or const pointer
class ChildIterator {
  typedef std::bidirectional_iterator_tag iterator_category;
  typedef typename iterator_traits<NodePointerType>::value_type value_type;
  typedef typename iterator_traits<NodePointerType>::difference_type difference_type;
  typedef typename iterator_traits<NodePointerType>::reference reference;
  typedef typename iterator_traits<NodePointerType>::pointer pointer;

protected:
  NodePointerType m_root = nullptr;
  NodePointerType m_current = nullptr;

public:
  ChildIterator() {}

  explicit ChildIterator(const NodePointerType& nodePointer)
      : m_root(nodePointer), m_current(nodePointer == nullptr ? nullptr : nodePointer->firstChild) {}

  ChildIterator(const ChildIterator& other)
      : m_root(other.m_root), m_current(other.m_current) {}

  ChildIterator& operator=(const ChildIterator& other)
  {
    m_root = other.m_root;
    m_current = other.m_current;
    return *this;
  }

  pointer current() const
  {
    return m_current;
  }

  reference
  operator*() const
  {
    return *m_current;
  }

  pointer
  operator->() const
  {
    return m_current;
  }

  bool operator==(const ChildIterator& rhs) const
  {
    return m_current == rhs.m_current;
  }

  bool operator!=(const ChildIterator& rhs) const
  {
    return !(m_current == rhs.m_current);
  }

  bool isEnd() const
  {
    return m_current == nullptr;
  }

  ChildIterator&
  operator++()
  {
    m_current = m_current->next;
    if (m_current != nullptr && m_current->parent != m_root) {
      m_current = nullptr;
    }
    return *this;
  }

  ChildIterator
  operator++(int)
  {
    ChildIterator copy(*this);
    ++(*this);
    return copy;
  }

  ChildIterator&
  operator--()
  {
    m_current = m_current->previous;
    if (m_current != nullptr && m_current->parent != m_root) {
      m_current = nullptr;
    }
    return *this;
  }

  ChildIterator
  operator--(int)
  {
    ChildIterator copy(*this);
    --(*this);
    return copy;
  }
};

/**
 * Tree iterator.
 */
template <typename NodePointerType> // pointer or const pointer
class TreeIterator {
  typedef std::forward_iterator_tag iterator_category;
  typedef typename iterator_traits<NodePointerType>::value_type value_type;
  typedef typename iterator_traits<NodePointerType>::difference_type difference_type;
  typedef typename iterator_traits<NodePointerType>::reference reference;
  typedef typename iterator_traits<NodePointerType>::pointer pointer;

protected:
  NodePointerType m_root = nullptr;
  NodePointerType m_current = nullptr;
  int m_moduleIndexLevel = -1;
  unsigned int m_moduleIndex = 0;
  std::deque<unsigned int> m_path; // The child index path to current node
  unsigned int m_depth = 0;

public:
  TreeIterator() {}

  TreeIterator(NodePointerType nodePointer, int moduleIndexLevel = -1)
      : m_root(nodePointer),
        m_current(nodePointer),
        m_moduleIndexLevel(moduleIndexLevel)
  {
  }

  TreeIterator(const TreeIterator& other)
      : m_root(other.m_root),
        m_current(other.m_current),
        m_moduleIndexLevel(other.m_moduleIndexLevel),
        m_moduleIndex(other.m_moduleIndex),
        m_path(other.m_path),
        m_depth(other.m_depth)
  {
  }

  virtual ~TreeIterator() {}

  TreeIterator& operator=(const TreeIterator& other)
  {
    m_root = other.m_root;
    m_current = other.m_current;
    m_moduleIndexLevel = other.m_moduleIndexLevel;
    m_moduleIndex = other.m_moduleIndex;
    m_path = other.m_path;
    m_depth = other.m_depth;
    return *this;
  }


  pointer current() const
  {
    return m_current;
  }

  reference
  operator*() const
  {
    return *m_current;
  }

  pointer
  operator->() const
  {
    return m_current;
  }

  bool operator==(const TreeIterator& rhs) const
  {
    return m_current == rhs.m_current;
  }

  bool operator!=(const TreeIterator& rhs) const
  {
    return !(m_current == rhs.m_current);
  }

  const std::deque<unsigned int>& path() const
  {
    return m_path;
  }

  unsigned int moduleIndex() const
  {
    return m_moduleIndex;
  }

  unsigned int depth() const
  {
    return m_depth;
  }

  bool isEnd() const
  {
    return m_current == nullptr;
  }

  TreeIterator& operator++()
  {
    NodePointerType curr = m_current;
    NodePointerType infomapRoot = curr->getInfomapRoot();
    if (infomapRoot != nullptr) {
      curr = infomapRoot;
    }

    if (curr->firstChild != nullptr) {
      curr = curr->firstChild;
      ++m_depth;
      m_path.push_back(0);
    } else {
    // Current node is a leaf
    // Presupposes that the next pointer can't reach out from the current parent.
    tryNext:
      while (curr->next == nullptr) {
        if (curr->parent != nullptr) {
          curr = curr->parent;
          --m_depth;
          m_path.pop_back();
          if (curr == m_root) // Check if back to beginning
          {
            m_current = nullptr;
            return *this;
          }
          if (m_moduleIndexLevel < 0) {
            if (curr->isLeafModule()) // TODO: Generalize to -2 for second level to bottom
              ++m_moduleIndex;
          } else if (static_cast<unsigned int>(m_moduleIndexLevel) == m_depth)
            ++m_moduleIndex;
        } else {
          NodePointerType infomapOwner = curr->owner;
          if (infomapOwner != nullptr) {
            curr = infomapOwner;
            if (curr == m_root) // Check if back to beginning
            {
              m_current = nullptr;
              return *this;
            }
            goto tryNext;
          } else // null also if no children in first place
          {
            m_current = nullptr;
            return *this;
          }
        }
      }
      curr = curr->next;
      ++m_path.back();
    }
    m_current = curr;
    return *this;
  }

  TreeIterator
  operator++(int)
  {
    TreeIterator copy(*this);
    ++(*this);
    return copy;
  }

  TreeIterator& stepForward()
  {
    ++(*this);
    return *this;
  }
};


/**
 * Base node iterator.
 */
template <typename NodePointerType, typename iterator_tag = std::bidirectional_iterator_tag>
struct node_iterator_base {
  //	typedef typename iterator_traits<NodePointerType>::iterator_category	iterator_category; //random_access_iterator_tag
  //	typedef std::forward_iterator_tag										iterator_category;
  typedef iterator_tag iterator_category;
  typedef typename iterator_traits<NodePointerType>::value_type value_type;
  typedef typename iterator_traits<NodePointerType>::difference_type difference_type;
  typedef typename iterator_traits<NodePointerType>::reference reference;
  typedef typename iterator_traits<NodePointerType>::pointer pointer;

  node_iterator_base()
      : m_current(nullptr)
  {
  }

  explicit node_iterator_base(const NodePointerType& nodePointer)
      : m_current(nodePointer)
  {
  }

  node_iterator_base(const node_iterator_base& other)
      : m_current(other.m_current)
  {
  }

  node_iterator_base& operator=(const node_iterator_base& other)
  {
    m_current = other.m_current;
    return *this;
  }

  virtual ~node_iterator_base() {}

  pointer base() const
  {
    return m_current;
  }

  reference
  operator*() const
  {
    return *m_current;
  }

  pointer
  operator->() const
  {
    return m_current;
  }

  bool operator==(const node_iterator_base& rhs) const
  {
    return m_current == rhs.m_current;
  }

  bool operator!=(const node_iterator_base& rhs) const
  {
    return !(m_current == rhs.m_current);
  }

  bool isEnd() const
  {
    return m_current == nullptr;
  }

protected:
  NodePointerType m_current;
};

template <typename NodePointerType>
class DepthFirstIteratorBase : public node_iterator_base<NodePointerType> {
  typedef node_iterator_base<NodePointerType> Base;

public:
  DepthFirstIteratorBase()
      : Base(),
        m_root(nullptr),
        m_depth(0)
  {
  }

  explicit DepthFirstIteratorBase(const NodePointerType& nodePointer)
      : Base(nodePointer),
        m_root(nodePointer),
        m_depth(0)
  {
  }

  DepthFirstIteratorBase(const DepthFirstIteratorBase& other)
      : Base(other),
        m_root(other.m_root),
        m_depth(other.m_depth)
  {
  }

  DepthFirstIteratorBase& operator=(const DepthFirstIteratorBase& other)
  {
    Base::operator=(other);
    m_root = other.m_root;
    m_depth = other.m_depth;
    return *this;
  }

  unsigned int depth() const
  {
    return m_depth;
  }

protected:
  NodePointerType m_root;
  unsigned int m_depth;
  using Base::m_current;
};


/**
 * Pre processing depth first iterator
 * Note:
 * This iterator presupposes that the next pointer of a node can't reach a node with a different parent.
 */
template <typename NodePointerType, bool pre_t = true>
class DepthFirstIterator : public DepthFirstIteratorBase<NodePointerType> {
  typedef DepthFirstIteratorBase<NodePointerType> Base;

public:
  DepthFirstIterator() : Base() {}

  explicit DepthFirstIterator(const NodePointerType& nodePointer) : Base(nodePointer) {}

  DepthFirstIterator(const DepthFirstIterator& other) : Base(other) {}

  DepthFirstIterator& operator=(const DepthFirstIterator& other)
  {
    Base::operator=(other);
    return *this;
  }

  DepthFirstIterator&
  operator++()
  {
    NodePointerType curr = Base::m_current;
    if (curr->firstChild != nullptr) {
      curr = curr->firstChild;
      ++Base::m_depth;
    } else {
      // Presupposes that the next pointer can't reach out from the current parent.
      while (curr->next == nullptr) {
        curr = curr->parent;
        --Base::m_depth;
        if (curr == Base::m_root || curr == nullptr) // 0 if no children in first place
        {
          Base::m_current = nullptr;
          return *this;
        }
      }
      curr = curr->next;
    }
    Base::m_current = curr;
    return *this;
  }

  DepthFirstIterator
  operator++(int)
  {
    auto copy(*this);
    ++(*this);
    return copy;
  }

  DepthFirstIterator
  next()
  {
    auto copy(*this);
    return ++copy;
  }
};


/**
 * Post processing depth first iterator
 * Note:
 * This iterator presupposes that the next pointer of a node can't reach a node with a different parent.
 */
template <typename NodePointerType>
class DepthFirstIterator<NodePointerType, false> : public DepthFirstIteratorBase<NodePointerType> {
  typedef DepthFirstIteratorBase<NodePointerType> Base;

public:
  DepthFirstIterator() : Base() {}

  explicit DepthFirstIterator(const NodePointerType& nodePointer) : Base(nodePointer) { init(); }

  DepthFirstIterator(const DepthFirstIterator& other) : Base(other) {}

  DepthFirstIterator& operator=(const DepthFirstIterator& other)
  {
    Base::operator=(other);
    return *this;
  }

  void init()
  {
    if (Base::m_current != nullptr) {
      while (Base::m_current->firstChild != nullptr) {
        Base::m_current = Base::m_current->firstChild;
        ++Base::m_depth;
      }
    }
  }

  DepthFirstIterator&
  operator++()
  {
    // The root should be the last node
    if (Base::m_current == Base::m_root) {
      Base::m_current = nullptr;
      return *this;
    }

    NodePointerType curr = Base::m_current;
    if (curr->next != nullptr) {
      curr = curr->next;
      while (curr->firstChild != nullptr) {
        curr = curr->firstChild;
        ++Base::m_depth;
      }
    } else {
      curr = curr->parent;
      --Base::m_depth;
    }

    Base::m_current = curr;

    return *this;
  }

  DepthFirstIterator
  operator++(int)
  {
    DepthFirstIterator copy(*this);
    ++(*this);
    return copy;
  }

  DepthFirstIterator
  next()
  {
    DepthFirstIterator copy(*this);
    return ++copy;
  }
};


/**
 * Leaf node iterator
 */
template <typename NodePointerType>
class LeafNodeIterator : public DepthFirstIteratorBase<NodePointerType> {
  typedef DepthFirstIteratorBase<NodePointerType> Base;

public:
  LeafNodeIterator() : Base() {}

  explicit LeafNodeIterator(const NodePointerType& nodePointer) : Base(nodePointer) { init(); }

  LeafNodeIterator(const LeafNodeIterator& other) : Base(other) {}

  LeafNodeIterator& operator=(const LeafNodeIterator& other)
  {
    Base::operator=(other);
    return *this;
  }

  void init()
  {
    if (Base::m_current != nullptr) {
      while (Base::m_current->firstChild != nullptr) {
        Base::m_current = Base::m_current->firstChild;
        ++Base::m_depth;
      }
    }
  }

  LeafNodeIterator&
  operator++()
  {
    ASSERT(Base::m_current != nullptr);
    while (Base::m_current->next == nullptr || Base::m_current->next->parent != Base::m_current->parent) {
      Base::m_current = Base::m_current->parent;
      --Base::m_depth;
      if (Base::m_current == nullptr)
        return *this;
    }

    Base::m_current = Base::m_current->next;

    if (Base::m_current != nullptr) {
      while (Base::m_current->firstChild != nullptr) {
        Base::m_current = Base::m_current->firstChild;
        ++Base::m_depth;
      }
    }
    return *this;
  }

  LeafNodeIterator
  operator++(int)
  {
    LeafNodeIterator copy(*this);
    ++(*this);
    return copy;
  }

  LeafNodeIterator
  next()
  {
    LeafNodeIterator copy(*this);
    return ++copy;
  }
};


/**
 * Leaf module iterator
 */
template <typename NodePointerType>
class LeafModuleIterator : public DepthFirstIteratorBase<NodePointerType> {
  typedef DepthFirstIteratorBase<NodePointerType> Base;

public:
  LeafModuleIterator() : Base() {}

  explicit LeafModuleIterator(const NodePointerType& nodePointer) : Base(nodePointer) { init(); }

  LeafModuleIterator(const LeafModuleIterator& other) : Base(other) {}

  LeafModuleIterator& operator=(const LeafModuleIterator& other)
  {
    Base::operator=(other);
    init();
    return *this;
  }

  void init()
  {
    if (Base::m_current != nullptr) {
      if (Base::m_current->firstChild == nullptr) {
        Base::m_current = nullptr; // End directly if no module
      } else {
        while (Base::m_current->firstChild->firstChild != nullptr) {
          Base::m_current = Base::m_current->firstChild;
          ++Base::m_depth;
        }
      }
    }
  }

  LeafModuleIterator&
  operator++()
  {
    ASSERT(Base::m_current != nullptr);
    while (Base::m_current->next == nullptr || Base::m_current->next->parent != Base::m_current->parent) {
      Base::m_current = Base::m_current->parent;
      --Base::m_depth;
      if (Base::m_current == nullptr)
        return *this;
    }

    Base::m_current = Base::m_current->next;

    if (Base::m_current != nullptr) {
      if (Base::m_current->firstChild == nullptr) {
        Base::m_current = Base::m_current->parent;
      } else {
        while (Base::m_current->firstChild->firstChild != nullptr) {
          Base::m_current = Base::m_current->firstChild;
          ++Base::m_depth;
        }
      }
    }
    return *this;
  }

  LeafModuleIterator
  operator++(int)
  {
    LeafModuleIterator copy(*this);
    ++(*this);
    return copy;
  }

  LeafModuleIterator
  next()
  {
    LeafModuleIterator copy(*this);
    return ++copy;
  }
};


/**
 * Sibling iterator.
 */
template <typename NodePointerType> // pointer or const pointer
class SiblingIterator : public node_iterator_base<NodePointerType> {
  typedef node_iterator_base<NodePointerType> Base;

public:
  typedef SiblingIterator<NodePointerType> self_type;

  SiblingIterator() : Base() {}

  explicit SiblingIterator(const NodePointerType& nodePointer) : Base(nodePointer) {}

  SiblingIterator(const SiblingIterator& other) : Base(other) {}

  SiblingIterator& operator=(const SiblingIterator& other)
  {
    Base::operator=(other);
    return *this;
  }

  SiblingIterator&
  operator++()
  {
    ASSERT(Base::m_current != nullptr);
    Base::m_current = Base::m_current->next;
    return *this;
  }

  SiblingIterator
  operator++(int)
  {
    SiblingIterator copy(*this);
    ++(*this);
    return copy;
  }

  SiblingIterator&
  operator--()
  {
    ASSERT(Base::m_current != nullptr);
    Base::m_current = Base::m_current->previous;
    return *this;
  }

  SiblingIterator
  operator--(int)
  {
    SiblingIterator copy(*this);
    --(*this);
    return copy;
  }
};

} // namespace infomap

#endif /* TREEITERATORS_H_ */
