/*
 *  SparseQuadTree.hpp
 *
 *  Created by thant on 5/21/2013.
 *  Copyright 2013 TeleNav. All rights reserved.
 *
 *  The idea behind this data structure is to provide the ability to
 *  store values at arbitrary nodes of a quad tree. As items are
 *  inserted into the quad tree, the tree is grown as needed to
 *  instantiate the node where the item will be stored. As items are
 *  removed, the tree is trimmed. Every node is either a branch or a
 *  leaf. All nodes either contain an item, or are the ancestor of a
 *  node that contains an item. In other words, leaves will always be
 *  occupied, and branches may or may not be occupied.
 *
 *  This data structure is implemented in two parts. The first is the
 *  actual node SparseQuadNode. The second is a container class
 *  SparseQuadTree. The programmer can ascend and descend the tree
 *  using SparseQuadNode, but he is not allowed to modify the topology
 *  of the tree. The SparseQuadTree is responsible for that. This is
 *  to insure the tree is properly trimmed as items are removed, and
 *  that therefore there are no memory leaks.
 */

#ifndef Foundation_Math_SparseQuadTree_INCLUDED
#define Foundation_Math_SparseQuadTree_INCLUDED

#include <map>

#include <boost/shared_ptr.hpp>
#include <boost/enable_shared_from_this.hpp>
#include <boost/weak_ptr.hpp>
#include <boost/optional.hpp>
#include <boost/none_t.hpp>

#include <foundation/Math/QuadCoord.hpp>
#include <foundation/Thunk/Thunk.hpp>

namespace Tn {

namespace math {


template <typename T>
struct SparseQuadNode : boost::enable_shared_from_this<SparseQuadNode<T> > {

    typedef boost::shared_ptr<SparseQuadNode> ptr;
    typedef boost::weak_ptr<SparseQuadNode> weakPtr;
    typedef boost::shared_ptr<SparseQuadNode const> constPtr;

    static ptr create() {
        return ptr(new SparseQuadNode<T>());
    }

    void setValue(T const & value) {
        m_valueOpt = value;
    }

    boost::optional<T> getValue() const {
        return m_valueOpt;
    }

    bool isEmpty() const {
        return !m_valueOpt;
    }

    eQuadrant getQuadrant() const {
        return m_quadrant;
    }

    bool isLeaf() const {
        return !(m_bl || m_br || m_tl || m_tr);
    }

    ptr getOrCreateBranch(eQuadrant q) {
        switch (q) {

        case eQuadrant_BottomLeft:
            if (!m_bl) {
                ptr parent = SparseQuadNode<T>::shared_from_this();
                m_bl = ptr(new SparseQuadNode(parent, eQuadrant_BottomLeft));
            }
            return m_bl;

        case eQuadrant_BottomRight:
            if (!m_br) {
                ptr parent = SparseQuadNode<T>::shared_from_this();
                m_br = ptr(new SparseQuadNode(parent, eQuadrant_BottomRight));
            }
            return m_br;

        case eQuadrant_TopLeft:
            if (!m_tl) {
                ptr parent = SparseQuadNode<T>::shared_from_this();
                m_tl = ptr(new SparseQuadNode(parent, eQuadrant_TopLeft));
            }
            return m_tl;

        case eQuadrant_TopRight:
            if (!m_tr) {
                ptr parent = SparseQuadNode<T>::shared_from_this();
                m_tr = ptr(new SparseQuadNode(parent, eQuadrant_TopRight));
            }
            return m_tr;
        }
        return ptr();
    }


    constPtr getBranch(eQuadrant q) const {
        switch (q) {
        case eQuadrant_BottomLeft:
            return m_bl;
        case eQuadrant_BottomRight:
            return m_br;
        case eQuadrant_TopLeft:
            return m_tl;
        case eQuadrant_TopRight:
            return m_tr;
        }
        return ptr();
    }


    ptr getParent() {
        ptr parent = m_parent.lock();
        return parent;
    }

    constPtr getParent() const {
        constPtr parent = m_parent.lock();
        return parent;
    }


