#pragma once

#include <desc/Composition.hxx>
#include <desc/utilities/ModelPruning.hxx>
#include <elisa/ComplementDist.hxx>
#include <elisa/Elisa.hxx>
#include <elisa/FDR.hxx>

#include <type_traits>

namespace sd::disc
{

// template<class T>
// T gamma_pdf(T x, T a, T b)
// {
//     using std::pow, std::exp, std::tgamma;
//     return pow(a, b) * pow(x, a - 1) * exp(-b * x) / tgamma(a);
// }

template <typename Float_Type,
          typename Statistics = ElisaChernoff,
          typename FDR        = fdr::AlphaSpending<Float_Type>>
struct IElisaFDR : Statistics
{
    IElisaFDR(Float_Type alpha) : fdr(alpha) {}
    IElisaFDR(const IElisaFDR&) = default;
    IElisaFDR(IElisaFDR&&)      = default;
    IElisaFDR& operator=(const IElisaFDR&) = default;
    IElisaFDR& operator=(IElisaFDR&&) = default;
    IElisaFDR()                       = default;

    template <typename C, typename Config>
    void prepare(C& c, const Config& cfg)
    {
        sd::disc::prepare(c, cfg);
        log_alpha_0 = log(cfg.alpha);
        *this       = IElisaFDR<Float_Type, Statistics, FDR>(cfg.alpha);
    }
    template <typename C, typename Candidate, typename Config>
    auto is_allowed(const C& c, const Candidate& x, const Config& cfg) const
    {
        return fdr.alpha_i > 0 && DefaultPatternsetMinerInterface::is_allowed(c, x, cfg);
    }

    template <typename C, typename Candidate, typename Config>
    auto heuristic(C& c, Candidate& x, const Config&) const -> Float_Type
    {
        using std::log;
        return elisa_fdr_heuristic(c, x, -log_alpha_0); //-log(fdr.alpha_i)
        // Float_Type h = elisa_fdr_heuristic(c, x);
        // return h >= -log(fdr.alpha_i) ? h : Float_Type(0);
    }
    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();
        return !statistical_test(p, q, c.data.size(), static_cast<double>(log_alpha_0));
    }

