using System.Diagnostics.CodeAnalysis;
using System.Drawing;

namespace Quadratic.Carto.Craft2.Containers;

/// <summary>
/// Weight-balanced tree. Immutable.
/// 
/// Also see: <a href="https://en.wikipedia.org/wiki/Weight-balanced_tree">Wikipedia</a>,
/// the <a href="https://dl.acm.org/doi/pdf/10.1145/1499949.1500040">original paper</a>,
/// and the <a href="https://yoichihirai.com/bst.pdf">paper</a> by Hirai about selecting the right balance factor.
/// </summary>
public struct WBTree<T> where T : IComparable<T>
{
    Node? root;

    class Node(T value, Node? left = null, Node? right = null)
    {
        public int Size { get; } = NSize(left) + NSize(right) + 1;
        public T Value { get; set; } = value;
        public Node? Left { get; } = left;
        public Node? Right { get; } = right;

        public void Deconstruct(out T value, out Node? left, out Node? right)
        {
            value = this.Value;
            left = this.Left;
            right = this.Right;
        }

        public bool Balanced()
        {
            return IsBalanced(Left, Right) && IsBalanced(Right, Left) &&
                (Left?.Balanced() ?? true) && (Right?.Balanced() ?? true);
        }

        private static int NSize(Node? n) => n?.Size ?? 0;

        private static bool IsBalanced(Node? l, Node? r) =>
            DELTA * (NSize(l) + 1) >= (NSize(r) + 1);

        private static bool IsSingle(Node? l, Node? r) =>
            (NSize(l) + 1) < GAMMA * (NSize(r) + 1);

        private const int DELTA = 2;
        private const int GAMMA = 3;

        public static Node Insert(Node? node, T value)
        {
            if (node == null) return new Node(value);
            var cmp = value.CompareTo(node.Value);
            if (cmp < 0)
                return BalanceR(node.Value, Insert(node.Left, value), node.Right);
            else if (cmp > 0)
                return BalanceL(node.Value, node.Left, Insert(node.Right, value));
            else
                return new Node(value, node.Left, node.Right);
        }

        private static Node BalanceL(T v, Node? l, Node r)
        {
            if (IsBalanced(l, r)) return new Node(v, l, r);
            // rotateL
            else if (IsSingle(l, r))
            {
                var (k2, t2, t3) = r;
                return new Node(k2, new Node(v, l, t2), t3);
            }
            else
            {
                var (k2, rl, t4) = r;
                if (rl is null)
                    throw new ArgumentException("Unable to perform double rotation; r.Left is null");
                var (k3, t2, t3) = rl;
                return new Node(k3, new Node(v, l, t2), new Node(k2, t3, t4));
            }
        }

        private static Node BalanceR(T cVal, Node l, Node? r)
        {
            if (IsBalanced(r, l)) return new Node(cVal, l, r);
            // rotateR
            else if (IsSingle(r, l))
            {
                var (k2, t2, t3) = l;
                return new Node(k2, t2, new Node(cVal, t3, r));
            }
            else
            {
                var (k2, t2, lr) = l;
                if (lr is null)
                    throw new ArgumentException("Unable to perform double rotation; l.Right is null");
                var (k3, t3, t4) = lr;
                return new Node(k3, new Node(k2, t2, t3), new Node(cVal, t4, r));
            }
        }
    }

    private WBTree(Node? node)
    {
        this.root = node;
    }

    public static WBTree<T> Empty => new WBTree<T>(null);

    public static WBTree<T> Singleton(T value) => new WBTree<T>(new Node(value));

    public readonly int Count => root is null ? 0 : root.Size;

    public WBTree<T> Add(T value) => new WBTree<T>(Node.Insert(root, value));

    public bool TryGet(T value, [NotNullWhen(true)] out T? result)
    {
        var node = root;
        while (node != null)
        {
            var cmp = value.CompareTo(node.Value);
            if (cmp < 0)
                node = node.Left;
            else if (cmp > 0)
                node = node.Right;
            else
            {
                result = node.Value;
                return true;
            }
        }
        result = default;
        return false;
    }

    public bool Contains(T value) => TryGet(value, out _);

    public T Get(T value) => TryGet(value, out var result) ? result : throw new KeyNotFoundException();
}

/// <summary>
/// A mutable reference to a weight-balanced tree. Might be preferable to use in some cases.
/// </summary>
/// <typeparam name="T"></typeparam>
public class WBTreeRef<T> where T : IComparable<T>
{
    WBTree<T> impl;

    public WBTreeRef() => impl = WBTree<T>.Empty;

    public WBTreeRef(T value) => impl = WBTree<T>.Singleton(value);

    public int Count => impl.Count;

    public void Add(T value) => impl = impl.Add(value);

    public bool TryGet(T value, [NotNullWhen(true)] out T? result) => impl.TryGet(value, out result);

    public bool Contains(T value) => impl.Contains(value);

    public T Get(T value) => impl.Get(value);

    public WBTree<T> ToImmutable() => impl;

    public static implicit operator WBTreeRef<T>(WBTree<T> tree) => new() { impl = tree };

    public static implicit operator WBTree<T>(WBTreeRef<T> tree) => tree.impl;
}