#pragma once

#include "Binomial.hxx"
#include "FWER.hxx"
#include "desc/Composition.hxx"
#include "desc/utilities/ModelPruning.hxx"
#include "elisa/ComplementDist.hxx"
#include "math/chi_squared.hxx"
#include <desc/DescMDL.hxx>
#include <elisa/ElisaFDR.hxx>
#include <elisa/FDR.hxx>

#include <type_traits>
#include <vector>

namespace sd::disc
{

namespace fdr
{
template <class T>
struct CwLORD3
{
    using float_type = T;
    float_type alpha = 0.05;
    float_type w0    = 0.005;
    float_type b0    = 0.045;

    float_type alpha_i, w, w_tau;
    size_t     i   = 1;
    size_t     tau = 0;

    explicit CwLORD3(float_type alpha_0 = 0.05, float_type factor = 0.5) : alpha(alpha_0)
    {
        w0      = alpha * factor;
        b0      = alpha - w0;
        alpha_i = discount_xi<T>(1, alpha, b0) * w0;
        w_tau = w = w0 - alpha_i + b0;
    }

    bool test_next(const T& pvalue, T weight)
    {
        if (i == 1)
        {
            alpha_i *= weight;
            w_tau = w = w0 - alpha_i + b0;
        }

        bool reject = pvalue < alpha_i;

        if (reject)
        {
            tau   = i;
            w_tau = w;
        }

        ++i;
        alpha_i = weight * discount_xi<T>(i - tau, alpha, b0) * w_tau;
        w       = w - alpha_i + reject * b0;
        return reject;
    }
};
template <class T>
struct AlphaSpendingw
{
    T alpha_i;

    explicit AlphaSpendingw(T alpha_0 = 0.05) : alpha_i(alpha_0) {}

    bool test_next(const T& pvalue, T weight)
    {
        bool reject = pvalue < weight * alpha_i;
        alpha_i     = std::max<T>(0, alpha_i - pvalue);
        return reject;
    }
};

template <class T>
struct CwLOND
{
    using float_type = T;
    float_type alpha = 0.05;

    T      alpha_i        = 0.05;
    T      denom          = 1;
    size_t i              = 1;
    size_t cnt_rejections = 0;

    explicit CwLOND(float_type alpha_0 = 0.05) : alpha(alpha_0), alpha_i(alpha_0) {}

    bool test_next(const float_type pvalue, const float_type weight)
    {
        alpha_i = weight * alpha * discount_gamma<T>(i) / denom *
                  (std::max<size_t>(cnt_rejections, 1));
        // alpha_i = alpha * discount_gamma<T>(i) / denom * (cnt_rejections + 1);
        bool reject = pvalue < alpha_i;
        cnt_rejections += reject;
        // comment out if independent;
        denom += T(1) / i;
        return reject;
    }
};

} // namespace fdr

// struct ElisaRatio
// {
//     template <typename T>
//     static auto log_pvalue(const T& p, const T& q, const size_t n) -> T
//     {
//     }
//     template <typename T>
//     static auto pvalue(const T& p, const T& q, const size_t n) -> T
//     {

//         return log(q / p) - 2 * ()

//     }
// };

template <typename T = double, typename Candidate>
bool in_independence_ci(const Candidate& x, const Candidate& a, const Candidate& b, size_t n)
{
    auto qa  = T(a.support) / n;
    auto qb  = T(b.support) / n;
    auto cip = binomial_ci_agresti_coull_gz(static_cast<size_t>(n * qa * qb), n, 1.96);
    bool overlapping_ci = is_outside_of_ci<T>(T(x.support) / n, qa * qb, cip);
    return overlapping_ci;
}

template <typename T = double, typename X>
auto log_prior(Component<T> const& c, const X& x)
{
    using std::log;
    typename T::float_type acc = 0; // universal_code(count(x));
    auto                   n   = c.data.size();
    for_each_item(x, [&](size_t item) {
        // auto q = c.frequency[item];
        // acc += c.data.size() * (-q * log(q) - (1 - q) * log(1 - q));
        auto s = c.frequency[item] * n;
        acc -= log2(double(s + (s == 0)) / (n + (s == 0)));
        // acc += universal_code(static_cast<size_t>(c.frequency[item] *n));
        // acc += universal_code(static_cast<size_t>(c.frequency[item] * c.data.size()));
    });
    return acc;
}

template <typename T, typename X>
auto log_prior(const Composition<T>& c, const X& x)
{
    using float_type = typename T::float_type;
    using std::log2;
    float_type acc = 0; // universal_code(count(x)) + universal_code(c.data.num_components());
    for (size_t j = 0; j < c.data.num_components(); ++j)
    {
        float_type l_j = 0;
        for_each_item(x, [&](size_t i) {
            float_type s = c.frequency(i, j) * c.data.subset(j).size();
            l_j -= log2((s + (s == 0)) / (c.data.subset(j).size() + (s == 0)));
        });
        acc = j == 0 ? l_j : std::min(acc, l_j);
    }
    return acc;
}

template <typename Float_Type,
          typename Statistics = ElisaGaussianTT,
          typename FDR        = fdr::CwLORD3<Float_Type>>
struct IElisaCFDR : Statistics
{
    IElisaCFDR(Float_Type alpha) : fdr(alpha) {}
    IElisaCFDR(const IElisaCFDR&) = default;
    IElisaCFDR(IElisaCFDR&&)      = default;
    IElisaCFDR& operator=(const IElisaCFDR&) = default;
    IElisaCFDR& operator=(IElisaCFDR&&) = default;
    IElisaCFDR()                        = default;

