/*-
 * Copyright (c) 2000 Qube Software, Ltd.
 * All rights reserved.
 * 
 * 	$Id: atree.h,v 1.5 2000/06/16 13:20:23 hugh Exp $
 */

#ifndef __atree_h__
#define __atree_h__

#include        <iostream>
#include        <assert.h>
#include        "vector.h"

// #define TREE_DEBUG

/* Forward Decls */
template<class K, class V> struct AVLTree;
template<class K, class V> struct AVLTreeIter;

template<class K, class V> struct AVLTreeNode
{
    typedef AVLTreeNode<K,V>            nodeType;

    friend struct AVLTree<K, V>;
    friend struct AVLTreeIter<K, V>;

    // Constraints on K
    // int operator==(const K&) const;
    // int operator<=(const k&) const;
    // K::K();

    // Constraints on V
    // V::V();

    // Constructors
                                AVLTreeNode(const K& key, const V& value)
                                    : key_(key), value_(value) { init(); }

    // Destructor
                                ~AVLTreeNode();

    // Accessors
    const K&                    key() const { return key_; }
    const V&                    value() const { return value_; }
    int                         balance() const { return balance_; }
    nodeType*                   left() const { return left_; }
    nodeType*                   right() const { return right_; }

    // Modifiers
    void                        value(const V& v) { value_ = v; }

protected:

    void                        init();

private:

    nodeType*                   parent_;
    K                           key_;
    V                           value_;
    nodeType*                   left_;
    nodeType*                   right_;
    int                         balance_;
};


template<class K, class V> struct AVLTree
{
    typedef AVLTreeNode<K,V>       nodeType;
    typedef AVLTree<K,V>           thisType;
    typedef AVLTreeIter<K,V>       iterator;

    friend struct AVLTreeIter<K,V>;

    // Constructors
                                AVLTree() { top_ = 0; }

    // Destructor
                                ~AVLTree() { empty(); }

    // Accessors
    nodeType*                   top() const { return top_; }

    // Features
    unsigned int                size() const;
    bool                        insert(const K& key, const V& value);
    bool                        remove(const K&);
    void                        empty() { delete top_; top_ = 0; }
    iterator                    find(const K& key) const;
    void                        checkTree() const;
    int                         checkDepth(nodeType*) const;
    iterator                    intern(const K&);
    iterator                    operator[](const K& k) { return intern(k); }

private:

    nodeType*                   top_;
};

template <class K, class V> struct AVLTreeIter
{
    friend struct AVLTree<K,V>;
    typedef AVLTreeNode<K,V>       nodeType;

    // Constructors
                                AVLTreeIter() { init(); }
                                AVLTreeIter(const AVLTree<K,V>& t)
                                        { init(); tree(&t); reset(); }
                                AVLTreeIter(const AVLTreeIter<K,V>& ti) :
                                    tree_(ti.tree_)
                                      { init(); *this = ti; }

    // Accessors
    const AVLTree<K,V>*         tree() const { return tree_; }

    // Modifiers
    void                        tree(const AVLTree<K,V>* t)
                                  { tree_ = const_cast<AVLTree<K,V>*>(t); }

    // Operators
    AVLTreeIter<K,V>&           operator=(const AVLTreeIter<K,V>&);
    int                         operator==(const AVLTreeIter<K,V>& ti) const
                                    { return pos_ == ti.pos_; }

    // Features
    void                        reset() { top(); downLeft(); }
    void                        reset(const AVLTree<K,V>& t)
                                { init(); tree(&t); reset(); }

    void                        top();
    void                        bump(bool forwards = true);
    bool                        find(const K&);
    bool                        insert(const K&, const V&);
    nodeType*                   pos() const { return pos_; }
    bool                        valid() const { return pos_ != 0; }
    const V&                    value() const;
    V&                          value();
    const K&                    key() const;
    const V&                    operator*() const { return value(); }
    AVLTreeIter<K,V>&           operator++() { bump(); return *this; }
    AVLTreeIter<K,V>&           operator--() { bump(false); return *this; }
    bool                        up();
    bool                        adjustBalance();
    void                        rotateLeft();
    void                        rotateRight();
    bool                        left();
    bool                        right();
    void                        rightBalance();
    void                        leftBalance();
    bool                        remove();
    void                        downLeft();
    void                        downRight();
    const V&                    operator=(const V&);
                                operator V();

protected:

    void                        init();
    int                         here(nodeType*** = 0);

private:

    AVLTree<K, V>*              tree_;
    nodeType*                   pos_;
    nodeType**                  insertAddr_;
};

