#pragma once

#include <desc/storage/Dataset.hxx>
#include <desc/storage/Itemset.hxx>
#include <desc/utilities/counting_iterator.h>

#include <algorithm>
#include <optional>
#include <stdexcept>
#include <string>
#include <vector>

#if defined(HAS_EXECUTION_POLICIES)
#include <desc/utilities/counting_iterator.h>
#include <execution>
#endif

namespace sd
{

namespace par
{
template <typename Iter, typename Rel = std::less<>>
void sort(Iter first, Iter last, Rel less = Rel{})
{
#if HAS_EXECUTION_POLICIES
    std::sort(std::execution::par_unseq, first, last, less);
#else
    std::sort(first, last, less);
#endif
}

template <typename Container, typename Fn>
void erase_if(Container& x, Fn&& fn)
{
#if HAS_EXECUTION_POLICIES
    auto ptr =
        std::remove_if(std::execution::par_unseq, x.begin(), x.end(), std::forward<Fn>(fn));
    x.erase(ptr, x.end());
#else
    auto ptr = std::remove_if(x.begin(), x.end(), std::forward<Fn>(fn));
    x.erase(ptr, x.end());
#endif
}

template <typename Iter, typename Sentinel, typename visitor_t>
void for_each(Iter first, Sentinel last, visitor_t&& visit)
{
#if HAS_EXECUTION_POLICIES
    std::for_each(std::execution::par_unseq, first, last, visit);
#else
#pragma omp parallel for schedule(dynamic, 1)
    for (auto i = first; i != last; ++i) { visit(*i); }
#endif
}
} // namespace par

namespace disc
{

template <typename S, typename T, typename U>
struct SlimCandidate
{
    using pattern_type = S;
    using float_type   = T;

    itemset<pattern_type>     pattern;
    long_storage_container<U> row_ids;
    size_t                    support = 0;
    T                         score   = 0;

