//-----------------------------------------------------------------------------
// <copyright>
//   Copyright (c) 2009 Nesterovsky Bros. All rights reserved.
// </copyright>
// <description>
//   A bare red-black tree imlementation.
// </description>
//-----------------------------------------------------------------------------

namespace NesterovskyBros.Collections
{
  using System;

  /// <summary>
  /// A binary Red-Black tree node.
  /// 
  /// A set of linked nodes along with a container node constitues a 
  /// red-black tree.
  /// 
  /// This implementation stores tree's root under the container node.
  /// Program using this tree may store a container node rather than 
  /// root node, as root of the tree can be changed during
  /// tree balancing, while container stays firm.
  /// 
  /// Any node with no parent is considered to be a container.
  /// 
  /// Invariants:
  ///   root.Parent == container;
  ///   container.left == root;
  ///   container.right == root;
  /// </summary>
  public class RedBlackNode<Node>
    where Node: RedBlackNode<Node>
  {
    /// <summary>
    /// A parent node.
    /// </summary>
    public Node Parent
    {
      get { return parent; }
    }

    /// <summary>
    /// A left child.
    /// </summary>
    public Node Left
    {
      get { return left; }
    }

    /// <summary>
    /// A right child.
    /// </summary>
    public Node Right
    {
      get { return right; }
    }

    /// <summary>
    /// Indicates whether the current node is a container.
    /// </summary>
    public bool IsContainer
    {
      get { return parent == null; }
    }

    /// <summary>
    /// Indicates whether the current node is a tree root.
    /// </summary>
    public bool IsRoot
    {
      get { return (parent != null) && (parent.parent == null); }
    }

    /// <summary>
    /// Gets a tree container.
    /// </summary>
    /// <param name="node">A node value.</param>
    /// <returns>A tree container</returns>
    public static Node Container(Node node)
    {
      if (node != null)
      {
        while(node.parent != null)
        {
          node = node.parent;
        }
      }

      return node;
    }

    /// <summary>
    /// Returns the leftmost node of the subtree rooted by a node.
    /// </summary>
    /// <param name="node">a node value</param>
    /// <returns>the leftmost node of the subtree rooted by a node.</returns>
    public static Node Leftmost(Node node)
    {
      if (node != null)
      {
        while(node.left != null)
        {
          node = node.left;
        }
      }

      return node;
    }

    /// <summary>
    /// Returns the rightmost node of the subtree rooted by this node.
    /// </summary>
    /// <param name="node">a node value</param>
    /// <returns>the rightmost node of the subtree rooted by this node.</returns>
    public static Node Rightmost(Node node)
    {
      if (node != null)
      {
        while(node.right != null)
        {
          node = node.right;
        }
      }

      return node;
    }

    /// <summary>
    /// Returns the predecessor of the specified node, or null if no such.
    /// </summary>
    /// <param name="node">a node value.</param>
    /// <returns>
    /// the predecessor of the specified node, or null if no such.
    /// </returns>
    public static Node Predecessor(Node node)
    {
      if (node != null)
      {
        if (node.left != null)
        {
          node = node.left;

          while(node.right != null)
          {
            node = node.right;
          }
        }
        else
        {
          while((node.parent != null) && (node.parent.left != node))
          {
            node = node.parent;
          }
        }
      }

      return node;
    }

    /// <summary>
    /// Return the successor of the specified node, or null if no such.
    /// </summary>
    /// <param name="node">a node value.</param>
    /// <returns>
    /// the successor of the specified node, or null if no such.
    /// </returns>
    public static Node Successor(Node node)
    {
      if (node != null)
      {
        if (node.right != null)
        {
          node = node.right;

          while(node.left != null)
          {
            node = node.left;
          }
        }
        else
        {
          Node child;

          do
          {
            child = node;
            node = node.parent;
          }
          while((node != null) && (child == node.right));
        }
      }

      return node;
    }

    /// <summary>
    /// Insert a node into the tree before a reference node.
    /// </summary>
    /// <param name="reference">a reference node.</param>
    /// <param name="node">a node to add.</param>
    public static void InsertBefore(Node reference, Node node)
    {
      if (reference == null)
      {
        throw new ArgumentException("reference");
      }

      if ((node == null) ||
        (node.parent != null) ||
        (node.left != null) ||
        (node.right != null))
      {
        throw new ArgumentException("node");
      }

      if (reference.left != null)
      {
        reference = Predecessor(reference);
        reference.right = node;
      }
      else
      {
        reference.left = node;

        if (reference.IsContainer)
        {
          // Link both sides in container.
          reference.right = node;
        }
      }

      node.RebindParent(reference);
      node.parent = reference;
      FixAfterInsert(node);
    }

    /// <summary>
    /// Inserts a node into the tree after a reference node.
    /// </summary>
    /// <param name="reference">a reference node.</param>
    /// <param name="node">a node to add.</param>
    public static void InsertAfter(Node reference, Node node)
    {
      if (reference == null)
      {
        throw new ArgumentException("reference");
      }

      if ((node == null) ||
        (node.parent != null) ||
        (node.left != null) ||
        (node.right != null))
      {
        throw new ArgumentException("node");
      }

      if (reference.right != null)
      {
        reference = Successor(reference);
        reference.left = node;
      }
      else
      {
        reference.right = node;

        if (reference.IsContainer)
        {
          // Link both sides in container.
          reference.left = node;
        }
      }

      node.RebindParent(reference);
      node.parent = reference;
      FixAfterInsert(node);
    }

    /// <summary>
    /// Removes this node from the tree.
    /// </summary>
    /// <param name="node">a node to remove.</param>
    public static void Remove(Node node)
    {
      if ((node == null) || node.IsContainer)
      {
        throw new ArgumentException("node");
      }

      DeleteNode(node);
    }

    /// <summary>
    /// Rebinds node's parent.
    /// RebindParent is called before parent is changed.
    /// Override this method to adjust node before parent is changed.
    /// </summary>
    /// <param name="from">parent a new parent node.</param>
    protected virtual void RebindParent(Node parent)
    {
    }

    /// <summary>
    /// Gets node "color".
    /// </summary>
    /// <param name="node">a node</param>
    /// <returns>a node color or a value "black" is value is null.</returns>
    private static bool ColorOf(Node node)
    {
      return node == null ? black : node.color;
    }

    /// <summary>
    /// Returns a parent for node.
    /// </summary>
    /// <param name="node">a node to get parent for.</param>
    /// <returns>a parent for node or null if value is null.</returns>
    private static Node ParentOf(Node node)
    {
      return (node == null) || node.IsRoot ? null : node.parent;
    }

    /// <summary>
    /// Returns a left node.
    /// </summary>
    /// <param name="node">a node to get left node for.</param>
    /// <returns>a left node or null if value is null.</returns>
    private static Node LeftOf(Node node)
    {
      return node == null ? null : node.left;
    }

    /// <summary>
    /// Returns a right node.
    /// </summary>
    /// <param name="node">a node to get right node for.</param>
    /// <returns>a right node or null if value is null.</returns>
    private static Node RightOf(Node node)
    {
      return node == null ? null : node.right;
    }

    /// <summary>
    /// Sets a color for a node. Does nothing if node is null.
    /// </summary>
    /// <param name="node">a node to set color for.</param>
    /// <param name="color">a node color.</param>
    private static void SetColor(Node node, bool color)
    {
      if (node != null)
      {
        node.color = color;
      }
    }

    /// <summary>
    /// Returns second value if first is null, otherwise returns first value.
    /// </summary>
    /// <param name="first">fist value.</param>
    /// <param name="second">second value.</param>
    /// <returns>
    /// second value is first is null, otherwise returns first value.
    /// </returns>
    private static Node Coalesce(Node first, Node second)
    {
      return first == null ? second : first;
    }

    /// <summary>
    /// "Rotates" node to the left in the tree.
    /// </summary>
    /// <param name="node">a node.</param>
    /// <returns>A root node, if it's changed, or null otherwise.</returns>
    private static Node RotateLeft(Node node)
    {
      if (node == null)
      {
        return node;
      }

      Node parent = node.parent;
      Node left = node.left;
      Node right = node.right;
      Node rightleft = right.left;

      if (rightleft != null)
      {
        rightleft.RebindParent(node);
        rightleft.parent = node;
      }

      right.RebindParent(parent);
      right.parent = parent;
      right.left = node;

      node.RebindParent(right);
      node.parent = right;
      node.right = rightleft;

      if (!parent.IsContainer)
      {
        if (parent.left == node)
        {
          parent.left = right;
        }
        else
        {
          parent.right = right;
        }

        return null;
      }
      else
      {
        // Rebind root.
        parent.left = right;
        parent.right = right;

        return right;
      }
    }

    /// <summary>
    /// "Rotates" node to the right in the tree.
    /// </summary>
    /// <param name="node">a node.</param>
    /// <returns>A root node, if it's changed, or null otherwise.</returns>
    private static Node RotateRight(Node node)
    {
      if (node == null)
      {
        return node;
      }

      Node parent = node.parent;
      Node right = node.right;
      Node left = node.left;
      Node leftright = left.right;

      if (leftright != null)
      {
        leftright.RebindParent(node);
        leftright.parent = node;
      }

      left.RebindParent(parent);
      left.parent = parent;
      left.right = node;

      node.RebindParent(left);
      node.parent = left;
      node.left = leftright;

      if (!parent.IsContainer)
      {
        if (parent.left == node)
        {
          parent.left = left;
        }
        else
        {
          parent.right = left;
        }

        return null;
      }
      else
      {
        // Rebind root.
        parent.left = left;
        parent.right = left;

        return left;
      }
    }

    /// <summary>
    /// Fixes tree after insertion.
    /// </summary>
    /// <param name="node">a node value.</param>
    private static void FixAfterInsert(Node node)
    {
      if (node.IsRoot)
      {
        return;
      }

      Node root = null;

      node.color = red;

      while((node != null) && !node.IsRoot && (node.parent.color == red))
      {
        Node parent = node.parent;
        Node grand = ParentOf(parent);

        if (parent == LeftOf(grand))
        {
          Node uncle = RightOf(grand);

          if (ColorOf(uncle) == red)
          {
            SetColor(parent, black);
            SetColor(uncle, black);
            SetColor(grand, red);
            node = grand;

            if (node.IsRoot)
            {
              root = node;
            }
          }
          else
          {
            if (node == RightOf(parent))
            {
              node = parent;
              root = Coalesce(RotateLeft(node), root);
              parent = node.parent;
              grand = ParentOf(parent);
            }

            SetColor(parent, black);
            SetColor(grand, red);
            root = Coalesce(RotateRight(grand), root);
          }
        }
        else
        {
          Node uncle = LeftOf(grand);

          if (ColorOf(uncle) == red)
          {
            SetColor(parent, black);
            SetColor(uncle, black);
            SetColor(grand, red);
            node = grand;

            if (node.IsRoot)
            {
              root = node;
            }
          }
          else
          {
            if (node == LeftOf(parent))
            {
              node = parent;
              root = Coalesce(RotateRight(node), root);
              parent = node.parent;
              grand = ParentOf(parent);
            }

            SetColor(parent, black);
            SetColor(grand, red);
            root = Coalesce(RotateLeft(grand), root);
          }
        }
      }

      SetColor(root, black);
    }

    /// <summary>
    /// Fixes tree after deletion.
    /// </summary>
    /// <param name="node">a node value.</param>
    private static void FixAfterDelete(Node node)
    {
      Node root = null;

      while((node != null) && !node.IsRoot && (node.color == black))
      {
        Node parent = node.parent;

        if (node == LeftOf(parent))
        {
          Node sibling = RightOf(parent);

          if (ColorOf(sibling) == red)
          {
            SetColor(sibling, black);
            SetColor(parent, red);
            root = Coalesce(RotateLeft(parent), root);
            sibling = RightOf(parent);
          }

          if ((ColorOf(LeftOf(sibling)) == black) &&
            (ColorOf(RightOf(sibling)) == black))
          {
            SetColor(sibling, red);
            node = node.parent;

            if (node.IsRoot)
            {
              root = node;
            }
          }
          else
          {
            if (ColorOf(RightOf(sibling)) == black)
            {
              SetColor(LeftOf(sibling), black);
              SetColor(sibling, red);
              root = Coalesce(RotateRight(sibling), root);
              sibling = RightOf(parent);
            }

            SetColor(sibling, ColorOf(parent));
            SetColor(parent, black);
            SetColor(RightOf(sibling), black);
            root = Coalesce(RotateLeft(parent), root);
            node = root;
          }
        }
        else
        {
          // symmetric
          Node sibling = LeftOf(parent);

          if (ColorOf(sibling) == red)
          {
            SetColor(sibling, black);
            SetColor(ParentOf(node), red);
            root = Coalesce(RotateRight(ParentOf(node)), root);
            sibling = LeftOf(ParentOf(node));
          }

          if ((ColorOf(RightOf(sibling)) == black) &&
            (ColorOf(LeftOf(sibling)) == black))
          {
            SetColor(sibling, red);
            node = node.parent;

            if (node.IsRoot)
            {
              root = node;
            }
          }
          else
          {
            if (ColorOf(LeftOf(sibling)) == black)
            {
              SetColor(RightOf(sibling), black);
              SetColor(sibling, red);
              root = Coalesce(RotateLeft(sibling), root);
              sibling = LeftOf(ParentOf(node));
            }

            SetColor(sibling, ColorOf(ParentOf(node)));
            SetColor(ParentOf(node), black);
            SetColor(LeftOf(sibling), black);
            root = Coalesce(RotateRight(ParentOf(node)), root);
            node = root;
          }
        }
      }

      SetColor(node, black);
    }

    /// <summary>
    /// Deletes a node from tree.
    /// </summary>
    /// <param name="erased">a node value.</param>
    private static void DeleteNode(Node erased)
    {
      Node node = erased;
      Node parent = erased.parent;
      Node fixNode; // the node to recolor as needed

      if (node.left == null)
      {
        // must stitch up right subtree
        fixNode = node.right;
      }
      else if (node.right == null)
      {
        // must stitch up left subtree
        fixNode = node.left;
      }
      else
      {
        // two subtrees, must lift successor node to replace erased
        node = Successor(erased);  // node is successor node
        fixNode = RightOf(node);  // fixNode is its only subtree
      }

      bool erasedColor;

      if (node == erased)
      {
        // at most one subtree, relink it
        if (fixNode != null)
        {
          // link up
          fixNode.RebindParent(parent);
          fixNode.parent = parent;
        }

        if (parent.IsContainer)
        {
          // link down from root
          parent.left = fixNode;
          parent.right = fixNode;
        }
        else
        {
          if (parent.left == erased)
          {
            parent.left = fixNode; // link down to left
          }
          else
          {
            parent.right = fixNode;  // link down to right
          }
        }

        erasedColor = erased.color;
      }
      else
      {
        // erased has two subtrees, node is successor to erased
        Node erasedleft = erased.left;
        Node erasedright = erased.right;

        // link left up
        erasedleft.RebindParent(node);
        erasedleft.parent = node;

        if (parent.IsContainer)
        {
          // link down from root
          parent.left = node;
          parent.right = node;
        }
        else
        {
          if (parent.left == erased)
          {
            parent.left = node; // link down to left
          }
          else
          {
            parent.right = node;  // link down to right
          }
        }

        node.RebindParent(parent); // link successor up
        node.parent = parent;
        node.left = erasedleft; // link successor down

        if (node != erasedright)
        {
          // successor further down, link in place of erased
          Node nodeparent = node.parent;

          if (fixNode != null)
          {
            // link fix up
            fixNode.RebindParent(nodeparent);
            fixNode.parent = nodeparent;
          }

          nodeparent.left = fixNode; // link fix down
          node.right = erasedright; // link successor down

          // link right up
          erasedright.RebindParent(node);
          erasedright.parent = node;
        }

        // recolor it
        erasedColor = ColorOf(node);
        SetColor(node, ColorOf(erased));
      }

      erased.RebindParent(null);
      erased.parent = null;
      erased.left = null;
      erased.right = null;
      erased.color = black;

      if (erasedColor == black)
      {
        FixAfterDelete(fixNode);
      }
    }

    /// <summary>
    /// A parent node.
    /// </summary>
    private Node parent;

    /// <summary>
    /// A left child node.
    /// </summary>
    private Node left;

    /// <summary>
    /// A right child node.
    /// </summary>
    private Node right;

    /// <summary>
    /// A "color" of the node.
    /// </summary>
    private bool color = black;

    /// <summary>
    /// Red color.
    /// </summary>
    private const bool red = true;

    /// <summary>
    /// Black color.
    /// </summary>
    private const bool black = false;
  }

