#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 <typename Trait, typename Candidate, typename Interface = DefaultAssignment>
bool elisa_fdr_find_assignment_impl_first(Composition<Trait>& c,
                                          const Candidate&    x,
                                          const Config&       cfg,
                                          Interface&&         f = {})
{
    using float_type = typename Trait::float_type;

    if (x.score <= 0 || !c.models[0].is_allowed(x.pattern)) return false;

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

    auto conf = f.confidence(c, 0, q, x.pattern, cfg);

    if (conf > 0)
    {
        c.models[0].insert(q, x.pattern, true);
        c.assignment[0].insert(c.summary.size());
        c.confidence.push_back(conf);
        insert_pattern_to_summary(c, x);

        return true;
    }
    else
    {
        return false;
    }
}

template <typename Trait, typename Candidate, typename Interface = DefaultAssignment>
bool elisa_fdr_find_assignment(Composition<Trait>& c,
                               const Candidate&    x,
                               const Config&       cfg,
                               Interface&&         f = {})
{
    if (c.data.num_components() == 1)
    {
        return elisa_fdr_find_assignment_impl_first(c, x, cfg, f);
    }
    else
    {
        return find_assignment_impl(c, x, cfg, f);
    }
}

template <typename Trait, typename Candidate, typename Interface = DefaultAssignment>
bool elisa_fdr_find_assignment(Component<Trait>& c,
                               const Candidate&  x,
                               const Config&     cfg,
                               Interface&&       f = {})
{
    using float_type = typename Trait::float_type;

    if (x.score <= 0) return false;

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

    auto conf = f.confidence(c, q, x.pattern, cfg);

    if (conf > 0)
    {
        c.model.insert(q, x.pattern, true);
        c.summary.insert(x.pattern);
        c.frequency.push_back(q);
        c.confidence.push_back(conf);

        return true;
    }
    else
    {
        return false;
    }
}

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

    template <typename C, typename Config>
    void prepare(C& c, const Config& cfg)
    {
        *this = IElisaCI<Float_Type, FDR>(cfg.alpha);
        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
    {
        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& cfg) const -> Float_Type
    {
        return elisa_ci_heuristic(c, x, fdr.alpha_i);
    }
    template <typename C, typename Candidate, typename Config>
    auto insert_into_model(C& c, Candidate& x, const Config& cfg)
    {
        return elisa_fdr_find_assignment(c, x, cfg, *this);
    }
    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();
        // Float_Type r = l2_interval(p, q, n, fdr.alpha_i);
        // fdr.test_next(r > 0 ? 0 : 1);
        bool r = test_ci(q, p, n, fdr.alpha_i); //(p, q, n, fdr.alpha_i);
        if(r > 0) {
            fdr.test_next(0);
        }
        // return r;
        return -ElisaChernoff::log_pvalue(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();
        bool r = test_ci(q, p, n, fdr.alpha_i); //(p, q, n, fdr.alpha_i);
        if(r > 0) {
            fdr.test_next(0);
        }
        // return r;
        return -ElisaChernoff::log_pvalue(p, q, n);
    }

    static auto l2_interval_two(Float_Type q, Float_Type p, size_t n, Float_Type alpha)
        -> Float_Type
    {
        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);
        auto       d   = [](auto x, auto y) {
            using std::sqrt;
            return sqrt((x - y) * (x - y));
        };
        bool overlapping = (q == p) || ((q > p) && (q + ciq) >= (p - cip)) ||
                           ((q > p) && (q - ciq) <= (p + cip));
        return overlapping ? 0 : (q < p ? d(q + ciq, p - cip) : d(q - ciq, p + cip));
    }
    static auto l2_interval(Float_Type q, Float_Type p, size_t n, Float_Type alpha)
        -> Float_Type
    {
        Float_Type cip = binomial_ci_agresti_coull(static_cast<size_t>(n * p), n, alpha);
        auto       d   = [](auto x, auto y) {
            using std::sqrt;
            return sqrt((x - y) * (x - y));
        };

        return !is_outside_of_ci<Float_Type>(q, p, cip)
                   ? 0
                   : (q < p ? d(q, p - cip) : d(q, p + cip));
    }
    static auto test_ci(Float_Type q, Float_Type p, size_t n, Float_Type alpha)
        -> Float_Type
    {
        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);
    }

    template <typename Trait, typename Candidate>
    static auto
    elisa_ci_heuristic(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);
        return test_ci(q, p, n, alpha) ? -ElisaChernoff::log_pvalue(p, q, n) : 0;

        // return l2_interval(q, p, n, alpha);
    }

    template <typename Trait, typename Candidate>
    static auto
    elisa_ci_heuristic(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);
            // auto h = l2_interval(q, p, n_i, alpha);
            auto h = test_ci(q, p, n_i, alpha) ? -ElisaChernoff::log_pvalue(p, q, n_i) : 0;
            acc = std::max<float_type>(h, acc); // instead of min rho i compute max - log(roh)
        }

        return acc;
    }

    FDR fdr;
};

} // namespace sd::disc