    void clear()
    {
        pattern.clear();
        row_ids.clear();
        support = 0;
        score   = 0;
    }
};

struct candidate_ordering
{
    template <class S, class T>
    constexpr bool operator()(const S& a, const T& b) const noexcept
    {
        return a.score < b.score;
    }
};
struct subset_relation
{
    template <typename S, typename T>
    bool operator()(const S& x, const T& y) const noexcept
    {
        return is_subset(x, y);
    }
};
struct lex_relation
{
    template <typename S, typename T>
    bool operator()(const S& x, const T& y) const noexcept
    {
        return std::lexicographical_compare(
            x.container.begin(), x.container.end(), y.container.begin(), y.container.end());
    }
};
struct equals_pattern
{
    template <class S, class T>
    bool operator()(const S x, const T& y) const noexcept
    {
        return x.support == y.support && sd::equal(y.pattern, x.pattern);
    }
};
struct ConstantScoreFunction
{
    template <class S>
    constexpr int operator()(const S&) const noexcept
    {
        return 1;
    }
};

template <typename S, typename T, typename U>
void join(SlimCandidate<S, T, U>&       next,
          const SlimCandidate<S, T, U>  a,
          const SlimCandidate<S, T, U>& 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, typename U, typename score_fn = ConstantScoreFunction>
int combine_two(SlimCandidate<S, T, U>&       joined,
                const SlimCandidate<S, T, U>& next,
                const SlimCandidate<S, T, U>& other,
                size_t                        min_support,
                score_fn&&                    score = {})
{
    joined.clear();
    if (is_subset(other.pattern, next.pattern)) return 0;

    join(joined, next, other);
    if (joined.support < min_support) return 0;
    joined.score = score(joined);
    return joined.score > 0;
}
template <typename S, typename T, typename U, typename score_fn>
int combine_two_singletons(SlimCandidate<S, T, U>&       joined,
                           const SlimCandidate<S, T, U>& next,
                           const SlimCandidate<S, T, U>& other,
                           size_t                        min_support,
                           score_fn&&                    score)
{
    if (other.support < min_support) return 0;
    if (next.support < min_support) return 0;
    join(joined, next, other);
    if (joined.support < min_support) return 0;
    joined.score = score(joined);
    return 1;
}

template <typename Data, typename S, typename T, typename U>
void create_singleton_candidates(Data const&                          data,
                                 std::vector<SlimCandidate<S, T, U>>& singletons,
                                 size_t                               min_support)
{
    singletons.clear();
    singletons.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)
    {
        for_each_item(point(x), [&](size_t j) { singletons[j].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 score_fn, typename S, typename T, typename U>
void create_pair_candidates_batch(score_fn&&                                 score,
                                  std::vector<SlimCandidate<S, T, U>> const& singletons,
                                  std::vector<SlimCandidate<S, T, U>>&       novel,
                                  std::vector<SlimCandidate<S, T, U>>&       candidates,
                                  size_t                                     min_support)
{
    size_t n = singletons.size();
    novel.clear();
    novel.resize(static_cast<size_t>(std::ceil(n * static_cast<double>(n + 1) / 2)));

    auto sym_index = [n](auto i, auto j)
    {
        if (j < i) { std::swap(i, j); }
        return i * n - (i - 1) * i / 2 + j - i;
    };

    par::for_each(counting_iterator<size_t>(0),
                  counting_iterator<size_t>(n),
                  [&](size_t i)
                  {
                      if (singletons[i].support > min_support)
                      {
                          for (size_t j = i + 1; j < n; ++j)
                          {
                              combine_two_singletons(novel[sym_index(i, j)],
                                                     singletons[i],
                                                     singletons[j],
                                                     min_support,
                                                     score);
                          }
                      }
                  });

    // candidates.reserve(candidates.size() + novel.size());
    std::copy_if(std::make_move_iterator(novel.begin()),
                 std::make_move_iterator(novel.end()),
                 std::back_inserter(candidates),
                 [](const auto& x) { return x.score > 0; });
    novel.clear();
}

template <typename score_fn, typename S, typename T, typename U>
void create_pair_candidates_it(score_fn&&                                 score,
                               std::vector<SlimCandidate<S, T, U>> const& singletons,
                               std::vector<SlimCandidate<S, T, U>>&       novel,
                               std::vector<SlimCandidate<S, T, U>>&       candidates,
                               size_t                                     min_support)
{
    const auto n = singletons.size();
    novel.clear();
    novel.resize(n);

    for (size_t j = 0; j < n; ++j)
    {
        par::for_each(counting_iterator<size_t>(j + 1),
                      counting_iterator<size_t>(n),
                      [&](const size_t i) {
                          combine_two_singletons(
                              novel[i], singletons[i], singletons[j], min_support, score);
                      });

        // candidates.reserve(candidates.size() + novel.size());
        for (size_t i = j + 1; i < n; ++i)
        {
            if (novel[i].score > 0) { candidates.push_back(novel[i]); }
            novel[i].clear();
        }
    }
}
template <typename score_fn, typename S, typename T, typename U>
void create_singular_candidates(score_fn&&                                 score,
                                std::vector<SlimCandidate<S, T, U>> const& singletons,
                                std::vector<SlimCandidate<S, T, U>>&       novel,
                                std::vector<SlimCandidate<S, T, U>>&       candidates,
                                size_t                                     min_support)
{
    size_t n = singletons.size();
    novel.clear();
    novel.resize(n);

    par::for_each(counting_iterator<size_t>(0),
                  counting_iterator<size_t>(n),
                  [&](size_t i)
                  {
                      if (singletons[i].support > min_support)
                      {
                          novel[i]       = singletons[i];
                          novel[i].score = score(novel[i]);
                      }
                  });

    candidates.reserve(candidates.size() + novel.size());
    std::copy_if(std::make_move_iterator(novel.begin()),
                 std::make_move_iterator(novel.end()),
                 std::back_inserter(candidates),
                 [](const auto& x) { return x.score > 0; });
    novel.clear();
}

template <typename score_fn, typename S, typename T, typename U>
void expand_candidate(const SlimCandidate<S, T, U>&              next,
                      std::vector<SlimCandidate<S, T, U>> const& singletons,
                      std::vector<SlimCandidate<S, T, U>>&       novel,
                      score_fn&&                                 score,
                      size_t                                     min_support)
{
    par::for_each(counting_iterator<size_t>(0),
                  counting_iterator<size_t>(singletons.size()),
                  [&](const auto& i)
                  { combine_two(novel[i], next, singletons[i], min_support, score); });
}

template <typename State>
struct CandidateGeneratorImpl
{
    using state_type   = State;
    using pattern_type = typename state_type::pattern_type;

    template <typename Data>
    CandidateGeneratorImpl(const Data&           data,
                           size_t                min_supp,
                           std::optional<size_t> max_tree_depth)
        : max_depth(max_tree_depth), min_support(min_supp)
    {
        create_singleton_candidates(data, singletons, min_support);
    }
    CandidateGeneratorImpl(std::vector<State> const& singletons,
                           size_t                    min_supp,
                           std::optional<size_t>     max_tree_depth)
        : max_depth(max_tree_depth), min_support(min_supp), singletons(singletons)
    {
    }

    const state_type& top() const { return candidates.back(); }

    std::optional<state_type> next()
    {
        if (candidates.empty()) return {};
        auto ret = std::move(candidates.back());
        candidates.pop_back();
        return ret;
    }
    void order_candidates()
    {
        par::sort(candidates.begin(), candidates.end(), candidate_ordering{});
    }
    template <typename Fn>
    void prune(Fn&& fn)
    {
        par::erase_if(candidates, std::forward<Fn>(fn));
    }
    template <typename score_fn>
    void compute_scores(score_fn&& score)
    {
        par::for_each(
            std::begin(candidates), std::end(candidates), [&](auto& x) { x.score = score(x); });
    }
    template <typename score_fn>
    void compute_scores(const state_type& joined, score_fn&& score)
    {
        par::for_each(std::begin(candidates),
                      std::end(candidates),
                      [&](auto& x)
                      {
                          if (intersects(joined.pattern, x.pattern)) x.score = score(x);
                      });
    }
    bool   has_next() const { return !candidates.empty(); }
    size_t size() const { return candidates.size(); }
    // template <typename score_fn = ConstantScoreFunction>
    // void expand_candidate_into_tmp(const state_type& next, score_fn&& score = {})
    // {
    //     novel.resize(singletons.size());
    //     par::for_each(counting_iterator<size_t>(0),
    //                   counting_iterator<size_t>(singletons.size()),
    //                   [&, count_next = count(next.pattern)](const auto& i)
    //                   {
    //                       combine_two(novel[i],
    //                                   count_next,
    //                                   next,
    //                                   singletons[i],
    //                                   min_support,
    //                                   max_depth,
    //                                   known,
    //                                   score);
    //                   });
    // }
    template <typename score_fn, typename prune_fn>
    std::pair<size_t, bool>
    expand_candidate(const state_type& next, score_fn&& score, prune_fn&& prune)
    {
        if (max_depth && count(next.pattern) >= *max_depth) return {size_t(0), false};

        sd::disc::expand_candidate(next, singletons, novel, score, min_support);

        std::pair<size_t, typename state_type::float_type> res = {0, next.score};
        std::for_each(novel.begin(),
                      novel.end(),
                      [&](auto& x)
                      {
                          if (x.score > 0 && !prune(x))
                          {
                              if (res.second < x.score) res = {candidates.size(), x.score};
                              candidates.push_back(x);
                          }
                          x.clear();
                      });
        return {res.first, res.second >= next.score};
    }
    template <typename score_fn, typename prune_fn>
    void expand_1(score_fn&& score, prune_fn&& pred_fn)
    {
        if (!has_next()) return;
        top_state             = candidates.back(); // copy is intentional
        auto [cur, is_better] = expand_candidate(top_state, score, pred_fn);
        if (is_better) std::swap(candidates.back(), candidates[cur]);
    }
    template <typename score_fn, typename prune_fn>
    void expand_bfs(score_fn&& score, prune_fn&& prune, size_t max_layer_expansion)
    {
        if (!has_next()) return;
        for (size_t layer = 0; layer < max_layer_expansion; ++layer)
        {
            top_state             = candidates.back(); // copy is intentional
            auto [cur, is_better] = expand_candidate(top_state, score, prune);
            if (!is_better) break;
            std::swap(candidates.back(), candidates[cur]);
        }
        order_candidates();
    }
    template <typename score_fn>
    void init_candidates(score_fn&& score)
    {
        if (singletons.size() > 10'000)
        {
            create_pair_candidates_it(
                std::forward<score_fn>(score), singletons, novel, candidates, min_support);
        }
        else
        {
            create_pair_candidates_batch(
                std::forward<score_fn>(score), singletons, novel, candidates, min_support);
        }
        order_candidates();
        novel.resize(singletons.size());
    }

private:
    state_type                                                   top_state;
    std::optional<size_t>                                        max_depth;
    size_t                                                       min_support = 2;
    std::vector<state_type>                                      singletons;
    std::vector<state_type>                                      candidates;
    std::vector<state_type>                                      novel;

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

template <typename State>
struct CandidateGeneratorHeapImpl
{
    using state_type   = State;
    using pattern_type = typename state_type::pattern_type;

    template <typename Data>
    CandidateGeneratorHeapImpl(const Data&           data,
                               size_t                min_supp,
                               std::optional<size_t> max_tree_depth)
        : max_depth(max_tree_depth), min_support(min_supp)
    {
        create_singleton_candidates(data, singletons, min_support);
    }
    CandidateGeneratorHeapImpl(std::vector<State> const& singletons,
                               size_t                    min_supp,
                               std::optional<size_t>     max_tree_depth)
        : max_depth(max_tree_depth), min_support(min_supp), singletons(singletons)
    {
        novel.resize(singletons.size());
    }
    const state_type&         top() const { return candidates.back(); }
    std::optional<state_type> next()
    {
        // if (top_state.pattern.empty()) return {};
        // return top_state;

        if (candidates.empty()) return {};
        std::pop_heap(candidates.begin(), candidates.end(), candidate_ordering{});
        auto ret = std::move(candidates.back());
        candidates.pop_back();
        return ret;
    }
    void order_candidates()
    {
        std::make_heap(candidates.begin(), candidates.end(), candidate_ordering{});
    }
    template <typename Fn>
    void prune(Fn&& fn)
    {
        par::erase_if(candidates, std::forward<Fn>(fn));
    }
    template <typename score_fn>
    void compute_scores(score_fn&& score)
    {
        par::for_each(
            std::begin(candidates), std::end(candidates), [&](auto& x) { x.score = score(x); });
    }
    template <typename score_fn>
    void compute_scores(const state_type& joined, score_fn&& score)
    {
        par::for_each(std::begin(candidates),
                      std::end(candidates),
                      [&](auto& x)
                      {
                          if (intersects(joined.pattern, x.pattern)) x.score = score(x);
                      });
    }

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

    template <typename score_fn, typename prune_fn>
    bool combine_pairs_heap(const state_type& next, score_fn&& score, prune_fn&& prune)
    {

        if (max_depth && count(next.pattern) >= *max_depth) return false;
        auto res = next.score * 0;
        expand_candidate(next, singletons, novel, score, min_support);

        std::for_each(novel.begin(),
                      novel.end(),
                      [&](auto& x)
                      {
                          if (x.score > 0 && !prune(x))
                          {
                              res = std::max(res, x.score);
                              candidates.push_back(x);
                              std::push_heap(
                                  candidates.begin(), candidates.end(), candidate_ordering{});
                          }
                          x.clear();
                      });
        return res >= next.score;
    }
    template <typename score_fn, typename prune_fn>
    void expand_1(score_fn&& score, prune_fn&& pred_fn)
    {
        if (!has_next()) return;
        std::pop_heap(candidates.begin(), candidates.end(), candidate_ordering{});
        top_state = candidates.back(); // copy is intentional
        candidates.pop_back();
        // auto improved = 
        combine_pairs_heap(top_state, score, pred_fn);
        candidates.push_back(top_state);
        std::push_heap(candidates.begin(), candidates.end(), candidate_ordering{});
        // if (!improved)
        // {
        // }
    }
    template <typename score_fn, typename prune_fn>
    void expand_bfs(score_fn&& score, prune_fn&& pred_fn, size_t max_layer_expansion)
    {
        if (!has_next()) return;
        for (size_t layer = 0; layer < max_layer_expansion; ++layer)
        {
            std::pop_heap(candidates.begin(), candidates.end(), candidate_ordering{});
            top_state = candidates.back(); // copy is intentional
            candidates.pop_back();
            auto improved = combine_pairs_heap(top_state, score, pred_fn);
            candidates.push_back(top_state);
            std::push_heap(candidates.begin(), candidates.end(), candidate_ordering{});
            // std::cout << improved << " improved? " << top_state.score << std::endl;
            if (!improved)
            {
                break;
            }
        }
    }

    // state_type best;
    // template <typename score_fn, typename prune_fn, typename Acceptable_Decay>
    // bool combine_pairs_heap(const state_type&  next,
    //                         score_fn&&         score,
    //                         prune_fn&&         prune,
    //                         Acceptable_Decay&& acceptable_decay)
    // {

    //     if (max_depth && count(next.pattern) >= *max_depth) return false;

    //     auto res        = next.score;
    //     bool acceptable = false;
    //     expand_candidate(next, singletons, novel, score, min_support);
    //     best       = next;
    //     best.score = 0;

    //     std::for_each(novel.begin(),
    //                   novel.end(),
    //                   [&](auto& x)
    //                   {
    //                       if (x.score > 0 && !prune(x))
    //                       {
    //                           for_each_item(x.pattern, [](auto i) { std::cout << i << " ";
    //                           }); std::cout << ": " << next.score << " " << x.score <<
    //                           std::endl;

    //                           if (acceptable_decay(x))
    //                           {
    //                               std::cout << "is acceptable\n";
    //                               acceptable |= true;
    //                               if (best.score <= x.score) best = x;
    //                           }

    //                           res = std::max(res, x.score);
    //                           candidates.push_back(x);
    //                           std::push_heap(
    //                               candidates.begin(), candidates.end(),
    //                               candidate_ordering{});
    //                       }
    //                       //   if (x.score != -std::numeric_limits<double>::infinity())
    //                       //   {
    //                       //       known.insert(x.pattern);
    //                       //   }
    //                       x.clear();
    //                   });

    //     return acceptable; // res >= next.score;
    // }

    // template <typename score_fn, typename prune_fn, typename Acceptable_Decay>
    // std::optional<state_type> expand_bfs(score_fn&&         score,
    //                                      prune_fn&&         pred_fn,
    //                                      size_t             max_layer_expansion,
    //                                      Acceptable_Decay&& acceptable_decay)
    // {
    //     if (!has_next()) return {};
    //     std::pop_heap(candidates.begin(), candidates.end(), candidate_ordering{});
    //     auto bak = top_state = candidates.back(); // copy is intentional
    //     candidates.pop_back();
    //     bool was_expanded = false;

    //     for (size_t layer = 0; layer < max_layer_expansion; ++layer)
    //     {
    //         std::cout << "Layer: " << layer << std::endl;
    //         for_each_item(top_state.pattern, [](auto i) { std::cout << i << " "; });
    //         std::cout << std::endl;
    //         auto acceptable = combine_pairs_heap(top_state,
    //                                              score,
    //                                              pred_fn,
    //                                              [&](const auto& t)
    //                                              { return acceptable_decay(top_state, t); });
    //         std::cout << "best\t" << std::endl;
    //         for_each_item(best.pattern, [](auto i) { std::cout << i << " "; });
    //         std::cout << std::endl;
    //         if (!acceptable)
    //         {
    //             // std::cout << improved << " improved? " << top_state.score << std::endl;
    //             break;
    //         }
    //         else
    //         {
    //             par::erase_if(candidates,
    //                           [&](auto& x) { return sd::equal(x.pattern, best.pattern); });
    //             order_candidates();
    //             was_expanded = true;
    //             top_state    = best;
    //         }
    //     }
    //     if (!was_expanded)
    //     {
    //         candidates.push_back(bak);
    //         std::push_heap(candidates.begin(), candidates.end(), candidate_ordering{});
    //     }
    //     return top_state;
    //     std::cout << std::endl;
    // }

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

        auto                      K = std::min(candidates.size(), size_t(5));
        std::array<state_type, 5> s;
        for (size_t k = 0; k < K; ++k)
        {
            std::pop_heap(candidates.begin(), candidates.end(), candidate_ordering{});
            s[k] = candidates.back(); // copy is intentional
            candidates.pop_back();
        }

        for (size_t k = 0; k < K; ++k)
        {
            for (size_t layer = 0; layer < max_layer_expansion; ++layer)
            {
                auto improved = combine_pairs_heap(s[k], score, pred_fn);
                candidates.push_back(s[k]);
                std::push_heap(candidates.begin(), candidates.end(), candidate_ordering{});
                // std::cout << improved << " improved? " << top_state.score << std::endl;

                if (!improved) { break; }
            }
        }
    }

    template <typename score_fn>
    void init_candidates(score_fn&& score)
    {
        if (singletons.size() > 10'000)
        {
            create_pair_candidates_it(
                std::forward<score_fn>(score), singletons, novel, candidates, min_support);
        }
        else
        {
            create_pair_candidates_batch(
                std::forward<score_fn>(score), singletons, novel, candidates, min_support);
        }
        order_candidates();
        novel.resize(singletons.size());
    }

private:
    state_type                                                   top_state;
    std::optional<size_t>                                        max_depth;
    size_t                                                       min_support = 2;
    std::vector<state_type>                                      singletons;
    std::vector<state_type>                                      candidates;
    std::vector<state_type>                                      novel;

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

template <typename S, typename T, typename U>
using CandidateGenerator = CandidateGeneratorHeapImpl<SlimCandidate<S, T, U>>;

} // namespace disc
} // namespace sd
