/*
 * Copyright 2003-2007, Ingo Weinhold <bonefish@cs.tu-berlin.de>.
 * Distributed under the terms of the MIT License.
 */
#ifndef TWO_KEY_AVL_TREE_H
#define TWO_KEY_AVL_TREE_H

#include <util/AVLTreeMap.h>


// TwoKeyAVLTreeKey
template<typename PrimaryKey, typename SecondaryKey>
class TwoKeyAVLTreeKey {
public:
        inline TwoKeyAVLTreeKey(const PrimaryKey &primary,
                                                        const SecondaryKey &secondary)
                : primary(primary),
                  secondary(secondary),
                  use_secondary(true)
        {
        }

        inline TwoKeyAVLTreeKey(const PrimaryKey *primary)
                : primary(primary),
                  secondary(NULL),
                  use_secondary(false)
        {
        }

        PrimaryKey              primary;
        SecondaryKey    secondary;
        bool                    use_secondary;
};

// TwoKeyAVLTreeKeyCompare
template<typename PrimaryKey, typename SecondaryKey,
                 typename PrimaryKeyCompare, typename SecondaryKeyCompare>
class TwoKeyAVLTreeKeyCompare {
private:
        typedef TwoKeyAVLTreeKey<PrimaryKey, SecondaryKey> Key;

public:
        inline TwoKeyAVLTreeKeyCompare(const PrimaryKeyCompare &primary,
                                                                   const SecondaryKeyCompare &secondary)
                : fPrimaryKeyCompare(primary), fSecondaryKeyCompare(secondary) {}

        inline int operator()(const Key &a, const Key &b) const
        {
                int result = fPrimaryKeyCompare(a.primary, b.primary);
                if (result == 0 && a.use_secondary && b.use_secondary)
                        result = fSecondaryKeyCompare(a.secondary, b.secondary);
                return result;
        }

private:
        PrimaryKeyCompare       fPrimaryKeyCompare;
        SecondaryKeyCompare     fSecondaryKeyCompare;
};

// TwoKeyAVLTreeGetKey
template<typename Value, typename PrimaryKey, typename SecondaryKey,
                 typename GetPrimaryKey, typename GetSecondaryKey>
class TwoKeyAVLTreeGetKey
{
private:
        typedef TwoKeyAVLTreeKey<PrimaryKey, SecondaryKey> Key;

public:
        TwoKeyAVLTreeGetKey(const GetPrimaryKey &getPrimary,
                                                const GetSecondaryKey &getSecondary)
                : fGetPrimaryKey(getPrimary),
                  fGetSecondaryKey(getSecondary)
        {
        }

        inline Key operator()(const Value &a) const
        {
                return Key(fGetPrimaryKey(a), fGetSecondaryKey(a));
        }

private:
        GetPrimaryKey   fGetPrimaryKey;
        GetSecondaryKey fGetSecondaryKey;
};


// TwoKeyAVLTreeStandardCompare
template<typename Value>
class TwoKeyAVLTreeStandardCompare
{
public:
        inline int operator()(const Value &a, const Value &b) const
        {
                if (a < b)
                        return -1;
                else if (a > b)
                        return 1;
                return 0;
        }
};


// TwoKeyAVLTreeStandardGetKey
template<typename Value, typename Key>
class TwoKeyAVLTreeStandardGetKey
{
public:
        inline const Key &operator()(const Value &a) const
        {
                return a;
        }

        inline Key &operator()(Value &a) const
        {
                return a;
        }
};


// TwoKeyAVLTreeNodeStrategy
template <typename PrimaryKey, typename SecondaryKey, typename Value,
        typename PrimaryKeyCompare, typename SecondaryKeyCompare,
        typename GetPrimaryKey, typename GetSecondaryKey>
class TwoKeyAVLTreeNodeStrategy {
public:
        typedef TwoKeyAVLTreeKey<PrimaryKey, SecondaryKey> Key;

