#pragma once

#include <disc/storage/Dataset.hxx>
#include <disc/storage/Itemset.hxx>

#include <algorithm>
#include <optional>
#include <vector>

#if __has_include(<execution>) && defined(WITH_EXECUTION_POLICIES)
#define HAS_EXECUTION_POLICIES 1
#include <execution>
#include <tbb/parallel_sort.h>
#endif

namespace sd
{
namespace disc
{

#if defined(HAS_EXECUTION_POLICIES)
static constexpr bool has_execution_policies = true;
#else
static constexpr bool has_execution_policies = false;
#endif

template <typename S, typename T>
struct SlimCandidate
{
    using pattern_type = S;
    itemset<pattern_type>                pattern;
    long_storage_container<pattern_type> row_ids;
    size_t                               support = 0;
    T                                    score   = 0;
};

template <typename S, typename T>
void swap(SlimCandidate<S, T>& a, SlimCandidate<S, T>& b)
{
    using sd::swap;
    using std::swap;

    swap(a.support, b.support);
    swap(a.score, b.score);
    swap(a.pattern, b.pattern);
    swap(a.row_ids, b.row_ids);
}

template <typename S, typename T>
void join(SlimCandidate<S, T>& next, const SlimCandidate<S, T> a, const SlimCandidate<S, T>& b)
{
    next.row_ids.clear();
    intersection(b.row_ids, a.row_ids, next.row_ids);
    next.pattern.assign(a.pattern);
    next.pattern.insert(b.pattern);
    next.support = count(next.row_ids);
}

template <typename S, typename T>
SlimCandidate<S, T> join(const SlimCandidate<S, T> a, const SlimCandidate<S, T>& b)
{
    auto next = a;
    intersection(b.row_ids, a.row_ids, next.row_ids);
    next.pattern.insert(b.pattern);
    next.support = count(next.row_ids);
    return next;
}

template <typename State>
struct SlimGeneratorImpl
{
    using state_type = State;

    struct ordering
    {
        constexpr bool operator()(const state_type& a, const state_type& b) const noexcept
        {
            return a.score < b.score;
        }
    };

    struct lexless_pattern
    {
        template <class S, class T>
        bool operator()(const SlimCandidate<S, T> x, const SlimCandidate<S, T>& y) const
        {
            return std::lexicographical_compare(x.pattern.container.begin(),
                                                x.pattern.container.end(),
                                                y.pattern.container.begin(),
                                                y.pattern.container.end());
        }
    };

    struct equals_pattern
    {
        template <class S, class T>
        bool operator()(const SlimCandidate<S, T> x, const SlimCandidate<S, T>& y) const
        {
            return x.support == y.support && sd::equal(y.pattern, x.pattern);
        }
    };

    struct ConstantScoreFunction
    {
        constexpr int operator()(const state_type&) const noexcept { return 1; }
    };

    template <typename Data, typename score_fn = ConstantScoreFunction>
    SlimGeneratorImpl(const Data&           data,
                      size_t                min_support,
                      std::optional<size_t> max_width = {},
                      score_fn&&            score     = {})
        : max_width(max_width), min_support(min_support)
    {
        init(data, std::forward<score_fn>(score));
    }

    template <typename Data, typename Patternset, typename score_fn = ConstantScoreFunction>
    SlimGeneratorImpl(const Data&           data,
                      const Patternset&     summary,
                      size_t                min_support,
                      std::optional<size_t> max_width = {},
                      score_fn&&            score     = {})
        : max_width(max_width), min_support(min_support)
    {
        init_from_patternset(data, summary, std::forward<score_fn>(score));
    }

    std::optional<state_type> next()
    {
        if (candidates.empty())
            return {};

        // std::pop_heap(candidates.begin(), candidates.end(), ordering{});

        auto ret = std::move(candidates.back());
        candidates.pop_back();

        return ret;
    }

    static bool candidate_known(const state_type&              x,
                                size_t                         count_x,
                                const std::vector<state_type>& candidates)
    {
        return end(candidates) !=
               std::find_if(begin(candidates), end(candidates), [&](const auto& y) {
                   return x.support == y.support && y.pattern.count() == count_x &&
                          equal(y.pattern, x.pattern);
               });
    }