/*********** AVLTreeNode<K,V>, methods  **************************/

template <class K, class V> void AVLTreeNode<K, V>::init()
{
    parent_ = 0;
    left_ = 0;
    right_ = 0;
    balance_ = 0;
}

template<class K, class V> AVLTreeNode<K,V>::~AVLTreeNode()
{
    delete left_;
    delete right_;
}

/*********** AVLTree<K, V>, methods  ********************************/

template <class K, class V>
bool AVLTree<K, V>::insert(const K& key, const V& value)
{
    AVLTreeIter<K,V> ti;

    ti.tree(this);
    ti.top();

    if (!ti.find(key)) {

        /* Insert new pair */
        ti.insert(key, value);

        /* Now adjust balance */
        while (ti.adjustBalance() && ti.up()) ;

#ifdef TREE_DEBUG
        checkTree();
#endif
        return true;
    }
    return false;
}

template<class K, class V> bool AVLTree<K,V>::remove(const K& key)
{
    AVLTreeIter<K,V> ti;

    ti.tree(this);
    ti.top();

    bool v;
    if (ti.find(key)) {
        /* Node is not a subtree, remove it */
        v = ti.remove();
    }
    else v = false;

    return v;
}

template<class K, class V>
AVLTreeIter<K,V> AVLTree<K,V>::find(const K& key) const
{
    AVLTreeIter<K,V> ti;

    ti.tree(this);
    ti.top();
    if (!ti.find(key)) ti.pos_ = 0;
    return ti;
}

template<class K, class V>
AVLTreeIter<K,V> AVLTree<K,V>::intern(const K& key)
{
    AVLTreeIter<K,V> ti;

    ti.tree(this);
    ti.top();
    if (!ti.find(key)) {
        /* Make an entry for the given key */
        ti.insert(key, V());

        /* Now adjust balance */
        while (ti.adjustBalance() && ti.up()) ;

#ifdef TREE_DEBUG
        checkTree();
#endif

        /* Find again */
        ti.top();
        ti.find(key);
    }

    return ti;
}

template<class K, class V>
unsigned int AVLTree<K,V>::size() const
{
    unsigned int n = 0;
    for (iterator i(*this); i.valid(); ++i) {
        ++n;
    }
    return n;
}

#ifdef TREE_DEBUG
template<class K, class V> void AVLTree<K,V>::checkTree() const
{
    /* Check that all keys are in order */
    iterator i(*this);

    if (i.valid()) {
        K lastKey(i.key());
        while ((++i).valid()) {
            if (!(lastKey <= i.key())) {
                cout << "Tree out of order!" << endl;
                assert(0);
            }
            lastKey = i.key();
        }
    }

    checkDepth(top_);
}

template<class K, class V> int AVLTree<K,V>::checkDepth(nodeType* node) const
{
    if (node) {
        int dl = checkDepth(node->left_);
        int dr = checkDepth(node->right_);
        
        int bal = dr - dl;
        if (bal < -1 || bal > 1) {
            cout << "Tree out of balance, left = " << dl 
                 << " right = " << dr << endl;
            assert(0);
        }
        
        if (node->balance_ != bal) {
            cout << "Node balance not correct!" << endl;
            assert(0);
        }

        if (dr > dl) dl = dr;
        return dl + 1;
    }
    else {
        return 0;
    }
}

#endif /* TREE_DEBUG */

/*********** AVLTreeIter<K, V>, methods  ***************************/

template <class K, class V> void AVLTreeIter<K, V>::init()
{
    pos_ = 0;
    insertAddr_ = 0;
}

template<class K, class V>
AVLTreeIter<K,V>& AVLTreeIter<K,V>::operator=(const AVLTreeIter<K,V>& ti)
{
    tree_ = ti.tree_;
    pos_ = ti.pos_;
    insertAddr_ = ti.insertAddr_;
    return *this;
}

template <class K, class V> void AVLTreeIter<K, V>::top()
{
    /* Set position to top of tree */
    assert(tree_);
    pos_ = tree_->top_;
    insertAddr_ = &tree_->top_;
}

template <class K, class V> const K& AVLTreeIter<K,V>::key() const
{
    assert(pos_);
    return pos_->key_;
}

template <class K, class V> const V& AVLTreeIter<K,V>::value() const
{
    assert(pos_);
    return pos_->value_;
}

template <class K, class V> V& AVLTreeIter<K,V>::value()
{
    assert(pos_);
    return pos_->value_;
}