        TwoKeyAVLTreeNodeStrategy(
                const PrimaryKeyCompare& primaryKeyCompare = PrimaryKeyCompare(),
                const SecondaryKeyCompare& secondaryKeyCompare = SecondaryKeyCompare(),
                const GetPrimaryKey& getPrimaryKey = GetPrimaryKey(),
                const GetSecondaryKey& getSecondaryKey = GetSecondaryKey())
                : fPrimaryKeyCompare(primaryKeyCompare),
                  fSecondaryKeyCompare(secondaryKeyCompare),
                  fGetPrimaryKey(getPrimaryKey),
                  fGetSecondaryKey(getSecondaryKey)
        {
        }

        struct Node : AVLTreeNode {
                Node(const Value &value)
                        : AVLTreeNode(),
                          value(value)
                {
                }

                Value   value;
        };

        inline Node* Allocate(const Key& key, const Value& value) const
        {
                return new(nothrow) Node(value);
        }

        inline void Free(Node* node) const
        {
                delete node;
        }

        // internal use (not part of the strategy)
        inline Key GetValueKey(const Value& value) const
        {
                return Key(fGetPrimaryKey(value), fGetSecondaryKey(value));
        }

        inline Key GetKey(Node* node) const
        {
                return GetValueKey(node->value);
        }

        inline Value& GetValue(Node* node) const
        {
                return node->value;
        }

        inline AVLTreeNode* GetAVLTreeNode(Node* node) const
        {
                return node;
        }

        inline Node* GetNode(AVLTreeNode* node) const
        {
                return static_cast<Node*>(node);
        }

        inline int CompareKeyNode(const Key& a, const Node* b) const
        {
                return _CompareKeys(a, GetKey(const_cast<Node*>(b)));
        }

        inline int CompareNodes(const Node* a, const Node* b) const
        {
                return _CompareKeys(GetKey(const_cast<Node*>(a)),
                        GetKey(const_cast<Node*>(b)));
        }

private:
        inline int _CompareKeys(const Key& a, const Key& b) const
        {
                int result = fPrimaryKeyCompare(a.primary, b.primary);
                if (result == 0 && a.use_secondary && b.use_secondary)
                        result = fSecondaryKeyCompare(a.secondary, b.secondary);
                return result;
        }

        PrimaryKeyCompare       fPrimaryKeyCompare;
        SecondaryKeyCompare     fSecondaryKeyCompare;
        GetPrimaryKey           fGetPrimaryKey;
        GetSecondaryKey         fGetSecondaryKey;
};


// for convenience
#define TWO_KEY_AVL_TREE_TEMPLATE_LIST template<typename Value, \
        typename PrimaryKey, typename PrimaryKeyCompare, \
        typename GetPrimaryKey, typename SecondaryKey, \
        typename SecondaryKeyCompare, typename GetSecondaryKey>
#define TWO_KEY_AVL_TREE_CLASS_NAME TwoKeyAVLTree<Value, PrimaryKey, \
        PrimaryKeyCompare, GetPrimaryKey, SecondaryKey, \
        SecondaryKeyCompare, GetSecondaryKey>


// TwoKeyAVLTree
template<typename Value, typename PrimaryKey,
                 typename PrimaryKeyCompare, typename GetPrimaryKey,
                 typename SecondaryKey = Value,
                 typename SecondaryKeyCompare
                        = TwoKeyAVLTreeStandardCompare<SecondaryKey>,
                 typename GetSecondaryKey
                        = TwoKeyAVLTreeStandardGetKey<Value, SecondaryKey> >
class TwoKeyAVLTree {
public:
        typedef TwoKeyAVLTreeKey<PrimaryKey, SecondaryKey> Key;
        typedef TwoKeyAVLTreeNodeStrategy<PrimaryKey, SecondaryKey, Value,
                PrimaryKeyCompare, SecondaryKeyCompare, GetPrimaryKey,
                GetSecondaryKey> NodeStrategy;
        typedef AVLTreeMap<Key, Value, NodeStrategy> TreeMap;