    template <typename C, typename Candidate, typename Config>
    auto insert_into_model(C& c, Candidate& x, const Config& cfg)
    {
        return find_assignment(c, x, cfg, *this);
    }
    template <typename C, typename Config>
    auto finish(C&, const Config&)
    {
        // auto old = fdr;
        // fdr      = FDR(cfg.alpha);
        // prune_composition(c, cfg, *this);
        // fdr = std::move(old);
    }

    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&) -> 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);
    }

    template <typename Trait, typename Pattern>
    auto confidence(const Component<Trait>&           c,
                    const typename Trait::float_type& q,
                    const Pattern&                    x,
                    const Config&) -> typename Trait::float_type
    {
        const auto p = c.model.expectation(x);
        const auto n = c.data.size();
        return statistical_test(p, q, n);
    }

    // protected:
    template <typename T>
    auto statistical_test(const T& p, const T& q, const size_t n) -> T
    {
        using std::log;
        T pv = Statistics::pvalue(p, q, n);
        if (pv > fdr.alpha) return T(0);
        auto r = fdr.test_next(pv);
        return r ? T(-log(pv)) : T(0);
    }

    // template <typename Trait, typename Candidate>
    // static auto
    // elisa_fdr_heuristic_ci(const Component<Trait>& c, const Candidate& x, Float_Type alpha)
    // {
    //     using float_type = typename Trait::float_type;

    //     auto n = c.data.size();
    //     auto q = static_cast<float_type>(x.support) / n;
    //     auto p = c.model.expectation(x.pattern);

    //     // float_type cip = binomial_ci_agresti_coull(static_cast<size_t>(n * p), n, alpha);
    //     // return is_outside_of_ci<float_type>(q, p, cip) ? -Statistics::log_pvalue(p, q, n)
    //     //                                                : float_type(0);
    //     return -Statistics::log_pvalue(p, q, n);
    // }

    // template <typename Trait, typename Candidate>
    // static auto
    // elisa_fdr_heuristic_ci(const Composition<Trait>& c, const Candidate& x, Float_Type alpha)
    // {
    //     using float_type = typename Trait::float_type;

    //     auto k = c.data.num_components();

    //     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);
    //         // float_type cip =
    //         //     binomial_ci_agresti_coull(static_cast<size_t>(n_i * p), n_i, alpha);
    //         // auto h = is_outside_of_ci<float_type>(q, p, cip)
    //         //              ? -Statistics::log_pvalue(p, q, n_i)
    //         //              : float_type(0);
    //         auto h =  -Statistics::log_pvalue(p, q, n_i);
    //         acc = std::max<float_type>(h, acc); // instead of min rho i compute max -
    //         log(roh)
    //     }

    //     return acc;
    // }

    template <typename Trait, typename Candidate>
    static auto elisa_fdr_heuristic(const Component<Trait>& c, const Candidate& x)
    {
        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);
        return -Statistics::log_pvalue(p, q, c.data.size());
    }
    template <typename Trait, typename Candidate>
    static auto
    elisa_fdr_heuristic(const Component<Trait>& c, const Candidate& x, Float_Type th)
    {
        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 r = -Statistics::log_pvalue(p, q, c.data.size());
        return r > th ? r : 0;
    }

    template <typename Trait, typename Candidate>
    static auto elisa_fdr_heuristic(const Composition<Trait>& c, const Candidate& x)
    {
        using float_type = typename Trait::float_type;

        auto k = c.data.num_components();

        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 = -Statistics::log_pvalue(p, q, n_i);
            acc += std::max<float_type>(h, 0);
            // acc = std::max<float_type>(h, acc); // instead of min rho i compute max -
            // log(roh)
        }

        return acc / k;
    }

    template <typename Trait, typename Candidate>
    static auto
    elisa_fdr_heuristic(const Composition<Trait>& c, const Candidate& x, Float_Type th)
    {
        using float_type = typename Trait::float_type;

        auto k = c.data.num_components();

        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 = -Statistics::log_pvalue(p, q, n_i);
            // acc += std::max<float_type>(h, 0);
            acc += h > th ? h : 0;
            // acc = std::max<float_type>(h, acc); // instead of min rho i compute max -
            // log(roh)
        }

        return acc / k;
    }
    Float_Type log_alpha_0 = -2.9957;
    FDR        fdr;
};

template <typename Float_Type,
          typename Statistics = ElisaChernoff,
          typename FDR        = fdr::AlphaSpending<Float_Type>>
struct IElisaFDRI : Statistics
{
    // IElisaFDRI(Float_Type alpha) : fdr(alpha) {}
    IElisaFDRI(Float_Type alpha, size_t num_components = 1)
        : fdrs(num_components, FDR(alpha / num_components))
    {
    }
    IElisaFDRI(const IElisaFDRI&) = default;
    IElisaFDRI(IElisaFDRI&&)      = default;
    IElisaFDRI& operator=(const IElisaFDRI&) = default;
    IElisaFDRI& operator=(IElisaFDRI&&) = default;
    IElisaFDRI()                        = default;

    template <typename C, typename Config>
    void prepare(Component<C>& c, const Config& cfg)
    {
        sd::disc::prepare(c, cfg);
        *this       = IElisaFDRI<Float_Type, Statistics, FDR>(cfg.alpha);
        log_alpha_0 = log(cfg.alpha);
    }
    template <typename C, typename Config>
    void prepare(Composition<C>& c, const Config& cfg)
    {
        sd::disc::prepare(c, cfg);
        *this = IElisaFDRI<Float_Type, Statistics, FDR>(cfg.alpha, c.data.num_components());
        log_alpha_0 = log(cfg.alpha);
    }
    template <typename C, typename Candidate, typename Config>
    auto is_allowed(const C& c, const Candidate& x, const Config& cfg) const
    {
        return fdr.alpha_i > 0 && DefaultPatternsetMinerInterface::is_allowed(c, x, cfg);
    }

