#pragma once

#include <elisa/FWER.hxx>
#include <relent/Reaper.hxx>

#include <desc/distribution/MaxEntFactor.hxx>
#include <math/IncrementalStatistics.hxx>

#include <cmath>

namespace sd::disc
{

template <typename Transactions, typename Model, typename Pattern>
auto compute_expectation_variance(Transactions const& transactions,
                                  size_t              len,
                                  Model const&        model,
                                  Pattern const&      x)
{
    using float_type = typename Model::float_type;
    IncrementalDescription<float_type> p;
    for (size_t i = 0; i < len; ++i) {
        const auto& t = transactions[i];
        if (t.value != 0 && is_subset(x, t.cover)) {
            p += t.value * probability(model, t.cover);
            assert(0 <= p.sum() && p.sum() <= 1 && p.sum() == p.sum());
        }
    }
    assert(0 <= p.sum() && p.sum() <= 1);
    return std::pair<float_type, float_type>{p.sum(),
                                             p.variance() * (float_type(p.size()) - 1)};
}

template <typename S, typename T>
std::pair<T, T> expectation_variance_unknown(MaxEntFactor<S, T> const& model,
                                             itemset<S> const&         x)
{
    thread_local TempPartitionBuffer<S, T, 13> bf;

    auto& b   = bf.get(model.itemsets.set.size() + 1);
    auto  len = make_partitions_for_unknown(b, model, x);
    return compute_expectation_variance(b, len, model, x);
}

template <typename S, typename T>
std::pair<T, T> expectation_variance(MaxEntFactor<S, T> const& model, itemset<S> const& x)
{
    if (auto p = model.get_precomputed_expectation(x); p) {
        return {p.value(), 1};
    }
    else
    {
        return expectation_variance_unknown(model, x);
    }
}

template <typename M, typename X, class T = typename M::float_type>
auto expectation_variance_impl(const M& pr, const X& x) -> std::pair<T, T>
{
    // using float_type   = typename M::float_type;
    using pattern_type = typename M::pattern_type;

    thread_local itemset<pattern_type> part, x_buf;

    x_buf.assign(x);

    T p = 1, mm = 1, var = 1;

    factorize(pr, x, [&](const auto& phi_i, size_t, bool singleton) {
        T mu = 0, s2 = 0.5;
        if (singleton) {
            mu = phi_i.factor.singletons.set.front().probability;
        }
        else
        {
            part.clear();
            intersection(x, phi_i.range, part);
            std::tie(mu, s2) = expectation_variance(phi_i.factor, part);
        }
        p *= mu;
        mm *= mu * mu;
        var *= mu * mu + s2; // numerically very unstable
    });

    var -= mm;

    return {p, var};
}
template <typename S, typename T, typename X>
std::pair<T, T> expectation_variance(const ClassicStaticFactorModel<S, T>& fm, const X& x)
{
    return expectation_variance_impl(fm, x);
}
template <typename S, typename T, typename X>
std::pair<T, T> expectation_variance(const StaticFactorModel<S, T>& fm, const X& x)
{
    return expectation_variance_impl(fm, x);
}
template <typename S, typename T, typename X>
std::pair<T, T> expectation_variance(const RelaxedFactorModel<S, T>& fm, const X& x)
{
    return expectation_variance_impl(fm, x);
}

// template <typename pattern_type, typename float_type, typename Underlying, typename
// query_type> std::pair<float_type, float_type> expectation_variance(const
// StaticFactorModel<pattern_type, float_type, Underlying>& fm,
//                      const query_type&                                              t)
// {
//     thread_local itemset<pattern_type> part;
//     part.clear();
//     part.reserve(fm.dim);

//     float_type p = 1, mm = 1, var = 1;
//     for (size_t i = 0; i < fm.factors.size(); ++i)
//     {
//         const auto& f = fm.factors[i];
//         if (intersects(t, f.range))
//         {
//             intersection(t, f.range, part);
//             auto [mu, s2] = expectation_variance(f.factor, part);
//             p *= mu;
//             mm *= mu * mu;
//             var *= mu * mu + s2; // numerically very unstable
//         }
//     }
//     var -= mm;

//     return {p, var};
// }

// template <typename pattern_type, typename float_type, typename W, typename query_type>
// std::pair<float_type, float_type>
// expectation_variance(const RelaxedFactorModel<pattern_type, float_type, W>& pr,
//                      const query_type&                                      x)
// {

//     thread_local itemset<pattern_type> part, x_buf;

//     x_buf.assign(x);

//     float_type p = 1, mm = 1, var = 1;

//     factorize_model_itemsets(pr, pr.dim, x_buf, [&](size_t index) {
//         part.clear();
//         intersection(x, pr.factors[index].range, part);
//         auto [mu, s2] = expectation_variance(pr.factors[index].factor, part);
//         p *= mu;
//         mm *= mu * mu;
//         var *= mu * mu + s2; // numerically very unstable
//     });
//     factorize_model_singletons(pr, pr.dim, x_buf, [&](size_t index) {
//         auto mu = pr.factors[index].factor.singletons.set.front().probability;
//         auto s2 = 1;
//         p *= mu;
//         mm *= mu * mu;
//         var *= mu * mu + s2; // numerically very unstable
//     });
//     var -= mm;

//     return {p, var};
// }

// template <typename float_t>
// auto ztest_pvalue(const float_t& mu,
//                   const float_t& var,
//                   const float_t& q,
//                   const size_t   sample_size) -> float_t
// {
//     using std::abs;
//     using std::erf;
//     using std::sqrt;

//     // static constexpr auto eps = float_t(1e-4);

//     const auto cdf = [mu = 0, sig = 1](auto x) {
//         return float_t(0.5) * (1 + erf((x - mu) / (sig * 1.4142)));
//     };
//     // if(abs(q - mu) < 1.96 * sqrt(var) / sqrt(sample_size)) return float_t(0);
//     // if (abs(q - mu) < float_t(0.2) / sqrt(data_size)) { return 0; } // CLT
//     const auto z = (q - mu) / sqrt(var);
//     // auto pval = 2*cdf(-abs(z  * sample_size)); // two-sided test
//     auto pval = cdf(-z * sample_size); // one-sided test

//     // std::cout << pval << " " << mu << " " << var << " " << z << std::endl;
//     return pval;
// }
// template <typename T>
// auto ztest_log_pvalue_tt(const T& mu, const T& var, const T& q, const size_t sample_size) ->
// T
// {
//     return log(ztest_pvalue_tt(mu, var, q, sample_size));
// }

template <typename T>
auto elisa_ci(const T& mu, const T& var, const T& alpha, const size_t n) -> std::pair<T, T>
{
    using std::abs, std::sqrt;
    T w = abs(standard_normal_quantile(1 - alpha / 2)) * sqrt(var / n);
    return {mu - w, mu + w};
}

template <typename T>
auto ztest_pvalue_tt(const T& mu, const T& var, const T& q, const size_t n) -> T
{
    using std::abs, std::sqrt;
    return 2 * standard_normal_cdf(-abs(n * (q - mu) / (sqrt(var / n))));
}
template <typename T>
auto ztest_log_pvalue_tt(const T& mu, const T& var, const T& q, const size_t sample_size) -> T
{
    return log(ztest_pvalue_tt(mu, var, q, sample_size));
}

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

template <typename Trait, typename Candidate>
auto elisa_ztest_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();
    const auto[mu, var] = expectation_variance(c.model.model, x.pattern);
    auto th             = fwer::log_alpha_correction<float_type>(cfg.alpha, layer, c.data.dim);
    return ztest(mu, var, q, c.data.size(), th);
}