    // std::vector<FDR> fdrs;
    // Float_Type mu = 0;
    // std::vector<Float_Type> supp;

    template <typename C, typename Config>
    void prepare(C& c, const Config& cfg)
    {
        *this       = IElisaCFDR<Float_Type, Statistics, FDR>(cfg.alpha);
        log_alpha_0 = log(cfg.alpha);
        sd::disc::prepare(c, cfg);

        // supp = compute_singleton_supports(c.data.dim, c.data);
        // mu = 0;
        // size_t n = c.data.size();
        // auto lnp1 =  log(n + 1);
        // for(auto && s : supp) {
        //     mu += log(s + 1) - lnp1;
        // }

        // fdrs.resize(c.data.num_components());
        // for(size_t i = 0; i < c.data.num_components(); ++i)
        // {
        //     fdrs[i] = FDR(cfg.alpha * Float_Type(c.data.subset(i).size()) / c.data.size());
        // }
    }
    template <typename C, typename Candidate, typename Config>
    auto is_allowed(const C& c, const Candidate& x, const Config& cfg) const
    {
        return x.score > 0 && fdr.alpha_i > 0 &&
               DefaultPatternsetMinerInterface::is_allowed(c, x, cfg);
    }
    template <typename C, typename Candidate, typename Config>
    auto heuristic(const C& c, const Candidate& x, const Config& cfg) const -> Float_Type
    {
        using std::log;
        // Float_Type h = elisa_fdr_heuristic(c, x, fdr.alpha_i);
        Float_Type h = elisa_fdr_heuristic(c, x, cfg.alpha * weight(c, x.pattern));
        return h >= -log(fdr.alpha_i) ? h : Float_Type(0);

        // // std::cout << h << "\t" << (-log_alpha_0 - log_weight(c, x.pattern)) << std::endl;
        // return h >= (-log_alpha_0 - log_weight(c, x.pattern)) ? h : Float_Type(0);
    }

    template <typename C, typename Candidate, typename Config>
    auto heuristic(const C&         c,
                   const Candidate& x,
                   const Candidate& a,
                   const Candidate& b,
                   const Config&    cfg) const
    {
        return heuristic(c, x, cfg);
        // puts("yay");
        // using std::log;
        // Float_Type h = elisa_fdr_heuristic2(c, x, a, b, fdr.alpha_i * weight(c, x.pattern));
        // // std::cout << h << "\t" << (-log_alpha_0 - log_weight(c, x.pattern)) << std::endl;
        // return h >= (-log(fdr.alpha_i) - log_weight(c, x.pattern)) ? h : Float_Type(0);
    }

    template <typename C, typename Candidate, typename Config>
    auto insert_into_model(C& c, Candidate& x, const Config& cfg)
    {
        x.score = heuristic(c, x, cfg);
        if (x.score <= 0) return false;

        auto assigned = elisa_fdr_find_assignment(c, x, cfg, *this);

        // auto h1 = heuristic(c, x, cfg);
        // std::cout << x.score << " " << h1 << " " << assigned << std::endl;

        return assigned;
    }
    template <typename C, typename Config>
    auto finish(C& c, const Config& cfg)
    {
        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 this->statistical_test(p, q, n, weight(c, x));

        // using std::log;
        // auto w = weight(c, x);
        // auto    pv = Statistics::pvalue(p, q, n);

        // auto    ci = binomial_ci_agresti_coull(static_cast<size_t>(n * p), n,
        // fdrs[index].alpha_i); bool r  = is_outside_of_ci(q, p, ci) && fdr.test_next(pv, w);
        // return r ? Float_Type(-log(pv)) : Float_Type(0);
    }