  /// <summary>
  /// A binary tree node with node index.
  /// </summary>
  public class IndexedNode<Node>: RedBlackNode<Node>
    where Node: IndexedNode<Node>
  {
    /// <summary>
    /// Node index.
    /// </summary>
    public int Index
    {
      get
      {
        int index = 0;
        IndexedNode<Node> node = this;

        do
        {
          index += node.index;
          node = node.Parent;
        }
        while (node != null);

        return index;
      }
    }

    /// <summary>
    /// Gets number of nodes in the tree.
    /// </summary>
    public static int Count(Node container)
    {
      return Rightmost(container).Index;
    }

    /// <summary>
    /// Gets node with a specified index.
    /// Note that container has index equal to 0.
    /// </summary>
    /// <param name="container">a container node.</param>
    /// <param name="index">node index.</param>
    /// <returns>
    /// node value, or null if there is no node with a specified index.
    /// </returns>
    public static Node Get(Node container, int index)
    {
      int nodeIndex = 0;
      Node node = container;

      while(true)
      {
        if (index > nodeIndex)
        {
          node = node.Right;
        }
        else if (index < nodeIndex)
        {
          node = node.Left;
        }
        else
        {
          return node;
        }

        if (node == null)
        {
          return null;
        }

        nodeIndex += node.index;
      }
    }

    /// <summary>
    /// Insert a node into the tree before a reference node.
    /// </summary>
    /// <param name="reference">a reference node.</param>
    /// <param name="node">a node to add.</param>
    public static new void InsertBefore(Node reference, Node node)
    {
      RedBlackNode<Node>.InsertBefore(reference, node);
      OnIndexChanged(node, 1, false);
    }

    /// <summary>
    /// Inserts a node into the tree after a reference node.
    /// </summary>
    /// <param name="reference">a reference node.</param>
    /// <param name="node">a node to add.</param>
    public static new void InsertAfter(Node reference, Node node)
    {
      RedBlackNode<Node>.InsertAfter(reference, node);
      OnIndexChanged(node, 1, true);
    }

    /// <summary>
    /// Removes this node from the tree.
    /// </summary>
    /// <param name="node">a node to remove.</param>
    public static new void Remove(Node node)
    {
      OnIndexChanged(node, -1, true);
      RedBlackNode<Node>.Remove(node);
    }

    /// <summary>
    /// Rebinds node's parent.
    /// RebindParent is called before parent is changed.
    /// </summary>
    /// <param name="from">parent a new parent node.</param>
    protected override void RebindParent(Node parent)
    {
      if ((parent == null) || (Parent == null))
      {
        index = 0;
      }
      else
      {
        index = Index - parent.Index;
      }
    }

    /// <summary>
    /// Adjusts node index.
    /// </summary>
    /// <param name="reference">a reference node.</param>
    /// <param name="offset">change offset.</param>
    /// <param name="before">
    /// true when index is increased or decreased before this node, and
    /// false when after it.
    /// </param>
    private static void OnIndexChanged(Node reference, int offset, bool before)
    {
      Node node = reference;
      bool left;

      if (before)
      {
        left = node.Left == null;

        if (!left)
        {
          node = Predecessor(node);
        }
      }
      else
      {
        left = node.Right != null;

        if (left)
        {
          node = Successor(node);
        }
      }

      while(!node.IsContainer)
      {
        Node parent = node.Parent;
        bool leftNode = parent.Right != node;

        if (left != leftNode)
        {
          if (leftNode)
          {
            node.index -= offset;
          }
          else
          {
            node.index += offset;
          }
        }

        node = parent;
        left = leftNode;
      }
    }
    
    /// <summary>
    /// Index offset.
    /// </summary>
    private int index;
  }
}