#pragma once

#include <boost/math/distributions/complement.hpp>
#include <cmath>
#include <elisa/ElisaZtest.hxx>

#include <boost/math/distributions/students_t.hpp>

namespace sd::disc
{

template <typename T>
auto students_t_pvalue(T mu, T sd, size_t n_samples, T x)
{
    // unequal variances:
    T v = n_samples - 2;
    T t = (mu - x) / (sd * sd * sqrt(T(1) / n_samples));

    using boost::math::students_t_distribution, boost::math::complement, boost::math::cdf,
        std::abs;

    T q = cdf(complement(students_t_distribution<T>(v), abs(t)));
    return 2 * q; // prob. that difference is due to chance
}
template <typename T>
auto students_t_log_pvalue(T mu, T sd, size_t n_samples, T x)
{
    using std::log;
    return log(students_t_pvalue(mu, sd, n_samples, x));
}

template <typename T>
auto students_t_test(
    const T& mu, const T& sd, const T& q, const size_t sample_size, const T& log_alpha) -> T
{
    auto h = -students_t_log_pvalue(mu, sd, sample_size, q);
    return h <= -log_alpha ? T(0) : h;
}

template <typename Trait, typename Candidate>
auto elisa_students_heuristic(const Component<Trait>& c,
                              const Candidate&        x,
                              const Config&           cfg,
                              size_t                  layer)
{
    using float_type = typename Trait::float_type;

    const auto q        = static_cast<float_type>(x.support) / c.data.size();
    const auto[mu, var] = expectation_variance(c.model.model, x.pattern);
    const auto th       = fwer::log_alpha_correction<float_type>(cfg.alpha, layer, c.data.dim);
    return students_t_test(mu, std::sqrt(var), q, c.data.size(), th);
}

template <typename Trait, typename Candidate>
auto elisa_students_heuristic(const Composition<Trait>& c,
                              const Candidate&          x,
                              const Config&             cfg,
                              size_t                    layer)
{
    using float_type = typename Trait::float_type;

    auto k  = c.data.num_components();
    auto th = fwer::log_alpha_correction<float_type>(cfg.alpha / k, layer, c.data.dim);

    float_type acc = 0;
    for (size_t i = 0; i < k; ++i) {
        auto n_i = c.data.subset(i).size();
        auto q   = static_cast<float_type>(size_of_intersection(x.row_ids, c.masks[i])) / n_i;
        const auto[mu, var] = expectation_variance(c.models[i].model, x.pattern);
        acc = std::max<float_type>(acc, students_t_test(mu, std::sqrt(var), q, n_i, th));
    }

    return acc;
}

template <typename T>
struct IElisaStudentsT
{
    IElisaStudentsT(T alpha, size_t d)
        : layer(1), log_alpha(fwer::log_alpha_correction(alpha, 1, d))
    {
    }
    IElisaStudentsT(const IElisaStudentsT&) = default;
    IElisaStudentsT(IElisaStudentsT&&)      = default;
    IElisaStudentsT& operator=(const IElisaStudentsT&) = default;
    IElisaStudentsT& operator=(IElisaStudentsT&&) = default;
    IElisaStudentsT()                             = default;

    template <typename C, typename Config>
    void prepare(C& c, const Config& cfg)
    {
        *this = IElisaStudentsT<T>(cfg.alpha, c.data.dim);
        sd::disc::prepare(c, cfg);
    }

    template <typename C, typename Candidate, typename Config>
    auto is_allowed(const C& c, const Candidate& x, const Config& cfg)
    {
        if (x.score < -log_alpha) return false;
        return DefaultPatternsetMinerInterface::is_allowed(c, x, cfg);
    }

    template <typename C, typename Candidate, typename Config>
    auto heuristic(C& c, Candidate& x, const Config& cfg) const
    {
        return elisa_students_heuristic(c, x, cfg, std::max(layer, count(x.pattern)));
    }

    template <typename C, typename Candidate, typename Config>
    auto insert_into_model(C& c, Candidate& x, const Config& cfg)
    {
        update_level(c, x, cfg);
        if (x.score < -log_alpha) return false;
        return sd::disc::find_assignment(c, x, cfg, *this);
    }

    template <typename C, typename Config>
    auto finish(C& c, const Config& cfg)
    {
        prune_composition(c, cfg, *this);
    }

    template <typename Trait, typename Pattern>
    auto confidence(const Composition<Trait>&         c,
                    size_t                            index,
                    const typename Trait::float_type& q,
                    const Pattern&                    x,
                    const Config&) const -> typename Trait::float_type
    {
        using float_type    = typename Trait::float_type;
        const auto[mu, var] = disc::expectation_variance(c.models[index].model, x);
        const auto n        = c.data.subset(index).size();
        return students_t_test(mu, std::sqrt(var), q, n, static_cast<float_type>(log_alpha));
    }

    template <typename Trait, typename Pattern>
    auto confidence(const Component<Trait>&           c,
                    const typename Trait::float_type& q,
                    const Pattern&                    x,
                    const Config&) const -> typename Trait::float_type
    {
        using float_type    = typename Trait::float_type;
        const auto[mu, var] = disc::expectation_variance(c.model.model, x);
        const auto n        = c.data.size();
        return students_t_test(mu, std::sqrt(var), q, n, static_cast<float_type>(log_alpha));
    }

protected:
    template <typename Trait, typename Candidate, typename Config>
    void update_level(Component<Trait> const& c, Candidate& x, const Config& cfg)
    {
        layer     = std::max(layer, count(x.pattern));
        log_alpha = fwer::log_alpha_correction(cfg.alpha, layer, c.data.dim);
    }
    template <typename Trait, typename Candidate, typename Config>
    void update_level(Composition<Trait> const& c, Candidate& x, const Config& cfg)
    {
        layer = std::max(layer, count(x.pattern));
        log_alpha =
            fwer::log_alpha_correction(cfg.alpha / c.data.num_components(), layer, c.data.dim);
    }

    size_t layer     = 1;
    T      log_alpha = 1;
};

struct ElisaStudentsTCharacterization
{
    template <typename Trait, typename Pattern>
    static auto confidence(const Composition<Trait>&         c,
                           size_t                            index,
                           const typename Trait::float_type& q,
                           const Pattern&                    x,
                           const Config&                     cfg) -> typename Trait::float_type
    {
        using float_type = typename Trait::float_type;
        using std::log;
        const auto& pr      = c.models[index];
        const auto[mu, var] = disc::expectation_variance(c.models[index].model, x);
        auto       k        = c.data.num_components();
        auto       ni       = c.data.subset(index).size();
        auto       s        = c.summary.size() - c.data.dim;
        float_type th       = log(cfg.alpha) - log(k) - log(s);
        return students_t_test(mu, std::sqrt(var), q, ni, th);
    }

    template <typename Trait, typename Pattern>
    static auto confidence(const Component<Trait>&           c,
                           const typename Trait::float_type& q,
                           const Pattern&                    x,
                           const Config&                     cfg) -> typename Trait::float_type
    {
        using std::log, std::sqrt;
        using float_type    = typename Trait::float_type;
        const auto& pr      = c.model;
        const auto[mu, var] = disc::expectation_variance(c.model.model, x);
        auto       ni       = c.data.size();
        auto       s        = c.summary.size() - c.data.dim;
        float_type th       = log(cfg.alpha) - log(s);
        return students_t_test(mu, sqrt(var), q, ni, th);
    }
};

} // namespace sd::disc