using System;
using System.Collections.Generic;
using System.Runtime.CompilerServices;
using UnityEngine;

public class SimpleHeap<T>
{
    const int ROOT_IDX = 0;
    const int INVALID_IDX = -1;

    T[] data = new T[0];
    int count = 0;
    Func<T, T, bool> compareFunc;

    public int Count => count;

    public SimpleHeap() { }

    public SimpleHeap(Func<T, T, bool> compareFunc)
    {
        this.compareFunc = compareFunc;
    }

    public SimpleHeap(IEnumerable<T> collection, Func<T, T, bool> compareFunc)
    {
        this.compareFunc = compareFunc;
        foreach (var item in collection)
        {
            Push(item);
        }
    }

    public void SetCompareFunc(Func<T, T, bool> compareFunc)
    {
        this.compareFunc = compareFunc;
    }

    public void Push(T value)
    {
        InsertBack(value);
        RecurcieveClimb(count - 1);
    }

    public T Pop()
    {
        Debug.Assert(count > 0, "Heap is empty");

        T result = data[ROOT_IDX];
        Swap(0, --count);
        RecurcieveFall(0);

        return result;
    }

    public T GetRootValue()
    {
        Debug.Assert(count > 0, "Heap is empty");

        return data[ROOT_IDX];
    }

    public bool TryGetRootValue(out T value)
    {
        if (IsEmpty())
        {
            value = default;
            return false;
        }

        value = data[ROOT_IDX];
        return true;
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public bool IsEmpty() => count == 0;

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public void Clear() => count = 0;

    void RecurcieveClimb(int curIdx)
    {
        if (curIdx == ROOT_IDX) return;

        int parentIdx = (curIdx - 1) >> 1;

        if (compareFunc(data[curIdx], data[parentIdx]))
        {
            Swap(curIdx, parentIdx);
            RecurcieveClimb(parentIdx);
        }
        else
        {
            return;
        }
    }

    void RecurcieveFall(int curIdx)
    {
        int leftChildIdx = GetLeftChildIdx(curIdx);
        int rightChildIdx = GetRightChildIdx(curIdx);

        int aptIdx = curIdx;
        if (leftChildIdx != INVALID_IDX && rightChildIdx != INVALID_IDX)
        {
            aptIdx = GetAptIndex(leftChildIdx, rightChildIdx);
            aptIdx = GetAptIndex(curIdx, aptIdx);
        }
        else if (leftChildIdx != INVALID_IDX)
        {
            aptIdx = GetAptIndex(curIdx, leftChildIdx);
        }

        if (aptIdx == curIdx)
        {
            return;
        }

        Swap(curIdx, aptIdx);
        RecurcieveFall(aptIdx);
    }

    void InsertBack(T value)
    {
        if (count >= data.Length)
            Array.Resize(ref data, count + 1);

        data[count++] = value;
    }

    int GetLeftChildIdx(int curIdx)
    {
        int result = (curIdx << 1) + 1;

        if (result >= count)
            result = INVALID_IDX;

        return result;
    }

    int GetRightChildIdx(int curIdx)
    {
        int result = (curIdx << 1) + 2;

        if (result >= count)
            result = INVALID_IDX;

        return result;
    }

    int GetAptIndex(int value1Index, int value2Index) => compareFunc(data[value1Index], data[value2Index]) ? value1Index : value2Index;

    void Swap(int sourceIdx, int destinationIdx)
    {
        T tmp = data[sourceIdx];
        data[sourceIdx] = data[destinationIdx];
        data[destinationIdx] = tmp;
    }
}