#pragma once

#include <desc/PatternsetMiner.hxx>
#include <desc/utilities/ModelPruning.hxx>
#include <elisa/FWER.hxx>
#include <elisa/Pvalues.hxx>

namespace sd::disc
{

// template <typename T>
// auto binomial_log_pvalue(const T& p, const T& q, const size_t n) -> T
// {

//     if (n < 100)
//     {
//         // using std::log;
//         // return q < p ? T(0) : log(binomial_cdf_exact(static_cast<size_t>(q * n), p, n));
//         return log(binomial_cdf_exact(static_cast<size_t>(q * n), p, n));
//     }
//     else
//     {
//         return binomial_log_cdf_chernoff(static_cast<size_t>(q * n), p, n);
//         // return q < p ? T(0) : (T(n) * -kl_bernoulli(q, p));
//     }
// }

// template <typename T>
// auto binomial_pvalue(const T& p, const T& q, const size_t n) -> T
// {

//     if (n < 100)
//     {
//         // using std::log;
//         // return q < p ? T(0) : log(binomial_cdf_exact(static_cast<size_t>(q * n), p, n));
//         return binomial_cdf_exact(static_cast<size_t>(q * n), p, n);
//     }
//     else
//     {
//         return binomial_cdf_chernoff(static_cast<size_t>(q * n), p, n);
//         // return q < p ? T(0) : (T(n) * -kl_bernoulli(q, p));
//     }

//     // if(n < 100)
//     // {
//     //     return q < p ? T(1) : binomial_cdf_exact(static_cast<size_t>(q * n), p, n);
//     // }
//     // else
//     // {
//     //     return q < p ? T(0) : std::exp(T(n) * -kl_bernoulli(q, p));
//     // }
// }

// template <typename T>
// auto binomial_test(const T& p, const T& q, const size_t n, const T& log_alpha) -> T
// {
//     // if (q <= p)
//     //     return T(0); // insignificant.
//     auto h = -binomial_log_pvalue(p, q, n);
//     return h <= -log_alpha ? T(0) : h; // rejects hypothesis if p-value > a
// }

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

//     auto q  = static_cast<float_type>(x.support) / c.data.size();
//     auto p  = c.model.expectation(x.pattern);
//     auto th = fwer::log_alpha_correction<float_type>(cfg.alpha, layer, c.data.dim);
//     return binomial_test(p, q, c.data.size(), th);
// }

// template <typename Trait, typename Candidate>
// auto elisa_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; auto p   = c.models[i].expectation(x.pattern); auto h   = binomial_test(p, q,
//         n_i, th); acc      = std::max<float_type>(h, acc); // instead of min rho i compute
//         max - log(roh)
//     }

//     return acc;
// }
/*
template <typename Float_Type>
struct IElisa
{
    IElisa(Float_Type alpha, size_t d) : layer(1), log_alpha(fwer::log_alpha_correction(alpha,
1, d)) {} IElisa(const IElisa&) = default; IElisa(IElisa&&)      = default; IElisa&
operator=(const IElisa&) = default; IElisa& operator=(IElisa&&) = default; IElisa() = default;

    template <typename C, typename Config>
    void prepare(C& c, const Config& cfg)
    {
        *this = IElisa<Float_Type>(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) const
    {
        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_heuristic(c, x, cfg, std::max(count(x.pattern), this->layer));
    }

    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 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 p = c.models[index].expectation(x);
        const auto n = c.data.subset(index).size();
        return binomial_test(p, 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 p = c.model.expectation(x);
        const auto n = c.data.size();
        return binomial_test(p, 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;
    Float_Type log_alpha = 1;
};
*/

// template<typename Statistics>
// struct StatisticalAssignment
// {
//     size_t layer = 1;

//     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&                     cfg) const -> typename
//                     Trait::float_type
//     {
//         using float_type = typename Trait::float_type;

//         const auto p = c.models[index].expectation(x);

//         auto a  = cfg.alpha / c.data.num_components();
//         auto ni = c.data.subset(index).size();
//         auto th = fwer::log_alpha_correction<float_type>(a, layer, c.data.dim);

//         return binomial_test(p, q, ni, th);
//     }

//     template <typename Trait, typename Pattern>
//     auto confidence(const Component<Trait>&           c,
//                     const typename Trait::float_type& q,
//                     const Pattern&                    x,
//                     const Config&                     cfg) const -> typename
//                     Trait::float_type
//     {
//         using float_type = typename Trait::float_type;

//         const auto& pr = c.model;
//         const auto  p  = pr.expectation(x);

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

//         return binomial_test(p, q, ni, th);
//     }
// };

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

