using System;
using Unity.Collections;

namespace Utility.Collections {
    /// <summary>
    /// A Priority Queue implemented as a Binary Heap. Enqueue and dequeue both have a complexity of O(log(n)).
    ///
    /// The code is intended to maximise performance. It intentionally does not include safety checks,
    /// and only offers the minimal API required to use it as the basis for Dijkstra or A* implementations.
    /// </summary>
    /// <typeparam name="V">The value type.</typeparam>
    public struct NativeBinaryHeap<V> : IDisposable {
        private readonly struct Element {
            public readonly int key;
            public readonly V   value;

            public Element(int key, V value) {
                this.key   = key;
                this.value = value;
            }
        }

        private          NativeArray<Element> data;
        private readonly int                  capacity;
        private readonly IntOrder             order;

        private readonly struct IntOrder : IOrder<int> {
            public int Compare(int lhs, int rhs) => lhs.CompareTo(rhs);

            public int Minimum => int.MinValue;
            public int Maximum => int.MaxValue;
        }

        /// <summary>
        /// Create a new <c>BinaryHeap</c>.
        /// </summary>
        /// <param name="capacity">The maximum number of elements that can be inserted.</param>
        /// <param name="allocator">allocator</param>
        public NativeBinaryHeap(int capacity, Allocator allocator) {
            order              = new IntOrder();
            data               = new NativeArray<Element>(capacity + 2, allocator);
            data[0]            = new Element(order.Minimum, default);
            this.capacity      = capacity;
            data[capacity + 1] = new Element(order.Maximum, default);
            Count              = 0;
            Clear();
        }

        /// <summary>
        /// Removes all keys and values from the <c>BinaryHeap</c>.
        /// </summary>
        public void Clear() {
            Count = 0;
            for (int i = 1; i <= capacity; ++i) {
                data[i] = new Element(order.Maximum, default);
            }
        }

        /// <summary>
        /// Returns the number of elements contained in the <c>BinaryHeap</c>.
        /// </summary>
        public int Count { get; private set; }

        /// <summary>
        /// Enqueues an element.
        ///
        /// Throws if the <c>BinaryHeap</c> is already full.
        /// </summary>
        /// <param name="value">The value to insert.</param>
        /// <param name="key">The key (priority) to insert.</param>
        public void Enqueue(V value, int key) {
            ++Count;
            int hole = Count;
            int pred = hole >> 1;
            while (order.Compare(data[pred].key, key) > 0) {
                data[hole] =   data[pred];
                hole       =   pred;
                pred       >>= 1;
            }

            data[hole] = new Element(key, value);
        }

        /// <summary>
        /// Removes the element with the smallest key (priority), and returns its value.
        ///
        /// Use <c>Count()</c> to determine the current size of the queue. If the queue is empty this method will return a nonsensical value.
        /// </summary>
        /// <returns>The value that was dequeued.</returns>
        public V Dequeue() {
            V min = data[1].value;

            int hole = 1;
            int succ = 2;
            int sz   = Count;

            while (succ < sz) {
                if (order.Compare(data[succ].key, data[succ + 1].key) > 0)
                    succ++;
                data[hole] =   data[succ];
                hole       =   succ;
                succ       <<= 1;
            }

            int pred = hole >> 1;
            while (order.Compare(data[pred].key, data[sz].key) > 0) {
                data[hole] =   data[pred];
                hole       =   pred;
                pred       >>= 1;
            }

            data[hole] = data[sz];

            data[Count] = new Element(order.Maximum, default);
            Count       = sz - 1;

            return min;
        }

        public void Dispose() {
            data.Dispose();
        }
    }
}