    template <typename score_fn = ConstantScoreFunction>
    size_t combine_two(state_type&       joined,
                       size_t            count_next,
                       const state_type& next,
                       const state_type& other,
                       score_fn&&        score = {})
    {
        if (is_subset(other.pattern, next.pattern))
            return 0;

        join(joined, next, other);
        if (joined.support < min_support)
            return 0;

        auto count_joined = count(joined.pattern);
        // if (is_subset(joined.pattern, next.pattern))         return 0;
        if (count_joined <= count_next)
            return 0;
        if (max_width && count_joined > *max_width)
            return 0;
        joined.score = score(joined);
        return count_joined;
    }

    // void insert_if(size_t                      count_joined,
    //                std::optional<state_type>&& joined,
    //                std::vector<state_type>&    candidates,
    //                bool                        check_if_exists)
    // {
    //     // if (joined)
    //     // {
    //     //     candidates.emplace_back(std::move(*joined));
    //     // }

    //     if (joined && (!check_if_exists || !candidate_known(*joined, count_joined,
    //     candidates, parallel)))
    //     {
    //         candidates.emplace_back(std::move(*joined));
    //     }
    // }

    template <typename score_fn = ConstantScoreFunction>
    void combine_pairs_allocate_tmp(const state_type& next, score_fn&& score = {})
    {
        const size_t count_next = count(next.pattern);

#pragma omp parallel
        {
            std::vector<state_type> novel;
            novel.reserve(singletons.size() * 0.33 / 24);
            state_type joined;
#pragma omp for nowait
            for (size_t i = 0; i < singletons.size(); ++i)
            {
                if (intersects(singleton_blacklist[i], next.pattern))
                    continue;
                auto count_joined = combine_two(joined, count_next, next, singletons[i], score);

                if (count_joined > 0)
                {
                    novel.push_back(joined);
                }
            }

#pragma omp critical
            {
                candidates.reserve(candidates.size() + novel.size());
                candidates.insert(candidates.end(),
                                  std::move_iterator(novel.begin()),
                                  std::move_iterator(novel.end()));
            }
        }
    }

    template <typename score_fn = ConstantScoreFunction>
    void combine_pairs_locking(const state_type& next, score_fn&& score = {})
    {
        const size_t count_next = count(next.pattern);

        candidates.reserve(candidates.size() + singletons.size() * 0.2);
        state_type joined;

#pragma omp parallel for firstprivate(joined)
        for (size_t i = 0; i < singletons.size(); ++i)
        {
            if (intersects(singleton_blacklist[i], next.pattern))
                continue;
            auto count_joined = combine_two(joined, count_next, next, singletons[i], score);
            if (count_joined > 0)
            {
#pragma omp critical
                {
                    candidates.push_back(joined);
                }
            }
        }
    }

    template <typename score_fn = ConstantScoreFunction>
    void combine_pairs(const state_type& next, bool make_unique, score_fn&& score = {})
    {
        if (singletons.size() > 128)
            combine_pairs_allocate_tmp(next, std::forward<score_fn>(score));
        else
            combine_pairs_locking(next, std::forward<score_fn>(score));

        if (make_unique)
        {
            remove_duplicates();
        }
    }

    void remove_duplicates()
    {

#if HAS_EXECUTION_POLICIES
        if constexpr (has_execution_policies)
        {
            if (candidates.size() > 1024 * 10)
            {
                // GCC 10's std::sort(std::execution::par uses a parallel_stable_sort that might
                // allocate memory

                std::sort(std::execution::par_unseq,
                          candidates.begin(),
                          candidates.end(),
                          lexless_pattern{});
                // tbb::parallel_sort(candidates.begin(), candidates.end(), lexless_pattern{});

                auto ptr = std::unique(std::execution::par_unseq,
                                       candidates.begin(),
                                       candidates.end(),
                                       equals_pattern{});
                candidates.erase(ptr, candidates.end());
                return;
            }
        }
#endif
        std::sort(candidates.begin(), candidates.end(), lexless_pattern{});
        auto ptr = std::unique(candidates.begin(), candidates.end(), equals_pattern{});
        candidates.erase(ptr, candidates.end());
    }