    template <typename C, typename Config>
    void prepare(Component<C>& c, const Config& cfg)
    {
        sd::disc::prepare(c, cfg);
        *this = IElisa<Float_Type, Statistics>(cfg.alpha, c.data.dim);
    }
    template <typename C, typename Config>
    void prepare(Composition<C>& c, const Config& cfg)
    {
        sd::disc::prepare(c, cfg);
        *this = IElisa<Float_Type, Statistics>(cfg.alpha / c.data.num_components(), c.data.dim);
    }
    template <typename C, typename Candidate, typename Config>
    auto is_allowed(const C& c, const Candidate& x, const Config& cfg) const
    {
        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_heuristic(c, x, cfg, std::max(count(x.pattern), this->layer));
    }

    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;
        else
            return find_assignment(c, x, cfg, *this);
    }

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

    // template <class C, class S, class T>
    // bool acceptable_decay(const C& c, const S& s, const T& t)
    // {
    //     auto q  = double(s.support) / c.data.size();
    //     auto p  = double(t.support) / c.data.size();
    //     auto pv = (1 - Statistics::pvalue(p, q, c.data.size()));
    //     std::cout << q << " " << p << " " << pv << " " << (pv < 0.05) << std::endl;
    //     return pv < 0.05;
    // }

    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 p = c.models[index].expectation(x);
        const auto n = c.data.subset(index).size();
        return statistical_test(p, 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 p = c.model.expectation(x);
        const auto n = c.data.size();
        return statistical_test(p, q, n, static_cast<float_type>(log_alpha));
    }

    // protected:
    template <typename T>
    static auto statistical_test(const T& p, const T& q, size_t n, const T& log_alpha) -> T
    {
        auto h = -Statistics::log_pvalue(p, q, n);
        return h <= -log_alpha ? T(0) : h;
    }

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

        auto q  = static_cast<float_type>(x.support) / c.data.size();
        auto p  = c.model.expectation(x.pattern);
        auto th = fwer::log_alpha_correction<float_type>(cfg.alpha, layer, c.data.dim);
        return statistical_test(p, q, c.data.size(), th);
    }

    template <typename Trait, typename Candidate>
    static auto elisa_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 / float_type(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;
            auto p = c.models[i].expectation(x.pattern);
            auto h = statistical_test(p, q, n_i, th);
            // instead of min rho i compute max - log(roh)
            // acc = std::max<float_type>(h, acc);
            acc += h;
        }

        return acc / k;
    }

    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;
    Float_Type log_alpha = 1;
};

// template <typename Float_Type, typename Statistics = ElisaBinomial>
// struct IElisa_NX : Statistics
// {

//     IElisa_NX(Float_Type alpha, size_t d)
//         : layer(1), log_alpha(fwer::log_alpha_correction(alpha, 1, d))
//     {
//     }
//     IElisa_NX(const IElisa_NX&) = default;
//     IElisa_NX(IElisa_NX&&)      = default;
//     IElisa_NX& operator=(const IElisa_NX&) = default;
//     IElisa_NX& operator=(IElisa_NX&&) = default;
//     IElisa_NX()                    = default;

//     template <typename C, typename Config>
//     void prepare(Component<C>& c, const Config& cfg)
//     {
//         sd::disc::prepare(c, cfg);
//         *this = IElisa_NX<Float_Type, Statistics>(cfg.alpha, c.data.dim);
//     }
//     template <typename C, typename Config>
//     void prepare(Composition<C>& c, const Config& cfg)
//     {
//         sd::disc::prepare(c, cfg);
//         *this = IElisa_NX<Float_Type, Statistics>(cfg.alpha / c.data.num_components(),
//         c.data.dim);
//     }
//     template <typename C, typename Candidate, typename Config>
//     auto is_allowed(const C& c, const Candidate& x, const Config& cfg) const
//     {
//         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_heuristic(c, x, cfg, std::max(count(x.pattern), this->layer));
//     }

//     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;
//         else
//             return 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 p = c.models[index].expectation(x);
//         const auto n = c.data.subset(index).size();
//         return statistical_test(p, 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 p = c.model.expectation(x);
//         const auto n = c.data.size();
//         return statistical_test(p, q, n, static_cast<float_type>(log_alpha));
//     }

// // protected:
//     template <typename T>
//     static auto statistical_test(const T& p, const T& q, size_t n, const T& log_alpha) -> T
//     {
//         auto h = -Statistics::log_pvalue(p, q, n);
//         return h <= -log_alpha ? T(0) : h;
//     }

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

//         auto q  = static_cast<float_type>(x.support) / c.data.size();
//         auto p  = c.model.expectation(x.pattern);
//         auto th = fwer::log_alpha_correction_tight<float_type>(cfg.alpha, layer, c.data.dim);
//         return statistical_test(p, q, c.data.size(), th);
//     }

//     template <typename Trait, typename Candidate>
//     static auto elisa_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_tight<float_type>(cfg.alpha / float_type(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; auto p = c.models[i].expectation(x.pattern); auto h = statistical_test(p, q,
//             n_i, th);
//             // instead of min rho i compute max - log(roh)
//             // acc = std::max<float_type>(h, acc);
//             acc += h;
//         }

//         return acc / k;
//     }

//     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_tight(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_tight(cfg.alpha / c.data.num_components(), layer,
//             c.data.dim);
//     }

//     size_t     layer     = 1;
//     Float_Type log_alpha = 1;
// };

template <typename Statistics = ElisaBinomial>
struct StatisticalCharacterization
{
    template <typename T>
    static T statistical_test(const T& p, const T& q, size_t n, const T& log_alpha)
    {
        auto h = -Statistics::log_pvalue(p, q, n);
        return h <= -log_alpha ? T(0) : h; // rejects hypothesis if p-value > a
    }

    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  p  = pr.expectation(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 statistical_test(p, 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 float_type = typename Trait::float_type;
        using std::log;

        const auto& pr = c.model;
        const auto  p  = pr.expectation(x);

        auto ni = c.data.size();
        auto s  = c.summary.size() - c.data.dim;

        float_type th = log(cfg.alpha); // - log(s);

        return statistical_test(p, q, ni, th);
    }
};

} // namespace sd::disc