template <class K, class V> bool AVLTreeIter<K, V>::find(const K& key)
{
    /* Search down from the current position for a match on the
     * given `key'.
     * If found, return true and leave the cursor position on the match.
     *
     */

    bool result = false;

    nodeType* p = pos_;
    while (p) {
        if (p->key_ == key) {
            result = true;
            break;
        }
        else {
            if (key <= p->key_) {       // actually <
                insertAddr_ = &pos_->left_;
                p = pos_->left_;
                if (p) pos_ = p;
            }
            else {                          // otherwise >
                insertAddr_ = &pos_->right_;
                p = pos_->right_;
                if (p) pos_ = p;
            }
        }
    }
    return result;
}

template <class K, class V>
bool AVLTreeIter<K, V>::insert(const K& key, const V& value)
{
    if (insertAddr_) {
        /* `find' has left us with `pos' on the parent and
         * `insertAddr' the child pointer of pos where the insert
         * should be performed.
         */
        AVLTreeNode<K,V>* node = new AVLTreeNode<K,V>(key, value);
        node->parent_ = pos_;
        *insertAddr_ = node;
    }

    return true;
}

template <class K, class V> void AVLTreeIter<K, V>::downRight()
{
    if (pos_) {
        while (pos_->right_) pos_ = pos_->right_;
    }
}

template <class K, class V> void AVLTreeIter<K, V>::downLeft()
{
    if (pos_) {
        while (pos_->left_) pos_ = pos_->left_;
    }
}

template <class K, class V> void AVLTreeIter<K, V>::bump(bool forwards)
{
    /* Move `forwards' or backwards in tree order.
     */
    if (pos_) {
        if (forwards && pos_->right_) {
            pos_ = pos_->right_;
            downLeft();

        }
        else if (!forwards && pos_->left_) {
            pos_ = pos_->left_;
            downRight();
        }
        else {
            /* Have no child in our direction, move up to parent */
            nodeType* p = pos_;
            pos_ = pos_->parent_;
            while (pos_) {
                if (forwards) {
                    if (pos_->left_ == p) {
                        /* We were the left branch, stay. */
                        break;
                    }
                }
                else if (!forwards) {
                    if (pos_->right_ == p) {
                        /* We were the right branch and want to go left,
                         * stay.
                         */
                        break;
                    }
                }
                p = pos_;
                pos_ = pos_->parent_;
            }
        }
    }
}

template <class K, class V> bool AVLTreeIter<K, V>::up()
{
    if (pos_) {
        here();
        pos_ = pos_->parent_;
        return true;
    }
    return false;
}

template <class K, class V>
int AVLTreeIter<K, V>::here(nodeType*** sibAddr)
{
    /* Result is:
     *  1 if we were right child.
     * -1 if left child.
     * 0  if no parent (we are root).
     */
    int result = 0;
    if (pos_) {
        nodeType* p = pos_->parent_;
        if (p) {
            if (p->left_ == pos_) {
                insertAddr_ = &p->left_;
                result = -1;
                if (sibAddr) *sibAddr = &p->right_;
            }
            else {
                insertAddr_ = &p->right_;
                result = 1;
                if (sibAddr) *sibAddr = &p->left_;
            }
        }
        else {
            insertAddr_ = &tree_->top_;
        }
    }
    return result;
}

template <class K, class V> bool AVLTreeIter<K, V>::left()
{
    bool result = false;
    if (pos_) {
        insertAddr_ = &pos_->left_;
        pos_ = pos_->left_;
        result = true;
    }
    return result;
}

template <class K, class V> bool AVLTreeIter<K, V>::right()
{
    bool result = false;
    if (pos_) {
        insertAddr_ = &pos_->right_;
        pos_ = pos_->right_;
        result = true;
    }
    return result;
}

template <class K, class V> bool AVLTreeIter<K, V>::adjustBalance()
{
    bool taller = false;

    if (pos_) {
        if (insertAddr_ == &pos_->left_) {

            /* If we've just inserted on the left .. */
            if (!pos_->balance_) taller = true;
            --pos_->balance_;

            if (pos_->balance_ < -1) leftBalance();
        }
        else if (insertAddr_ == &pos_->right_) {

            if (!pos_->balance_) taller = true;
            ++pos_->balance_;
            if (pos_->balance_ > 1) rightBalance();
        }
    }
    return taller;
}