    bool clearValue() {
        bool set = !isEmpty();
        m_valueOpt = boost::optional<T>();
        return set;
    }


    void removeBranch(eQuadrant q) {
        switch (q)
        {
        case eQuadrant_BottomLeft:
            m_bl.reset();
            return;
        case eQuadrant_BottomRight:
            m_br.reset();
            return;
        case eQuadrant_TopLeft:
            m_tl.reset();
            return;
        case eQuadrant_TopRight:
            m_tr.reset();
            return;
        }
    }

private:

    SparseQuadNode() : m_quadrant(eQuadrant(0) /* invalid */) {}  // root

    SparseQuadNode(typename SparseQuadNode<T>::ptr const & parent, 
                   eQuadrant quadrant) :
        m_parent(parent),
        m_quadrant(quadrant) {}

    weakPtr m_parent;

    eQuadrant m_quadrant;

    ptr m_bl;
    ptr m_br;
    ptr m_tl;
    ptr m_tr;

    boost::optional<T> m_valueOpt;
};


template <typename T>
struct SparseQuadTree {

    /*  This keeps both a sparse quad tree, and a map of pointers to
     *  occupied nodes indexed by QuadCoord. This allows for both
     *  tree-like operations on the quad-tree itself, and efficient
     *  searces and iteration through the occupied nodes.
     *
     *  
     */

    typedef std::map<QuadCoord, typename SparseQuadNode<T>::ptr> Nodes;
    typedef std::pair<QuadCoord, typename SparseQuadNode<T>::ptr> NodesEntry;


    SparseQuadTree() :
        m_root(SparseQuadNode<T>::create()) {

    }

    // Returns true if node was created, false otherwise.

    bool set(QuadCoord const & qc, T const & value) {

        typename Nodes::iterator it = m_nodes.find(qc);
        if (it != m_nodes.end()) {
            it->second->setValue(value);
            return false;
        }

        typename SparseQuadNode<T>::ptr node = findOrCreate(m_root, qc);
        assert(node);
        node->setValue(value);
        m_nodes.insert(NodesEntry(qc, node));
        return true;
    }

    boost::optional<T> get(QuadCoord const & qc) {

        typename Nodes::iterator it = m_nodes.find(qc);

        if (it != m_nodes.end()) {
            return it->second->getValue();
        }

        return boost::optional<T>();
    }

    // Returns true if node existed and value was removed. False otherwise.
    bool remove(QuadCoord const & qc) {

        typename Nodes::iterator it = m_nodes.find(qc);

        if (it == m_nodes.end()) {
            return false;
        }

        typename SparseQuadNode<T>::ptr node = it->second;
        if (!node->clearValue()) {
            return false;
        }

        m_nodes.erase(it);
        trim(node);

        return true;
    }

    /*  This applies a function to every occupied node.
     */

    void apply(typename Tn::Foundation::Thunk1<T,void>::ptr const & thunk) {

        typename Nodes::iterator begin = m_nodes.begin();
        typename Nodes::iterator end = m_nodes.end();

        apply(thunk, begin, end);
    }


    /*  This folds and map through all occupied nodes.
     */
    template <typename Result>
    Result foldl(typename Tn::Foundation::Thunk2<T,Result,Result>::ptr const & thunk, Result const & init) {

        typename Nodes::iterator begin = m_nodes.begin();
        typename Nodes::iterator end = m_nodes.end();

        return foldl(thunk, init, begin, end);
    }

    /*  This only returns a valid pointer if the node actually exists
     *  and has a value.
     */
    typename SparseQuadNode<T>::constPtr find(QuadCoord const & qc) const {
        return find(m_root, qc);
    }

    /*  This is only for testing. This checks that:
     *
     *  1) All leaves have values. No leaves should be empty. (An
     *  empty leaf implies that the tree isn't being trimmed properly
     *  as values are removed.)
     *
     *  2) The number of live nodes in the tree equals the number of
     *  values in the live dictionary.
     */