    template <typename C, typename Candidate, typename Config>
    auto heuristic(C& c, Candidate& x, const Config&) const -> Float_Type
    {
        using std::log;
        return elisa_fdr_heuristic_th(c, x);
        // Float_Type h = elisa_fdr_heuristic(c, x);
        // return h >= -log(fdr.alpha_i) ? h : Float_Type(0);
    }

    template <typename C, typename Candidate, typename Config>
    auto insert_into_model(C& c, Candidate& x, const Config& cfg)
    {
        return find_assignment(c, x, cfg, *this);
    }
    template <typename C, typename Config>
    auto finish(Component<C>, const Config&)
    {
        // auto old = *this;
        // *this = IElisaFDRI<Float_Type, Statistics, FDR>(cfg.alpha);
        // prune_composition(c, cfg, *this);
        // *this = std::move(old);
    }

    template <typename C, typename Config>
    auto finish(Composition<C>&, const Config&)
    {
        // auto old = *this;
        // *this = IElisaFDRI<Float_Type, Statistics, FDR>(cfg.alpha, c.data.num_components());
        // prune_composition(c, cfg, *this);
        // *this = std::move(old);
    }

    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&) -> 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, index);
    }

    template <typename Trait, typename Pattern>
    auto confidence(const Component<Trait>&           c,
                    const typename Trait::float_type& q,
                    const Pattern&                    x,
                    const Config&) -> typename Trait::float_type
    {
        const auto p = c.model.expectation(x);
        const auto n = c.data.size();
        return statistical_test(p, q, n);
    }

    // protected:
    template <typename T>
    auto statistical_test(const T& p, const T& q, const size_t n) -> T
    {
        using std::log;
        T pv = Statistics::pvalue(p, q, n);
        if (pv > fdrs[0].alpha) return T(0);
        auto r = fdrs[0].test_next(pv);
        return r ? T(-log(pv)) : T(0);
    }
    template <typename T>
    auto statistical_test(const T& p, const T& q, const size_t n, size_t index) -> T
    {
        using std::log;
        T pv = Statistics::pvalue(p, q, n);
        if (pv > fdrs[index].alpha) return T(0);
        auto r = fdrs[index].test_next(pv);
        return r ? T(-log(pv)) : T(0);
    }

    // template <typename Trait, typename Candidate>
    // static auto
    // elisa_fdr_heuristic_ci(const Component<Trait>& c, const Candidate& x, Float_Type alpha)
    // {
    //     using float_type = typename Trait::float_type;

    //     auto n = c.data.size();
    //     auto q = static_cast<float_type>(x.support) / n;
    //     auto p = c.model.expectation(x.pattern);

    //     // float_type cip = binomial_ci_agresti_coull(static_cast<size_t>(n * p), n, alpha);
    //     // return is_outside_of_ci<float_type>(q, p, cip) ? -Statistics::log_pvalue(p, q, n)
    //     //                                                : float_type(0);
    //     return -Statistics::log_pvalue(p, q, n);
    // }

    // template <typename Trait, typename Candidate>
    // static auto
    // elisa_fdr_heuristic_ci(const Composition<Trait>& c, const Candidate& x, Float_Type alpha)
    // {
    //     using float_type = typename Trait::float_type;

    //     auto k = c.data.num_components();

    //     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);
    //         // float_type cip =
    //         //     binomial_ci_agresti_coull(static_cast<size_t>(n_i * p), n_i, alpha);
    //         // auto h = is_outside_of_ci<float_type>(q, p, cip)
    //         //              ? -Statistics::log_pvalue(p, q, n_i)
    //         //              : float_type(0);
    //         auto h =  -Statistics::log_pvalue(p, q, n_i);
    //         acc = std::max<float_type>(h, acc); // instead of min rho i compute max -
    //         log(roh)
    //     }

    //     return acc;
    // }

    template <typename Trait, typename Candidate>
    auto elisa_fdr_heuristic(const Component<Trait>& c, const Candidate& x) const
    {
        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);
        return -Statistics::log_pvalue(p, q, c.data.size());
    }
    template <typename Trait, typename Candidate>
    auto elisa_fdr_heuristic_th(const Component<Trait>& c, const Candidate& x) const
    {
        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 r = -Statistics::log_pvalue(p, q, c.data.size());
        // const auto th = -log(fdrs[0].alpha_i);
        const auto th = -log_alpha_0;
        return r > th ? r : 0;
    }

    template <typename Trait, typename Candidate>
    auto elisa_fdr_heuristic(const Composition<Trait>& c, const Candidate& x) const
    {
        using float_type = typename Trait::float_type;

        auto k = c.data.num_components();

        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 = -Statistics::log_pvalue(p, q, n_i);
            acc += std::max<float_type>(h, 0);
            // acc = std::max<float_type>(h, acc); // instead of min rho i compute max -
            // log(roh)
        }

        return acc / k;
    }

    template <typename Trait, typename Candidate>
    auto elisa_fdr_heuristic_th(const Composition<Trait>& c, const Candidate& x) const
    {
        using float_type = typename Trait::float_type;

        auto k = c.data.num_components();

        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 = -Statistics::log_pvalue(p, q, n_i);
            // acc += std::max<float_type>(h, 0);
            // const auto th = -log(fdrs[k].alpha_i);
            const auto th = -log_alpha_0;

            acc += h > th ? h : 0;
            // acc = std::max<float_type>(h, acc); // instead of min rho i compute max -
            // log(roh)
        }

        return acc / k;
    }
    std::vector<FDR> fdrs;
    Float_Type       log_alpha_0 = -2.9957;
    FDR              fdr;
};