    template <class C, class X, typename T>
    void print_me(const C& c, const X& x, T p, T q, size_t n) const
    {
#if 0
        std::cout << p << '\t' 
                  << q << '\t' 
                  << Statistics::pvalue(p, q, n) << '\t'
                  << -Statistics::log_pvalue(p, q, n) << '\t' 
                  << count(x) << '\t'
                  << exp(log_weight_geo(c, x)) << '\t' 
                  << exp(log_weight_exp(c, x)) << '\t'
                  << weight(c, x, static_cast<size_t>(q * n)) << '\t' 
                  << weight(c, x) << '\t'
                  << log_weight(c, x) << '\t' 
                  << std::endl;
#endif
    }
    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 this->template statistical_test<Float_Type>(p, q, n, weight(c, x));
    }

protected:
    template <class C, class X>
    Float_Type log_weight_geo(const C& c, const X& x) const
    {
        Float_Type k = count(x);
        Float_Type N = c.data.size();
        return -k * (log(N) - log(k) + 1);
    }
    template <class C, class X>
    Float_Type log_weight_exp(const C& c, const X& x) const
    {
        return -Float_Type(count(x));
    }
    template <class C, class X>
    Float_Type log_weight_exp2(const C& c, const X& x) const
    {
        return -Float_Type(count(x));
    }
    template <class C, class X>
    Float_Type log_weight(const C& c, const X& x) const
    {
        // DeltaMDL weight:
        return -log_prior(c, x);

        // Float_Type w = 0;
        // size_t n = c.data.size();
        // for_each_item(x, [&](auto i) {
        //     w += log(supp[i] + 1) / (n + 1);
        // });
        // return w - log(mu);
    }

    template <class C, class X>
    Float_Type log_weight(const C& c, const X& x, size_t s) const
    {
        return -log_prior(c, x) - additional_cost_mdl(s);
    }

    template <class C, class X>
    Float_Type log_weightb(const C& c, const X& x, size_t s) const
    {
        return reg_ibeta<Float_Type>(count(x), 0.5, 0.5);
        // return -log_prior(c, x) - additional_cost_mdl(s);
    }

    template <class C, class X>
    Float_Type weight(const C& c, const X& x, size_t s) const
    {
        using std::exp;
        return exp(log_weight(c, x, s));
    }

    template <class C, class X>
    Float_Type weight(const C& c, const X& x) const
    {
        using std::exp;
        return exp(log_weight(c, x));
    }

    template <typename T>
    auto statistical_test(T p, T q, size_t n, T w) -> T
    {
        using std::log;
        T    pv = Statistics::pvalue(p, q, n);
        T    ci = binomial_ci_agresti_coull<T>(static_cast<size_t>(n * p), n, fdr.alpha_i);
        bool r  = is_outside_of_ci(q, p, ci) && fdr.test_next(pv, w);
        // std::cout << q << '\t' << p << '\t' << ci << '\t' << -Statistics::log_pvalue(p, q, n)
        //           << '\t' << -log(fdr.alpha_i * w) << '\t' << w << '\t' << std::endl;
        return r ? T(-log(pv)) : T(0);
    }

    template <typename Trait, typename Candidate>
    auto elisa_fdr_heuristic(const Component<Trait>&    c,
                             const Candidate&           x,
                             typename Trait::float_type alpha) const
    {
        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);

        // auto ci = binomial_ci_agresti_coull_gz<float_type>(static_cast<size_t>(n * p), n,
        // gauss_ppf<T>(1. - alpha / 2, 0, 1000));

        // auto cip             = binomial_ci_agresti_coull(static_cast<size_t>(n * p), n,
        // alpha);

        print_me(c, x.pattern, p, q, n);

        // float_type w_alpha = alpha * weight(c, x.pattern);
        float_type cip = binomial_ci_agresti_coull(static_cast<size_t>(n * p), n, alpha);
        // float_type ciq = binomial_ci_agresti_coull(static_cast<size_t>(n * q), n, alpha);
        // bool       overlapping_ci = (q == p) || ((q > p) && (q + ciq) >= (p - cip)) ||
        //                       ((q > p) && (q - ciq) <= (p + cip));
        bool overlapping_ci = is_outside_of_ci<float_type>(q, p, cip);

        // std::cout << q << " p/m " << ciq << '\t' << p << " p/m " << cip << '\t'
        //           << -Statistics::log_pvalue(p, q, n) << '\t' <<
        //           c.model.probability(x.pattern)
        //           << std::endl;

        auto pv = -Statistics::log_pvalue(p, q, n); // + log_weight(c, x.pattern);