    void order_candidates()
    {
#if HAS_EXECUTION_POLICIES
        if constexpr (has_execution_policies)
        {
            if (candidates.size() > 1024 * 10)
            {
                // tbb::parallel_sort(candidates.begin(), candidates.end(), ordering{});
                std::sort(std::execution::par_unseq,
                          candidates.begin(),
                          candidates.end(),
                          ordering{});
                return;
            }
        }
#endif
        std::sort(candidates.begin(), candidates.end(), ordering{});
    }

    template <typename Fn>
    void prune(Fn&& fn)
    {
#if HAS_EXECUTION_POLICIES
        if constexpr (has_execution_policies)
        {
            if (candidates.size() > 1024 * 10)
            {
                auto ptr = std::remove_if(std::execution::par_unseq,
                                          candidates.begin(),
                                          candidates.end(),
                                          std::forward<Fn>(fn));
                candidates.erase(ptr, candidates.end());
                return;
            }
        }
#endif
        auto ptr = std::remove_if(candidates.begin(), candidates.end(), std::forward<Fn>(fn));
        candidates.erase(ptr, candidates.end());
    }

    template <typename score_fn>
    void compute_scores(score_fn&& score)
    {
#if HAS_EXECUTION_POLICIES
        if constexpr (has_execution_policies)
        {
            std::for_each(std::execution::par_unseq,
                          std::begin(candidates),
                          std::end(candidates),
                          [&](auto& x) { x.score = score(x); });
        }
        else
#endif
        {

#pragma omp parallel for
            for (size_t i = 0; i < candidates.size(); ++i)
            {
                candidates[i].score = score(candidates[i]);
            }
        }
    }

    template <typename score_fn>
    void compute_scores(const state_type& joined, score_fn&& score)
    {
#if HAS_EXECUTION_POLICIES
        if constexpr (has_execution_policies)
        {
            std::for_each(std::execution::par_unseq,
                          std::begin(candidates),
                          std::end(candidates),
                          [&](auto& x) {
                              if (intersects(joined.pattern, x.pattern))
                                  x.score = score(x);
                          });
        }
        else
#endif
        {
#pragma omp parallel for
            for (size_t i = 0; i < candidates.size(); ++i)
            {
                if (intersects(joined.pattern, candidates[i].pattern))
                    candidates[i].score = score(candidates[i]);
            }
        }
    }

    template <typename score_fn>
    void expand_from(const state_type& next, score_fn&& score)
    {
        combine_pairs(next, true);
        // compute_scores(std::forward<score_fn>(score));
        compute_scores(next, std::forward<score_fn>(score));
        order_candidates();
    }

    template <typename score_fn, typename prune_fn>
    void expand_bfs(score_fn&& score, prune_fn&& prune_pred, size_t max_layer_expansion)
    {
        if (!has_next())
            return;

        auto best_score = candidates.back().score;
        auto best_point = candidates.back().pattern;

        for (size_t layer = 0; layer < max_layer_expansion; ++layer)
        {
            auto top = candidates.back(); // copy is intentional

            // expand_from(top, score);
            combine_pairs(top, true, score);
            order_candidates();

            prune(prune_pred);

            if (!has_next())
                break;

            if (best_score == candidates.back().score &&
                sd::equal(best_point, candidates.back().pattern))
            {
                break;
            }

            best_point = candidates.back().pattern;
        }
    }

    template <typename score_fn, typename prune_fn>
    void expand_from(const state_type& next,
                     score_fn&&        score,
                     prune_fn&&        prune_pred,
                     size_t            max_search_depth)
    {
        expand_from(next, score);

        if (max_search_depth > 1)
        {
            expand_bfs(score, std::forward<prune_fn>(prune_pred), max_search_depth - 1);
        }
        else
        {
            prune(std::forward<prune_fn>(prune_pred));
        }
    }

    bool   has_next() const { return !candidates.empty(); }
    size_t size() const { return candidates.size(); }