template <class K, class V> void AVLTreeIter<K, V>::rightBalance()
{
    assert(insertAddr_ && *insertAddr_);
    assert((*insertAddr_)->balance_ != 0);
    assert(pos_);

    nodeType* insNode = *insertAddr_;
    if (insNode->balance_ > 0) {
        
        /* Reset the balances that will be the case _after_ the
         * rotate left.
         */
        pos_->balance_ = 0;
        insNode->balance_ = 0;

        /* Move up and rotate */
        up();
        rotateLeft();
    }
    else { // < 0

        nodeType* leftNode = insNode->left_;
        assert(leftNode);

        /* Adjust all the balances to be right after the double
         * rotation.
         */
        if (!leftNode->balance_) {
            pos_->balance_ = 0;
            insNode->balance_ = 0;
        }
        else if (leftNode->balance_ < 0) {
            pos_->balance_ = 0;
            insNode->balance_ = 1;
        }
        else {
            pos_->balance_ = -1;
            insNode->balance_ = 0;
        }
        leftNode->balance_ = 0;

        /* Double rotation */
        rotateRight();
        up();
        rotateLeft();
    }
}

template <class K, class V> void AVLTreeIter<K, V>::leftBalance()
{
    assert(insertAddr_ && *insertAddr_);
    assert((*insertAddr_)->balance_ != 0);
    assert(pos_);

    nodeType* insNode = *insertAddr_;
    if (insNode->balance_ < 0) {

        pos_->balance_ = 0;
        insNode->balance_ = 0;

        up();
        rotateRight();
    }
    else { // > 0

        nodeType* rightNode = insNode->right_;
        assert(rightNode);

        /* Adjust all the balances to be right after the double
         * rotation.
         */
        if (!rightNode->balance_) {
            pos_->balance_ = 0;
            insNode->balance_ = 0;
        }
        else if (rightNode->balance_ > 0) {
            pos_->balance_ = 0;
            insNode->balance_ = -1;
        }
        else {
            pos_->balance_ = 1;
            insNode->balance_ = 0;
        }
        rightNode->balance_ = 0;

        /* Perform double rotation */
        rotateLeft();
        up();
        rotateRight();
    }
}

template <class K, class V> void AVLTreeIter<K, V>::rotateLeft()
{
    assert(insertAddr_ && *insertAddr_);
    
    nodeType* insNode = *insertAddr_;
    nodeType* rChild = insNode->right_;
    assert(rChild);

    insNode->right_ = rChild->left_;
    if (rChild->left_) rChild->left_->parent_ = insNode;

    rChild->left_ = insNode;
    insNode->parent_ = rChild;

    *insertAddr_ = rChild;
    rChild->parent_ = pos_;
}

template <class K, class V> void AVLTreeIter<K, V>::rotateRight()
{
    assert(insertAddr_ && *insertAddr_);

    nodeType* insNode = *insertAddr_;
    nodeType* lChild = insNode->left_;
    assert(lChild);

    insNode->left_ = lChild->right_;
    if (lChild->right_) lChild->right_->parent_ = insNode;

    lChild->right_ = insNode;
    insNode->parent_ = lChild;

    *insertAddr_ = lChild;
    lChild->parent_ = pos_;
}