        // std::min<float_type>(-Statistics::log_pvalue(p, q, n), -Statistics::log_pvalue((1-q),
        // (1-p), n))

        return overlapping_ci ? pv : float_type(0);
    }

    template <typename Trait, typename Candidate>
    auto elisa_fdr_heuristic(const Composition<Trait>&  c,
                             const Candidate&           x,
                             typename Trait::float_type alpha) 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);

            float_type w_alpha = (alpha / k) * weight(c, x.pattern, size_t(q * n_i));

            auto       h = -Statistics::log_pvalue(p, q, n_i); // + log_weight(c, x.pattern);
            float_type cip =
                binomial_ci_agresti_coull(static_cast<size_t>(n_i * p), n_i, w_alpha);
            // auto ciq = binomial_ci_agresti_coull(static_cast<size_t>(n_i * q), n_i, alpha);
            // bool overlapping_ci = (q == p) || (q > p) && (q + ciq) >= (p - cip) ||
            //                       (q > p) && (q - ciq) <= (p + cip);
            bool overlapping_ci = is_outside_of_ci<float_type>(q, p, cip);

            if (!overlapping_ci) { acc = std::max<float_type>(h, acc); }
        }

        return acc;
    }

    template <typename Trait, typename Candidate>
    auto elisa_fdr_heuristic2(const Component<Trait>&    c,
                              const Candidate&           x,
                              const Candidate&           a,
                              const Candidate&           b,
                              typename Trait::float_type alpha) const
    {
        using float_type = typename Trait::float_type;

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

        auto qa   = static_cast<float_type>(a.support) / n;
        auto qb   = static_cast<float_type>(b.support) / n;
        auto ciab = binomial_ci_agresti_coull_gz(static_cast<size_t>(n * qa * qb), n, 1.96);
        auto p    = c.model.expectation(x.pattern);

        //  std::cout << -Statistics::log_pvalue(p, q, n) << '\t' << - log_weight(c, x.pattern)
        //               << std::endl;

        if (!is_outside_of_ci<float_type>(q, qa * qb, ciab)) return float_type(0);

        // std::cout << -Statistics::log_pvalue(p, q, n) << '\t' << -log_weight(c, x.pattern)
        //           << std::endl;
        // auto ci = binomial_ci_agresti_coull_gz<float_type>(static_cast<size_t>(n * p), n,
        // gauss_ppf<T>(1. - alpha / 2, 0, 1000));
        auto ciq            = binomial_ci_agresti_coull(static_cast<size_t>(n * q), n, alpha);
        auto cip            = binomial_ci_agresti_coull(static_cast<size_t>(n * p), n, alpha);
        bool overlapping_ci = (q == p) || ((q > p) && (q + ciq) >= (p - cip)) ||
                              ((q > p) && (q - ciq) <= (p + cip));

        auto pv = -Statistics::log_pvalue(p, q, n); // + log_weight(c, x.pattern);

        // std::min<float_type>(-Statistics::log_pvalue(p, q, n), -Statistics::log_pvalue((1-q),
        // (1-p), n))

        return overlapping_ci ? pv : float_type(0);
    }

    template <typename Trait, typename Candidate>
    auto elisa_fdr_heuristic2(const Composition<Trait>&  c,
                              const Candidate&           x,
                              const Candidate&           a,
                              const Candidate&           b,
                              typename Trait::float_type alpha) 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 qa =
                static_cast<float_type>(size_of_intersection(a.row_ids, c.masks[i])) / n_i;
            auto qb =
                static_cast<float_type>(size_of_intersection(b.row_ids, c.masks[i])) / n_i;
            auto ciab =
                binomial_ci_agresti_coull_gz(static_cast<size_t>(n_i * qa * qb), n_i, 1.96);
            if (!is_outside_of_ci<float_type>(q, qa * qb, ciab)) continue;

            auto p = c.models[i].expectation(x.pattern);
            auto h = -Statistics::log_pvalue(p, q, n_i); // + log_weight(c, x.pattern);

            auto ciq = binomial_ci_agresti_coull(static_cast<size_t>(n_i * q), n_i, alpha);
            auto cip = binomial_ci_agresti_coull(static_cast<size_t>(n_i * p), n_i, alpha);
            bool overlapping_ci = (q == p) || ((q > p) && (q + ciq) >= (p - cip)) ||
                                  ((q > p) && (q - ciq) <= (p + cip));

            if (!overlapping_ci) { acc = std::max<float_type>(h, acc); }
        }

        return acc;
    }

    Float_Type log_alpha_0 = -2.9957;
    FDR        fdr;
};

} // namespace sd::disc