template <typename T>
void un_initialize(Composition<T>& s, const Config& cfg)
{
    s.models.assign(s.data.num_components(), make_distribution(s, cfg));
}
// template <typename T>
// void un_initialize(Component<T>& s, const Config& cfg)
// {
//     s.model = make_distribution(s, cfg);
// }

// template <typename C,
//           typename I    = DefaultPatternsetMinerInterface,
//           typename Info = EmptyCallback>
// void discover_patterns_generic_no_pairs(C& s, const Config& cfg, I fn = {}, Info&& info = {})
// {
//     using patter_type = typename C::pattern_type;
//     using float_type  = typename C::float_type;
//     using generator   = CandidateGenerator<patter_type, float_type>;
//     using clk         = std::chrono::high_resolution_clock;

//     assert(s.data.dim != 0);

//     fn.prepare(s, cfg);

//     info(std::as_const(s));

//     auto score_fn = [&](auto& x) { return fn.heuristic(s, x, cfg); };
//     auto prune_fn = [&](auto& x) { return x.score <= 0 || !fn.is_allowed(s, x, cfg); };

//     auto max_depth = cfg.max_pattern_size.value_or(cfg.max_factor_width);
//     auto gen       = generator(s.data, cfg.min_support, max_depth);

//     auto bak = s;

//     un_initialize(s, cfg);
//     gen.create_singular_candidates(score_fn);

//     if (cfg.search_depth > 1) { gen.expand_bfs(score_fn, prune_fn, cfg.search_depth); }
//     // s = std::move(bak);

//     size_t     items_used = 0;
//     size_t     patience   = cfg.max_patience;
//     const auto start_time = clk::now();

//     for (size_t it = 0; it < cfg.max_iteration && gen.has_next(); ++it)
//     {
//         auto c = gen.next();

//         if (c && fn.insert_into_model(s, *c, cfg))
//         {
//             patience   = std::min(patience * 2, cfg.max_patience);
//             items_used = items_used + 1;

//             info(std::as_const(s));
//         }
//         else if (patience-- == 0)
//             break;

//         if (cfg.max_time && (clk::now() - start_time) > *cfg.max_time) { break; }

//         if (cfg.max_patternset_size && items_used >= *cfg.max_patternset_size) { break; }

//         if (c) { gen.expand_from(*c, score_fn, prune_fn, cfg.search_depth); }
//     }

//     fn.finish(s, cfg);
// }

} // namespace sd::disc