#pragma once

#include <disc/storage/Dataset.hxx>
#include <disc/utilities/Support.hxx>

#include <random>

namespace sd::disc
{

template <typename T, typename RNG>
void generate_pattern(size_t from, size_t to, size_t items, itemset<T>& out, RNG& rng)
{
    out.clear();
    std::uniform_int_distribution<size_t> randint(from, to);
    for (size_t i = 0; i < items; ++i)
        out.insert(from + randint(rng));
}

template <typename Data>
void generate_singletons(size_t from, size_t to, double noise, Data& set)
{
    itemset<typename Data::pattern_type> buf;
    buf.reserve(to);

    for (size_t i = from; i <= to; ++i)
    {
        buf.clear();
        buf.insert(i);
        set.insert(noise, buf);
    }
}

template <typename Data, typename RNG, typename Fr>
void generate_patterns(size_t num_patterns,
                       size_t from,
                       size_t to,
                       Data&  set,
                       RNG&   rng,
                       Fr&&   fr,
                       size_t min_width = 2,
                       size_t max_width = 5)
{
    assert(max_width > min_width);

    itemset<typename Data::pattern_type> buf;
    buf.reserve(to + max_width - min_width);

    std::uniform_int_distribution<size_t> randint(min_width, max_width);

    for (size_t i = 0; i < num_patterns; ++i)
    {
        size_t items = randint(rng);
        generate_pattern<typename Data::pattern_type>(from, to, items, buf, rng);
        set.insert(fr(rng, buf), buf);
    }
}

template <typename Data, typename RNG>
void generate_patterns(size_t num_patterns,
                       size_t from,
                       size_t to,
                       Data&  set,
                       RNG&   rng,
                       size_t min_width = 2,
                       size_t max_width = 5)
{
    auto k = std::max(num_patterns, size_t(5));
    auto u = std::uniform_real_distribution<double>(2, 5);

    generate_patterns(
        num_patterns,
        from,
        to,
        set,
        rng,
        [&](auto& rng, const auto&) { return u(rng) / k; },
        min_width,
        max_width);
}

template <typename Patterns, typename Data, typename RNG>
void plant_patterns(size_t n, Patterns const& s, Data& out, RNG& rng)
{
    std::uniform_real_distribution<double> uniform(0, 1);
    itemset<typename Data::pattern_type>   row;
    while (out.size() < n)
    {
        row.clear();
        for (size_t j = 0; j < s.size(); ++j)
        {
            if (uniform(rng) < s.label(j))
            {
                // row.insert(s.point(j));
                foreach(s.point(j), [&](size_t item) { row.insert(item); });
            }
        }
        if (!row.empty())
            out.insert(row);
    }
}

template <typename Patterns, typename Data, typename RNG>
void plant_patterns_discrete_sampler(size_t          n,
                                     Patterns const& s,
                                     Data&           out,
                                     double          noise,
                                     size_t          from,
                                     size_t          to,
                                     RNG&            rng,
                                     size_t          min_pattern_plant,
                                     size_t          max_pattern_plant)
{
    const auto&                           weights = s.template col<0>();
    std::discrete_distribution<>          d(weights.begin(), weights.end());
    std::uniform_int_distribution<size_t> patterns_per_row(min_pattern_plant,
                                                           max_pattern_plant);
    std::uniform_real_distribution<>      uniform(0, 1);

    itemset<typename Data::pattern_type> row;
    while (out.size() < n)
    {
        auto howmany = patterns_per_row(rng);
        row.clear();

        for (size_t j = from; j < to; ++j)
        {
            if (uniform(rng) < noise)
                row.insert(j);
        }

        for (size_t j = 0; j < howmany; ++j)
        {
            foreach(s.point(d(rng)), [&](size_t item) { row.insert(item); });
            // row.insert(s.point(d(rng)));
        }

        if (!row.empty())
            out.insert(row);
    }
}

template <typename S, typename T>
void compute_frequencies(const Dataset<S>& data, LabeledDataset<T, S>& patterns)
{
    for (auto [q, _] : patterns)
        q = 0;

    for (const auto& [t] : data)
    {
        for (auto [q, x] : patterns)
        {
            if (is_subset(x, t))
            {
                q += 1;
            }
        }
    }

    for (auto [q, _] : patterns)
        q /= data.size();
}

template <typename S, typename T, typename RNG = std::default_random_engine>
void generate_dataset(size_t                rows,
                      size_t                dim,
                      size_t                n_patterns,
                      Dataset<S>&           out,
                      LabeledDataset<T, S>& patterns,
                      double                noise     = 0.1,
                      RNG                   rng       = RNG{},
                      size_t                offset    = 0,
                      size_t                min_width = 2,
                      size_t                max_width = 5)
{
    generate_singletons(offset, dim, noise, patterns);
    generate_patterns(n_patterns, offset, dim, patterns, rng, min_width, max_width);
    plant_patterns(rows, patterns, out, rng);
    compute_frequencies(out, patterns);
}

template <typename S, typename T, typename RNG = std::default_random_engine>
void generate_dataset_discrete_sampler(size_t                rows,
                                       size_t                dim,
                                       size_t                n_patterns,
                                       Dataset<S>&           out,
                                       LabeledDataset<T, S>& patterns,
                                       double                noise     = 0.1,
                                       RNG                   rng       = RNG{},
                                       size_t                offset    = 0,
                                       size_t                min_width = 2,
                                       size_t                max_width = 5,
                                       size_t                min_plant = 2,
                                       size_t                max_plant = 5)
{
    generate_patterns(n_patterns, offset, dim, patterns, rng, min_width, max_width);
    plant_patterns_discrete_sampler(
        rows, patterns, out, noise, offset, dim, rng, min_plant, max_plant);
    generate_singletons(offset, dim, noise, patterns);
    compute_frequencies(out, patterns);
}

} // namespace sd::disc