        typedef typename TreeMap::Iterator      TreeMapIterator;
        typedef typename NodeStrategy::Node Node;

        class Iterator;

        TwoKeyAVLTree();
        TwoKeyAVLTree(const PrimaryKeyCompare &primaryCompare,
                                  const GetPrimaryKey &getPrimary,
                                  const SecondaryKeyCompare &secondaryCompare,
                                  const GetSecondaryKey &getSecondary);
        ~TwoKeyAVLTree();

        inline int CountItems() const   { return fTreeMap.Count(); }

        Value *FindFirst(const PrimaryKey &key, Iterator *iterator = NULL);
        Value *FindLast(const PrimaryKey &key, Iterator *iterator = NULL);
        inline Value *Find(const PrimaryKey &primaryKey,
                                           const SecondaryKey &secondaryKey,
                                           Iterator *iterator = NULL);

        inline void GetIterator(Iterator *iterator);

        inline status_t Insert(const Value &value, Iterator *iterator = NULL);
        inline status_t Remove(const PrimaryKey &primaryKey,
                                                   const SecondaryKey &secondaryKey);

private:
        TreeMap                         fTreeMap;
        PrimaryKeyCompare       fPrimaryKeyCompare;
        GetPrimaryKey           fGetPrimaryKey;
};


// Iterator
TWO_KEY_AVL_TREE_TEMPLATE_LIST
class TWO_KEY_AVL_TREE_CLASS_NAME::Iterator {
public:
        typedef typename TWO_KEY_AVL_TREE_CLASS_NAME::TreeMapIterator
                TreeMapIterator;

        inline Iterator()
                : fIterator()
        {
        }

        inline ~Iterator()
        {
        }

        inline Value *GetCurrent()
        {
                return fIterator.CurrentValuePointer();
        }

        inline Value *GetNext()
        {
                fIterator.Next();
                return GetCurrent();
        }

        inline Value *GetPrevious()
        {
                fIterator.Previous();
                return GetCurrent();
        }

        inline void Remove()
        {
                fIterator.Remove();
        }

private:
        friend class TWO_KEY_AVL_TREE_CLASS_NAME;

        Iterator(const Iterator& other)
                : fIterator(other.fIterator)
        {
        }

        Iterator &operator=(const Iterator& other)
        {
                fIterator = other.fIterator;
        }

        Iterator(const TreeMapIterator &iterator)
                : fIterator()
        {
        }

        inline void _SetTo(const TreeMapIterator &iterator)
        {
                fIterator = iterator;
        }

private:
        TWO_KEY_AVL_TREE_CLASS_NAME::TreeMapIterator fIterator;
};


// constructor
TWO_KEY_AVL_TREE_TEMPLATE_LIST
TWO_KEY_AVL_TREE_CLASS_NAME::TwoKeyAVLTree()
        : fTreeMap(NodeStrategy(PrimaryKeyCompare(), SecondaryKeyCompare(),
                GetPrimaryKey(), GetSecondaryKey())),
          fPrimaryKeyCompare(PrimaryKeyCompare()),
          fGetPrimaryKey(GetPrimaryKey())
{
}


// constructor
TWO_KEY_AVL_TREE_TEMPLATE_LIST
TWO_KEY_AVL_TREE_CLASS_NAME::TwoKeyAVLTree(
        const PrimaryKeyCompare &primaryCompare, const GetPrimaryKey &getPrimary,
        const SecondaryKeyCompare &secondaryCompare,
        const GetSecondaryKey &getSecondary)
        : fTreeMap(NodeStrategy(primaryCompare, secondaryCompare, getPrimary,
                getSecondary)),
          fPrimaryKeyCompare(primaryCompare),
          fGetPrimaryKey(getPrimary)
{
}


// destructor
TWO_KEY_AVL_TREE_TEMPLATE_LIST
TWO_KEY_AVL_TREE_CLASS_NAME::~TwoKeyAVLTree()
{
}