template <typename Trait, typename Candidate>
auto elisa_ztest_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 s        = size_of_intersection(x.row_ids, c.masks[i]);
        auto q        = static_cast<float_type>(s) / n_i;
        auto[mu, var] = expectation_variance(c.models[i].model, x.pattern);

        acc = std::max<float_type>(acc, ztest(mu, var, q, n_i, th));
    }

    return acc;
}

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

    template <typename C, typename Config>
    void prepare(C& c, const Config& cfg)
    {
        *this = IElisaZ<T>(cfg.alpha, c.data.dim);
        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_ztest_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 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] = expectation_variance(c.models[index].model, x);
        const auto n        = c.data.subset(index).size();
        return ztest(mu, 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] = expectation_variance(c.model.model, x);
        const auto n        = c.data.size();
        return ztest(mu, 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 ElisaZCharacterization
{
    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] = 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 ztest(mu, 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 float_type = typename Trait::float_type;
        using std::log;

        const auto& pr      = c.model;
        const auto[mu, var] = 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 ztest(mu, var, q, ni, th);
    }
};

// template <typename D, typename Candidate, typename float_t>
// auto z_test_bin_score(const Distribution<D>& pr,
//                       const Candidate&       x,
//                       const float_t&            q,
//                       const size_t,
//                       const size_t data_size,
//                       float_t         alpha = float_t(0.05)) -> float_t
// {
//     using std::abs;
//     using std::erf;
//     using std::sqrt;

