// Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved.
//
// NVIDIA CORPORATION and its licensors retain all intellectual property
// and proprietary rights in and to this software, related documentation
// and any modifications thereto.  Any use, reproduction, disclosure or
// distribution of this software and related documentation without an express
// license agreement from NVIDIA CORPORATION is strictly prohibited.
//
#pragma once
#include <algorithm>
#include <functional>
#include <initializer_list>
#include <vector>

namespace carb
{
namespace flatcache
{
template <class T, class Compare = std::less<T>>
struct set
{
    using value_type = T;

    std::vector<T> v;
    Compare cmp;
    using iterator =typename std::vector<T>::iterator;
    using const_iterator =typename std::vector<T>::const_iterator;
    iterator begin()
    {
        return v.begin();
    }
    iterator end()
    {
        return v.end();
    }
    const_iterator begin() const
    {
        return v.begin();
    }
    const_iterator end() const
    {
        return v.end();
    }
    set(const Compare& c = Compare()) : v(), cmp(c)
    {
    }
    template <class InputIterator>
    set(InputIterator first, InputIterator last, const Compare& c = Compare()) : v(first, last), cmp(c)
    {
        std::sort(begin(), end(), cmp);
    }

    set(std::initializer_list<T> _Ilist) : set(_Ilist.begin(), _Ilist.end())
    {
    }

    void reserve(size_t newCapacity)
    {
        v.reserve(newCapacity);
    }

    void clear()
    {
        v.clear();
    }

    iterator insert(const T& t)
    {
        iterator i = std::lower_bound(begin(), end(), t, cmp);
        if (i == end() || cmp(t, *i))
            i = v.insert(i, t);
        return i;
    }

    iterator insert(T&& t)
    {
        iterator i = std::lower_bound(begin(), end(), t, cmp);
        if (i == end() || cmp(t, *i))
            i = v.insert(i, std::move(t));
        return i;
    }

    template <class _Iter>
    void insert(_Iter _First, _Iter _Last)
    { // insert [_First, _Last) one at a time
        for (; _First != _Last; ++_First)
        {
            insert(*_First);
        }
    }

    iterator insert(const_iterator hint, const value_type& value)
    {
        // Measurements show it is faster to ignore hint in this application
        return insert(value);
    }

    void insert(std::initializer_list<T> _Ilist)
    {
        insert(_Ilist.begin(), _Ilist.end());
    }

    size_t erase(const T& key)
    {
        iterator removeElement = find(key);
        if (removeElement != v.end())
        {
            v.erase(removeElement);
            return 1;
        }
        else
        {
            return 0;
        }
    }

    iterator erase(iterator iter)
    {
        return v.erase(iter);
    }

    const_iterator find(const T& t) const
    {
        const_iterator i = std::lower_bound(begin(), end(), t, cmp);
        return i == end() || cmp(t, *i) ? end() : i;
    }

    iterator find(const T& t)
    {
        iterator i = std::lower_bound(begin(), end(), t, cmp);
        return i == end() || cmp(t, *i) ? end() : i;
    }

    bool contains(const T& t) const
    {
        const_iterator i = std::lower_bound(begin(), end(), t, cmp);
        return i != end() && !cmp(t, *i);
    }

    bool operator==(const set<T>& other) const
    {
        return v == other.v;
    }

    bool operator!=(const set<T>& other) const
    {
        return v != other.v;
    }

    size_t size() const
    {
        return v.size();
    }

    T* data()
    {
        return v.data();
    }

    const T* data() const
    {
        return v.data();
    }
};

template <class T, class Compare = std::less<T>>
bool operator<(const set<T, Compare>& left, const set<T, Compare>& right)
{
    return left.v < right.v;
}

template<typename T>
flatcache::set<T> nWayUnion(std::vector<flatcache::set<T>>& srcBuckets)
{
    flatcache::set<T> retval;

    // Calculate the maximum number of destination attributes
    // We could instead calculate it exactly by finding union of attribute names
    size_t maxDestAttrCount = 0;
    for (flatcache::set<T>& srcBucket : srcBuckets)
    {
        maxDestAttrCount += srcBucket.size();
    }
    retval.reserve(maxDestAttrCount);

    auto currentDest = std::back_inserter(retval.v);

    size_t bucketCount = srcBuckets.size();

    // Initialize invariant that nonEmpty is the vector of buckets that have
    // non-zero attribute counts
    struct NonEmptySegment
    {
        // Invariant is current!=end
        typename std::vector<T>::iterator current;
        typename std::vector<T>::iterator end;
    };
    std::vector<NonEmptySegment> nonEmpty;
    nonEmpty.reserve(bucketCount);

    for (size_t i = 0; i != bucketCount; i++)
    {
        if (srcBuckets[i].begin() != srcBuckets[i].end())
        {
            nonEmpty.push_back({ srcBuckets[i].begin(), srcBuckets[i].end() });
        }
    }

    // Keep going until there's only 1 non-empty bucket
    // At that point we can just copy its attributes to the output
    while (1 < nonEmpty.size())
    {
        // Find all the buckets that have the minimum element
        // These are the ones whose iterators will get advanced

        // By the loop guard and the invariant, we know that nonEmpty[0] exists
        // and that nonEmpty[0].current!=nonEmpty[0].end.
        // So *nonEmpty[0].current is a safe dereference
        T minSoFar = *nonEmpty[0].current;
        std::vector<size_t> indicesAtMin;
        indicesAtMin.reserve(nonEmpty.size());
        indicesAtMin.push_back(0);
        for (size_t i = 1; i != nonEmpty.size(); i++)
        {
            if (*nonEmpty[i].current < minSoFar)
            {
                minSoFar = *nonEmpty[i].current;
                indicesAtMin = { i };
            }
            else if (*nonEmpty[i].current == minSoFar)
            {
                indicesAtMin.push_back(i);
            }
        }

        // Copy minimum element to the output
        *currentDest = minSoFar;
        ++(*currentDest);

        // Advance the iterators that pointed to the min

        std::vector<NonEmptySegment> tempNonEmpty;
        tempNonEmpty.reserve(indicesAtMin.size());

        for (size_t i = 0; i != indicesAtMin.size(); i++)
        {
            nonEmpty[indicesAtMin[i]].current++;
        }

        // Maintain the invariant that nonEmpty are the non empty ones
        // Replace with O(n) copy into a temporary if necessary
        auto it = nonEmpty.begin();
        while (it != nonEmpty.end())
        {
            if (it->current == it->end)
            {
                it = nonEmpty.erase(it);
            }
            else
            {
                ++it;
            }
        }
    }

    // By the negation of the guard we know that nonEmpty has zero or one elements
    if (nonEmpty.size() == 1)
    {
        // If one bucket is left, copy its elements to the output
        std::copy(nonEmpty[0].current, nonEmpty[0].end, currentDest);
    }

    return retval;
}

}
}