// FindFirst
TWO_KEY_AVL_TREE_TEMPLATE_LIST
Value *
TWO_KEY_AVL_TREE_CLASS_NAME::FindFirst(const PrimaryKey &key,
        Iterator *iterator)
{
        const NodeStrategy& strategy = fTreeMap.GetNodeStrategy();
        Node *node = fTreeMap.RootNode();

        while (node) {
                int cmp = fPrimaryKeyCompare(key, fGetPrimaryKey(
                        strategy.GetValue(node)));
                if (cmp == 0) {
                        // found a matching node, now get the left-most node with that key
                        while (node->left && fPrimaryKeyCompare(key,
                                        fGetPrimaryKey(strategy.GetValue(
                                                strategy.GetNode(node->left)))) == 0) {
                                node = strategy.GetNode(node->left);
                        }
                        if (iterator)
                                iterator->_SetTo(fTreeMap.GetIterator(node));
                        return &strategy.GetValue(node);
                }

                if (cmp < 0)
                        node = strategy.GetNode(node->left);
                else
                        node = strategy.GetNode(node->right);
        }
        return NULL;
}

// FindLast
TWO_KEY_AVL_TREE_TEMPLATE_LIST
Value *
TWO_KEY_AVL_TREE_CLASS_NAME::FindLast(const PrimaryKey &key,
        Iterator *iterator)
{
        const NodeStrategy& strategy = fTreeMap.GetNodeStrategy();
        Node *node = fTreeMap.RootNode();

        while (node) {
                int cmp = fPrimaryKeyCompare(key, fGetPrimaryKey(
                        strategy.GetValue(node)));
                if (cmp == 0) {
                        // found a matching node, now get the right-most node with that key
                        while (node->right && fPrimaryKeyCompare(key,
                                        fGetPrimaryKey(strategy.GetValue(
                                                strategy.GetNode(node->right)))) == 0) {
                                node = strategy.GetNode(node->right);
                        }
                        if (iterator)
                                iterator->_SetTo(fTreeMap.GetIterator(node));
                        return &strategy.GetValue(node);
                }

                if (cmp < 0)
                        node = strategy.GetNode(node->left);
                else
                        node = strategy.GetNode(node->right);
        }
        return NULL;
}

// Find
TWO_KEY_AVL_TREE_TEMPLATE_LIST
Value *
TWO_KEY_AVL_TREE_CLASS_NAME::Find(const PrimaryKey &primaryKey,
        const SecondaryKey &secondaryKey, Iterator *iterator)
{
        
        TreeMapIterator it = fTreeMap.Find(Key(primaryKey, secondaryKey));
        if (iterator)
                iterator->_SetTo(it);
        return it.CurrentValuePointer();
}

// GetIterator
TWO_KEY_AVL_TREE_TEMPLATE_LIST
void
TWO_KEY_AVL_TREE_CLASS_NAME::GetIterator(Iterator *iterator)
{
        TreeMapIterator it = fTreeMap.GetIterator();
        it.Next();
                // Our iterator needs to point to the first entry already.
        iterator->_SetTo(it);
}

// Insert
TWO_KEY_AVL_TREE_TEMPLATE_LIST
status_t
TWO_KEY_AVL_TREE_CLASS_NAME::Insert(const Value &value, Iterator *iterator)
{
        NodeStrategy& strategy
                = const_cast<NodeStrategy&>(fTreeMap.GetNodeStrategy());

        TreeMapIterator it;
        status_t status = fTreeMap.Insert(strategy.GetValueKey(value), value, &it);
        if (status != B_OK || !iterator)
                return status;

        iterator->_SetTo(it);
        return B_OK;
}

// Remove
TWO_KEY_AVL_TREE_TEMPLATE_LIST
status_t
TWO_KEY_AVL_TREE_CLASS_NAME::Remove(const PrimaryKey &primaryKey,
        const SecondaryKey &secondaryKey)
{
        return fTreeMap.Remove(Key(primaryKey, secondaryKey));
}

#endif  // TWO_KEY_AVL_TREE_H