//     const auto p = expectation(pr.model, x.pattern);

//     if (q <= p)
//         return 0;

//     static constexpr auto sqrt_2 = 1.4142; // sqrt(float_t(2));
//     const auto            cdf    = [mu = 0, sig = 1](auto x) {
//         return float_t(0.5) * (1 + erf((x - mu) / (sig * sqrt_2)));
//     };

//     const auto z  = (data_size * (q - p) + 0.5) / sqrt(data_size * p * (1 - p));
//     const auto r  = -lg_fwer_geo<float_t>(alpha, count(x.pattern), pr.model.dim);
//     const auto pv = abs(1 - (cdf(z) - 1e-20));
//     const auto h  = -log(pv);

//     printf("q = %lf\tp = %lf\tz = %lf\tpv = %1.20lf\tlogpv = %2.10lf\tlog_th = %2.10lf\n",
//            (double)q * data_size,
//            (double)p * data_size,
//            (double)z,
//            (double)pv,
//            (double)h,
//            (double)r);
//     return h - r;
// }

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

//     auto fr              = static_cast<float_type>(x.support) / c.data.size();
//     const auto [mu, var] = expectation_variance(c.model.model, x.pattern);
//     return 0.5 * log2(2 * 3.14159 * 2.71828 * var) * fr;
// }

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

//     float_type acc = 0;
//     for (size_t i = 0; i < c.data.num_components(); ++i)
//     {
//         auto support         = size_of_intersection(x.row_ids, c.masks[i]);
//         auto fr              = static_cast<float_type>(support) / c.data.size();
//         const auto [mu, var] = expectation_variance(c.models[i].model, x.pattern);
//         acc += 0.5 * log2(2 * 3.14159 * 2.71828 * var) * fr;
//     }

//     return acc;
// }

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

//     auto fr = static_cast<float_type>(x.support) / c.data.size();
//     return z_test_pvalue(c.model, x, fr, x.support, c.data.size()) > 0.95;
// }

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

//     for (size_t i = 0; i < c.data.num_components(); ++i)
//     {
//         auto support = size_of_intersection(x.row_ids, c.masks[i]);
//         auto q       = static_cast<float_type>(support) / c.data.subset(i).size();

//         if (z_test_pvalue(c.models[i], x, q, support, c.data.size()) > 0.95)
//             return true;
//     }

//     return false;
// }

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

//     auto fr = static_cast<float_type>(x.support) / c.data.size();
//     return z_test_score(c.model, x, fr, x.support, c.data.size(), float_type(0.99999999999));
// }

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

//     float_type acc = 0;
//     for (size_t i = 0; i < c.data.num_components(); ++i)
//     {
//         auto support = size_of_intersection(x.row_ids, c.masks[i]);
//         auto fr      = static_cast<float_type>(support) / c.data.size();
//         acc += z_test_score(
//             c.models[i], x, fr, support, c.data.subset(i).size(), float_type(0.99999999999));
//     }

//     return acc;
// }

// template <typename Trait, typename Candidate>
// auto usage_z_test_heur(const Component<Trait>& c, const Candidate& x)
// {
//     using float_type = typename Trait::float_type;
//     auto fr          = static_cast<float_type>(x.support) / c.data.size();

//     float_type alpha(0.05);

//     if constexpr (!is_dynamic_factor_model<typename Trait::distribution_type>())
//     {
//         return z_test_bin_score(c.model, x, fr, x.support, c.data.size(), alpha);
//     }
//     else
//     {
//         auto u = usage_of_new_factor(x.pattern, c.data, c.model, x.row_ids);
//         return z_test_bin_score(c.model, x, fr, u, c.data.size(), alpha);
//     }
// }

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

//     thread_local itemset<typename Trait::pattern_type> rows;

//     float_type acc = 0;
//     for (size_t i = 0; i < c.data.num_components(); ++i)
//     {
//         rows.clear();
//         intersection(x.row_ids, c.masks[i], rows);
//         auto       support = count(rows);
//         auto       fr      = static_cast<float_type>(support) / c.data.subset(i).size();
//         float_type alpha(0.05);

//         if constexpr (!is_dynamic_factor_model<typename Trait::distribution_type>())
//         {
//             acc += z_test_bin_score(c.models[i], x, fr, support, c.data.size(), alpha);
//         }
//         else
//         {
//             auto u = usage_of_new_factor(x.pattern, c.data, c.models[i], x.row_ids);
//             acc += z_test_bin_score(c.models[i], x, fr, u, c.data.size(), alpha);
//         }
//     }

//     return acc;
// }

} // namespace sd::disc