    std::pair<bool, uint32_t> valid() const {

        typename SparseQuadNode<T>::constPtr root = m_root;
        std::pair<bool, uint32_t> result = valid(root);

        bool good = result.first && (result.second == m_nodes.size());
        uint32_t count = result.second;

        return std::pair<bool, uint32_t>(good, count);
    }

private:

    void trim(typename SparseQuadNode<T>::ptr const & node) {

        if (node && node->isEmpty() && node->isLeaf()) {
            eQuadrant quadrant = node->getQuadrant();
            typename SparseQuadNode<T>::ptr parent = node->getParent();
            if (parent) {
                parent->removeBranch(quadrant);
                trim(parent);
            }
        }
    }


    static void apply(typename Tn::Foundation::Thunk1<T,void>::ptr const & thunk,
                      typename Nodes::iterator const & i,
                      typename Nodes::iterator const & end) {

        if (i==end) return;

        typename SparseQuadNode<T>::ptr const & node = i->second;
        boost::optional<T> valueOpt = node->getValue();

        if (valueOpt) {

            T const& value = *valueOpt;
            (*thunk)(value);
            
            typename Nodes::iterator next = i;
            ++next;

            apply(thunk, next, end);
            return;
        }
        // shouldn't get here
        assert(false);
    }

    template <typename Result>
    static Result foldl(typename Tn::Foundation::Thunk2<T,Result,Result>::ptr const & thunk,
                        Result const & init,
                        typename Nodes::iterator const & i,
                        typename Nodes::iterator const & end) {

        if (i==end) return init;

        typename SparseQuadNode<T>::ptr const & node = i->second;
        boost::optional<T> valueOpt = node->getValue();

        if (valueOpt) {

            T const & value = *valueOpt;
            Result result = (*thunk)(value, init);

            typename Nodes::iterator next = i;
            ++next;

            return foldl(thunk, result, next, end);

        }
        // shouldn't get here
        assert(false);
        return init;
    }
            

    static 
    typename SparseQuadNode<T>::ptr 
    findOrCreate(typename SparseQuadNode<T>::ptr const & node,
                 QuadCoord const & qc,
                 uint32_t i = 0) {

        assert(node);

        if (i>=qc.level) {

            return node;
        }

        eQuadrant branch = qc[i];
        typename SparseQuadNode<T>::ptr next = node->getOrCreateBranch(branch);
        return findOrCreate(next, qc, i+1);
    }

    // Const version. Stops at leaves. Doesn't return empty nodes even
    // if they exist.
    static 
    typename SparseQuadNode<T>::constPtr 
    find(typename SparseQuadNode<T>::constPtr const & node,
         QuadCoord const & qc,
         uint32_t i = 0) {

        if (!node) return typename SparseQuadNode<T>::constPtr();

        if (i>=qc.level) {

            if (node->isEmpty()) return SparseQuadNode<T>::constPtr();

            return node;
        }

        eQuadrant branch = qc[i];
        typename SparseQuadNode<T>::constPtr next = node->getChild(branch);
        return find(next, qc, i+1);
    }


    typename SparseQuadNode<T>::ptr m_root;

    Nodes m_nodes;

    std::pair<bool, uint32_t> valid(typename SparseQuadNode<T>::constPtr const & node) const {

        if (!node) return std::pair<bool, uint32_t>(true, 0);

        std::pair<bool, uint32_t> bl = valid(node->getBranch(eQuadrant_BottomLeft));
        std::pair<bool, uint32_t> br = valid(node->getBranch(eQuadrant_BottomRight));
        std::pair<bool, uint32_t> tl = valid(node->getBranch(eQuadrant_TopLeft));
        std::pair<bool, uint32_t> tr = valid(node->getBranch(eQuadrant_TopRight));

        uint32_t sum 
            = bl.second 
            + br.second 
            + tl.second 
            + tr.second + 
            ( node->isEmpty() ? 0 : 1 );

        bool valid = 
            (!node->isEmpty() || !node->isLeaf()) &&
            bl.first && br.first && tl.first && tr.first;

        return std::pair<bool,uint32_t>(valid, sum);
    }
};


} // namespace math

} // namespace Tn

#endif
