#pragma once

#include <disc/disc/CharacterizeComponents.hxx>
#include <disc/disc/Desc.hxx>
#include <disc/disc/InsertMissingSingletons.hxx>
#include <relaxent/DESCVariants.hxx>

#include <chrono>
#include <random>
#include <tuple>

namespace sd::disc
{

template <typename T = std::chrono::milliseconds, typename Fn>
auto timeit(Fn&& fn)
{
    namespace ch = std::chrono;
    // using namespace std::chrono_literals;

    auto before = ch::high_resolution_clock::now();
    std::forward<Fn>(fn)();
    auto after = ch::high_resolution_clock::now();
    // return (after - before) / 1ms;
    return ch::duration_cast<T>(after - before);
}


template <typename DataType, typename ItemsetType>
void insert_missing_singletons_from_main_dataset(const DataType& data, ItemsetType& summary)
{
    // this function is needed because I use distributions to classify points that may lie in a
    // different domain

    itemset<tag_dense> set(std::max(data.dim, summary.dim));

    for (const auto& x : summary)
        if (is_singleton(point(x)))
            set.insert(front(point(x)));

    set.flip_all();
    foreach(set, [&](auto i) { summary.insert(0, i); });
}

template <typename Miner = DefaultPatternsetMinerInterface, typename Trait, typename Data>
auto training_desc(Data data, const MiningSettings& cfg)
{
    Composition<Trait> c;
    c.data = std::move(data);
    initialize_composition(c, cfg);
    discover_patterns_generic(c, cfg, Miner{});
    return c.models;
}

template <typename Miner = DefaultPatternsetMinerInterface, typename Trait, typename Data>
auto training_desc1(const Data& data, const MiningSettings& cfg)
{
    const size_t                                   k = data.num_components();
    std::vector<typename Trait::distribution_type> models;
    models.reserve(k);

    for (size_t i = 0; i < k; ++i)
    {
        auto c = PatternsetResult<Trait>(data.subset(i));
        insert_missing_singletons(c.data, c.summary);

        auto b = c.summary;
        insert_missing_singletons_from_main_dataset(data, c.summary);
        initialize_model(c);

        c.data.dim         = std::max(c.data.dim, c.summary.dim);
        c.initial_encoding = c.encoding = encode(c, cfg);

        c.summary = b;
        discover_patterns_generic(c, cfg, Miner{});
        models.emplace_back(std::move(c.model));
    }

    return models;
}

template <typename Models, typename Data, typename Predictions>
void classify(const Models& ps, const Data& test, Predictions& ys)
{
    using float_type =
        std::decay_t<decltype(ps[0].expected_generalized_frequency(test.point(0)))>;

    size_t j = 0;
    for (const auto&& x : test)
    {
        std::pair<size_t, float_type> opt = {test.num_components() + 1, -1.0};
        for (size_t i = 0; i < ps.size(); ++i)
        {
            const auto p = ps[i].expected_generalized_frequency(point(x));
            if (opt.second < p)
            {
                opt = {i, p};
            }
        }
        ys[j++] = opt.first;
    }
}

template <typename Data, typename RNG>
auto unbalanced_partition(const Data& data, double fraction, RNG&& rng)
{
    Data a, b;

    std::vector<size_t> s(data.size());
    std::iota(begin(s), end(s), 0);
    std::shuffle(begin(s), end(s), rng);

    size_t n = data.size() * fraction;

    for (size_t i = 0; i < n; ++i)
    {
        a.insert(data.point(s[i]), data.label(s[i]));
    }
    for (size_t i = n; i < data.size(); ++i)
    {
        b.insert(data.point(s[i]), data.label(s[i]));
    }

    a.group_by_label();
    b.group_by_label();

    a.dim = b.dim = data.dim;

    return std::make_pair(std::move(a), std::move(b));
}

template <typename Data, typename RNG>
auto balanced_partition(const Data& data, double fraction, RNG&& rng)
{
    Data a, b;

    std::vector<size_t> s;
    s.reserve(data.size());

    for (size_t j = 0; j < data.num_components(); ++j)
    {
        const auto& x = data.subset(j);

        s.resize(x.size());
        std::iota(begin(s), end(s), 0);
        std::shuffle(begin(s), end(s), rng);

        const size_t n = std::floor(x.size() * fraction);

        for (size_t i = 0; i < n; ++i)
        {
            a.insert(point(x[s[i]]), label(x[s[i]]));
        }
        for (size_t i = n; i < x.size(); ++i)
        {
            b.insert(point(x[s[i]]), label(x[s[i]]));
        }
    }

    a.group_by_label();
    b.group_by_label();

    a.dim = b.dim = data.dim;

    return std::make_pair(std::move(a), std::move(b));
}

template <typename X, typename Y>
double l1_error(const X& xs, const Y& ys)
{
    double err = 0;

    for (size_t i = 0; i < xs.size(); ++i)
    {
        // std::cout << xs[i] << " " << ys[i] << std::endl;
        err += static_cast<size_t>(xs[i]) != static_cast<size_t>(ys[i]);
    }

    return err / xs.size();
}

// template<typename Data>
// void print_data(const Data& data)
// {
//     for(const auto && x : data)
//     {
//         foreach(point(x), [](size_t i) { printf(" %d", i); });
//         puts("");
//     }
// }

template <typename Trait, typename S, typename TrainingFunction>
auto cross_validate_impl(PartitionedData<S>&   total,
                    TrainingFunction &&   do_training,
                    const double          train_fraction = 0.8,
                    const size_t          folds          = 10,
                    const size_t          seed           = 13,
                    const MiningSettings& cfg            = {})
{

    using distribution = typename Trait::distribution_type;

    total.group_by_label();
    simplify_labels(total);

    if (train_fraction > 1 || train_fraction < 0)
        throw std::runtime_error{"illegal fraction of training data (0 <= fraction <= 1)"};
    if (total.num_components() <= 1)
        throw std::runtime_error{"no class labels given"};

    std::vector<double> test_errors(folds), train_errors(folds);

    namespace ch = std::chrono;
    using ms     = ch::milliseconds;
    std::vector<ms> test_time(folds), train_time(folds);

    std::default_random_engine rng{seed};

    for (size_t fold = 0; fold < folds; ++fold)
    {
        auto [train, test] = balanced_partition(total, train_fraction, rng);
        // auto [train, test] = unbalanced_partition(total, train_fraction, rng);

        std::vector<size_t> ytrain(train.size()), ytest(test.size());

        std::vector<distribution> p;
        train_time[fold] = timeit([&] { p = do_training(train, cfg); });

        test_time[fold] = timeit([&] {
            classify(p, train, ytrain);
            classify(p, test, ytest);

            train_errors[fold] = l1_error(train.template col<0>(), ytrain);
            test_errors[fold]  = l1_error(test.template col<0>(), ytest);
        });
    }

    return std::make_tuple(train_errors, test_errors, train_time, test_time);
}


enum class TrainingType
{
    DESC,
    CDESC,
    UDESC,
    UCDESC,
    DESC1,
    CDESC1,
    UDESC1,
    UCDESC1,
};


template <typename Trait, typename S>
auto cross_validate(PartitionedData<S>&   total,
                    const double          train_fraction = 0.8,
                    const size_t          folds          = 10,
                    const size_t          seed           = 13,
                    const TrainingType    kind           = TrainingType::DESC,
                    const MiningSettings& cfg            = {})
{
    // clang-format off   
    switch(kind) {
    default:
    case TrainingType::DESC:    return cross_validate_impl<Trait>(total, [](const auto& train, const auto& cfg) { return training_desc <IDesc,   Trait> (train, cfg); }, train_fraction, folds, seed, cfg);  
    case TrainingType::CDESC:   return cross_validate_impl<Trait>(total, [](const auto& train, const auto& cfg) { return training_desc <ICDesc,  Trait> (train, cfg); }, train_fraction, folds, seed, cfg);  
    case TrainingType::UDESC:   return cross_validate_impl<Trait>(total, [](const auto& train, const auto& cfg) { return training_desc <IUDesc,  Trait> (train, cfg); }, train_fraction, folds, seed, cfg);  
    case TrainingType::UCDESC:  return cross_validate_impl<Trait>(total, [](const auto& train, const auto& cfg) { return training_desc <IUCDesc, Trait> (train, cfg); }, train_fraction, folds, seed, cfg);  
    case TrainingType::DESC1:   return cross_validate_impl<Trait>(total, [](const auto& train, const auto& cfg) { return training_desc1<IDesc,   Trait> (train, cfg); }, train_fraction, folds, seed, cfg);  
    case TrainingType::CDESC1:  return cross_validate_impl<Trait>(total, [](const auto& train, const auto& cfg) { return training_desc1<ICDesc,  Trait> (train, cfg); }, train_fraction, folds, seed, cfg);  
    case TrainingType::UDESC1:  return cross_validate_impl<Trait>(total, [](const auto& train, const auto& cfg) { return training_desc1<IUDesc,  Trait> (train, cfg); }, train_fraction, folds, seed, cfg);  
    case TrainingType::UCDESC1: return cross_validate_impl<Trait>(total, [](const auto& train, const auto& cfg) { return training_desc1<IUCDesc, Trait> (train, cfg); }, train_fraction, folds, seed, cfg);  
    }
    // clang-format on
}
} // namespace sd::disc