    template <typename Data>
    void init_singletons(Data const& data)
    {
        singletons = std::vector<state_type>(data.dim);
        singleton_blacklist.clear();
        singleton_blacklist.resize(data.dim);

        if (data.size() == 0 || data.dim == 0)
            return;

        size_t i = 0;
        for (auto& s : singletons)
        {
            s.row_ids.reserve(data.size());
            s.pattern.insert(i++);
        }

        size_t row_index = 0;
        for (const auto& x : data)
        {
            foreach(point(x), [&](size_t i) { singletons[i].row_ids.insert(row_index); });
            ++row_index;
        }

        for (auto& s : singletons)
        {
            s.support = count(s.row_ids);
        }

        singletons.erase(std::remove_if(singletons.begin(),
                                        singletons.end(),
                                        [&](const auto& s) { return s.support < min_support; }),
                         singletons.end());
    }

    template <typename Data, typename score_fn = ConstantScoreFunction>
    void init(Data const& data, score_fn&& score = {})
    {
        init_singletons(data);
        candidates.clear();
        candidates.reserve(data.dim * 3);

        state_type joined;
#pragma omp parallel for firstprivate(joined) schedule(dynamic, 1) // collapse (2)
        for (size_t i = 0; i < singletons.size(); ++i)
        {
            if (singletons[i].support >= min_support)
            {
                for (size_t j = i + 1; j < singletons.size(); ++j)
                {
                    if (singletons[j].support >= min_support)
                    {
                        join(joined, singletons[i], singletons[j]);
                        if (max_width && count(joined.pattern) > *max_width)
                            continue;
                        if (joined.support >= min_support)
                        {
                            joined.score = score(joined);
                            if (joined.score > 0)
                            {
#pragma omp critical
                                {
                                    candidates.push_back(joined);
                                }
                            }
                        }
                        else
                        {
#pragma omp critical
                            {
                                singleton_blacklist[i].insert(j);
                                singleton_blacklist[j].insert(i);
                            }
                        }
                    }
                }
            }
        }
        order_candidates();
    }

    template <typename Data, typename Patternset, typename score_fn = ConstantScoreFunction>
    void
    init_from_patternset(Data const& data, const Patternset& patternset, score_fn&& score = {})
    {
        init(data, score);
        for (const auto& x : patternset)
        {
            combine_pairs(point(x), false);
        }
        remove_duplicates();
        order_candidates();
    }

private:
    std::optional<size_t>   max_width;
    size_t                  min_support = 2;
    std::vector<state_type> singletons;
    std::vector<state_type> candidates;

    std::vector<itemset<typename state_type::pattern_type>> singleton_blacklist;

    static_assert(std::is_swappable_v<state_type>);
};

template <typename S, typename T>
using SlimGenerator = SlimGeneratorImpl<SlimCandidate<S, T>>;

} // namespace disc
} // namespace sd

// void remove_duplicates()
// {
//     std::sort(
//         begin(candidates)
//         end(candidates),
//         [](const auto& x, const auto& y) { return lexicographical_ordering(begin(x),
//         end(x), begin(y), end(y)); });

//     auto pos = std::unique(
//         begin(candidates)
//         end(candidates),
//         [](const auto& y) { return y.count() == count_joined && equal(y, joined); }
//     );
//     candidates.erase(pos, end(candidates));
// }

//         const size_t count_next = count(next.pattern);

// #pragma omp declare reduction (merge : std::vector<state_type> :
// omp_out.insert(omp_out.end(), omp_in.begin(), omp_in.end()))
// initializer(omp_priv=std::vector<state_type>{})

// #pragma omp parallel for  reduction(merge:candidates) if (singletons.size() > 10)
//         for (size_t i = 0; i < singletons.size(); ++i)
//         {
//             if (is_subset(singletons[i].pattern, next.pattern)) continue;

//             auto joined       = join(next, singletons[i]);
//             auto count_joined = count(joined.pattern);
//             // if (is_subset(joined.pattern, next.pattern))         continue;
//             if (count_joined <= count_next)             continue;
//             if (joined.support < min_support)           continue;
//             if (max_width && count_joined > *max_width) continue;
//                 candidates.emplace_back(std::move(joined));
//         }

//         if (!check_if_exists)
//         {
//              remove_duplicates();
//         }
//     }
// }
