Juliet

using System;
using System.Collections.Generic;

namespace Juliet.Collections.Immutable
{
    public abstract class BaseBinaryTree<TreeType, T>
        where TreeType : BaseBinaryTree<TreeType, T>
    {
        protected readonly Func<T, T, int> Comparer;

        public BaseBinaryTree(Func<T, T, int> comparer)
        {
            this.Comparer = comparer;
        }

        protected abstract TreeType Self();
        protected abstract U Match<U>(Func<U> nilFunc, Func<TreeType, T, TreeType, U> nodeFunc);
        protected abstract TreeType CreateNode(TreeType left, T value, TreeType right);

        public virtual TreeType Insert(T value)
        {
            return this.Match(
                () => CreateNode(Self(), value, Self()),
                (l, x, r) =>
                    {
                        int compare = this.Comparer(value, x);
                        if (compare < 0) { return CreateNode(l.Insert(value), x, r); }
                        else if (compare > 0) { return CreateNode(l, x, r.Insert(value)); }
                        else { return Self(); }
                    });
        }
    }

    public abstract class DerivedAvlTree<T> : BaseBinaryTree<DerivedAvlTree<T>, T>
    {
        #region Nodes
        private sealed class Nil : DerivedAvlTree<T>
        {
            public Nil(Func<T, T, int> comparer) : base(comparer) { }

            protected override U Match<U>(Func<U> nilFunc, Func<DerivedAvlTree<T>, T, DerivedAvlTree<T>, U> nodeFunc)
            {
                return nilFunc();
            }

            protected override int GetCount() { return 0; }
            protected override int GetHeight() { return 0; }
        }

        private sealed class Node : DerivedAvlTree<T>
        {
            DerivedAvlTree<T> left;
            T value;
            DerivedAvlTree<T> right;
            int count;
            int height;

            public Node(Func<T, T, int> comparer, DerivedAvlTree<T> left, T value, DerivedAvlTree<T> right)
                : base(comparer)
            {
                this.left = left;
                this.value = value;
                this.right = right;
                this.count = left.Count + right.Count + 1;
                this.height = Math.Max(left.Height, right.Height) + 1;
            }

            protected override U Match<U>(Func<U> nilFunc, Func<DerivedAvlTree<T>, T, DerivedAvlTree<T>, U> nodeFunc)
            {
                return nodeFunc(this.left, this.value, this.right);
            }

            protected override int GetCount() { return count; }
            protected override int GetHeight() { return height; }
        }
        #endregion

        public static DerivedAvlTree<T> Create(Func<T, T, int> comparer)
        {
            return new Nil(comparer);
        }

        protected DerivedAvlTree(Func<T, T, int> comparer)
            : base(comparer)
        {
        }

        #region BaseBinaryTree implementation
        protected override DerivedAvlTree<T> CreateNode(DerivedAvlTree<T> left, T value, DerivedAvlTree<T> right)
        {
            return new Node(this.Comparer, left, value, right);
        }

        protected override DerivedAvlTree<T> Self()
        {
            return this;
        }

        #endregion

        protected abstract int GetCount();
        protected abstract int GetHeight();

        protected int Height { get { return GetHeight(); } }
        public int Count { get { return GetCount(); } }

        protected int BalanceFactor
        {
            get
            {
                return this.Match(
                    () => 0,
                    (l, x, r) => l.Height - r.Height);
            }
        }

        public override DerivedAvlTree<T> Insert(T value)
        {
            return base.Insert(value).Balance();
        }

        protected DerivedAvlTree<T> Balance()
        {
            int balanceFactor = this.BalanceFactor;
            if (balanceFactor >= 2)
            {
                // heavy left side

                int innerBalanceFactor =
                    this.Match(
                        () => 0,
                        (l, x, r) => l.BalanceFactor);

                if (innerBalanceFactor >= 1)
                {
                    // left-left case
                    return this.RotateRight();
                }
                else if (innerBalanceFactor <= -1)
                {
                    // left-right case
                    return this.RotateDoubleRight();
                }
            }
            else if (balanceFactor <= -2)
            {
                // heavy right side
                int innerBalanceFactor =
                    this.Match(
                        () => 0,
                        (l, x, r) => r.BalanceFactor);

                if (innerBalanceFactor >= 1)
                {
                    // right-left case
                    return this.RotateDoubleLeft();
                }
                else if (innerBalanceFactor <= -1)
                {
                    // right-right case
                    return this.RotateLeft();
                }
            }

            return this;
        }

        /// <summary>
        /// Left tree rotation.
        /// </summary>
        /// <remarks>
        ///       q      Right Rotation       p
        ///     /   \    --------------+    /   \
        ///    p     c                     a     q
        ///   / \        Left Rotation          / \
        ///  a   b       +-------------        b   c
        /// </remarks>
        protected DerivedAvlTree<T> RotateLeft()
        {
            return this.Match(
                () => this,
                (a, p, q) => q.Match(
                    () => this,
                    (b, qx, c) => CreateNode(CreateNode(a, p, b), qx, c)));
        }

        /// <summary>
        /// Right tree rotation
        /// </summary>
        /// <remarks>
        ///       q      Right Rotation       p
        ///     /   \    --------------+    /   \
        ///    p     c                     a     q
        ///   / \        Left Rotation          / \
        ///  a   b       +-------------        b   c
        /// </remarks>
        protected DerivedAvlTree<T> RotateRight()
        {
            return this.Match(
                () => this,
                (p, q, c) => p.Match(
                    () => this,
                    (a, px, b) => CreateNode(a, px, CreateNode(b, q, c))));
        }

        /// <summary>
        /// Performs a double left rotation
        /// </summary>
        /// <remarks>
        ///     p                       p                       q
        ///   /   \                    / \                    /   \
        ///  a      r                 a   q                 p       r
        ///       /   \    ---------+    / \    -------+   / \     / \
        ///      q     d                b   r             a   b   c   d
        ///     / \                        / \
        ///    b   c                      c   d
        /// </remarks>
        protected DerivedAvlTree<T> RotateDoubleLeft()
        {
            return this.Match(
                () => this,
                (l, x, r) => CreateNode(l, x, r.RotateRight()).RotateLeft());
        }

        /// <summary>
        /// Performs a double right rotation
        /// </summary>
        /// <remarks>
        ///        r                        r                   q
        ///      /   \                     / \                /   \
        ///    p      d                   q   d             p       r
        ///  /   \         ---------+    / \    -------+   / \     / \
        /// a     q                     p   c             a   b   c   d
        ///      / \                   / \
        ///     b   c                 a   b
        /// </remarks>
        protected DerivedAvlTree<T> RotateDoubleRight()
        {
            return this.Match(
                () => this,
                (l, x, r) => CreateNode(l.RotateLeft(), x, r).RotateRight());
        }
    }
}