template <class K, class V>
bool AVLTreeIter<K,V>::remove()
{
    /* Remove node at `pos' from tree.
     */

    bool result = false;

    if (pos_) {
        if (pos_->left_ && pos_->right_) {
            /* We have two children  */
            AVLTreeIter<K,V> ti(*this);

            /* Find the preceding node (in logical order)
             * and replace this node with the content of the latter.
             * Continue by removing the preceding node.
             * The preceding node will be a leaf and not have two 
             * children. 
             */
            
            --ti;
            nodeType* pNode = ti.pos();
            assert(pNode);

            /* Copy the guts of preceding node into `pos'
             * (the one to be deleted).
             */
            pos_->key_ = pNode->key_;
            pos_->value_= pNode->value_;
            
            assert(ti.pos_ &&
                   !((ti.pos_->left_ != 0) && (ti.pos_->right_ != 0)));
            result = ti.remove();
        }
        else {
            /* We have at most one child */
            int balAdjust;
            nodeType**      sibling = 0;

            /* Find out whether we are left or right child of parent
             * and set up `insertAddr' to `pos' from parent.
             * Set `sibling' to the other branch.
             */
            balAdjust = -here(&sibling);
            
            /* On the basis that we have at most one child.
             * Change the parent's pointer from us to that child 
             * (or null).
             */
            if (pos_->left_) {
                pos_->left_->parent_ = pos_->parent_;
                *insertAddr_ = pos_->left_;
                pos_->left_ = 0;
            }
            else if (pos_->right_) {
                pos_->right_->parent_ = pos_->parent_;
                *insertAddr_ = pos_->right_;
                pos_->right_ = 0;
            }
            else *insertAddr_ = 0;

            nodeType* n = pos_;
            result = true;
            
            /* Move up to parent */
            pos_ =  pos_->parent_;

            /* Erase node */
            delete n;

            while (pos_) {
                if (pos_->balance_ == 0) {
                    /* Parent was equally balanced, not anymore! */
                    pos_->balance_ = balAdjust;
                    break;
                }
                else if (pos_->balance_ > 0 && balAdjust < 0 ||
                         pos_->balance_ < 0 && balAdjust > 0) {
                    /* Shortened the taller subtree */
                    pos_->balance_ = 0;

                    /* Move up */
                    balAdjust = -here(&sibling);
                    up();
                }
                else {
                    /* Shorter subtree was shortened.
                     * `sibling' is the taller tree. 
                     * `pos' is overbalanced, some rotations are needed
                     */
                    nodeType* sibNode = *sibling;
                    assert(sibNode);
                    if (sibNode->balance_ == 0) {
                        if (balAdjust > 0) {
                            sibNode->balance_ = -1;
                            up();
                            rotateLeft();
                        }
                        else {
                            sibNode->balance_ = 1;
                            up();
                            rotateRight();
                        }
                        break;
                    }
                    else if (sibNode->balance_ == pos_->balance_) {
                        pos_->balance_ = 0;
                        sibNode->balance_ = 0;
                        int ba = -here(&sibling);
                        up();
                        if (balAdjust > 0) rotateLeft();
                        else rotateRight();

                        /* Move up */
                        balAdjust = ba;

                    }
                    else {
                        /* Balance factors are opposite */
                        if (balAdjust > 0) {
                            insertAddr_ = sibling;
                            rotateRight();
                            
                            /* The node rotated under insertaddr will
                             * be the root of this subtree after the
                             * final rotation.
                             * It will have been balanced.
                             */
                            nodeType* node = *insertAddr_;
                            assert(node);
                            int oldBal = node->balance_;
                            node->balance_ = 0;

                            /* Fixup sibling node weight */
                            if (oldBal >= 0) {
                                sibNode->balance_ = 0;
                            }
                            else sibNode->balance_ = 1;

                            balAdjust = -here(&sibling);
                            up();

                            /* Now deal with original top node */
                            node = *insertAddr_;
                            assert(node);
                            if (oldBal <= 0) {
                                node->balance_ = 0;
                            }
                            else node->balance_ = -1;
                            
                            rotateLeft();
                        }
                        else {
                            /* Mirror image case */

                            insertAddr_ = sibling;
                            rotateLeft();
                            nodeType* node = *insertAddr_;
                            assert(node);
                            int oldBal = node->balance_;
                            node->balance_ = 0;

                            /* Fixup sibling node weight */
                            if (oldBal <= 0) {
                                sibNode->balance_ = 0;
                            }
                            else sibNode->balance_ = -1;

                            balAdjust = -here(&sibling);
                            up();

                            node = *insertAddr_;
                            assert(node);
                            if (oldBal >= 0) {
                                node->balance_ = 0;
                            }
                            else node->balance_ = 1;
                            
                            rotateRight();
                        }
                    }
                }
            }
        }
    }

#ifdef TREE_DEBUG
    if (result) tree_->checkTree();
#endif

    return result;
}

template<class K, class V>
const V& AVLTreeIter<K,V>::operator=(const V& val)
{
    if (pos_) {
        pos_->value(val);
    }
    return val;
}

template<class K, class V>
AVLTreeIter<K,V>::operator V() 
{
    if (pos_) {
        return pos_->value();
    }
    return V();
}

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

template<class K, class V>
std::ostream& operator<<(std::ostream& os, const AVLTreeNode<K,V>* node)
{
    os << '(';

    if (node) {

#ifdef TREE_DEBUG
        os << "balance:" << node->balance() << ' ';
#endif /* TREE_DEBUG */

        os << node->key();
        os << ':';
        os << node->value() << ' ';
        os << node->left() << ' ' << node->right();
    }
    os << ')';

    return os;
}

template<class K, class V>
std::ostream& operator<<(std::ostream& os, const AVLTree<K,V>& tree)
{
    return os << tree.top();
}

#endif /* __atree_h__ */


