#pragma once

#include <array>
#include <cmath>
#include <numeric>
#include <vector>

namespace sd::disc::fdr
{

template <class T>
T geometric_number(size_t k)
{
    return T(1) / pow(2, k);
}
template <class T>
T pi_sequence_number(size_t k)
{
    //     return T(1) / (k*k) ;
    return T(1) / T(k * k) * T(6) / T(3.1415);
}
template <class T>
T ln2_sequence_number(size_t k)
{
    return (T(1) / (T(2 * k + 1) * T(2 * k + 1))) / T(0.6931);
}
template <class T>
T lord_sequence_number(size_t k)
{
    using std::log, std::exp;
    return T(0.0722) * log(std::max<T>(k, 2)) / (k * exp(sqrt(std::max<T>(k, 1))));
}

template <class T>
T discount_gamma(size_t k)
{
    //                     return acc += alpha_0 * pi_sequence_number<T>(k - i + 1);
    //                     return acc += alpha_0 * geometric_number<T>(k - i + 1) / 2;
    //                     return acc += gamma_sequence<T>(k - i + 1);
    return lord_sequence_number<T>(k);
}

template <class T>
T discount_xi(size_t k, T alpha, T b_0)
{
    return (T(0.139307) * alpha / b_0) / T(k * log(std::max<size_t>(k, 2)));
}

// template <class T>
// T xi_faculty(size_t k, T alpha, T b_0)
// {
//     constexpr fak = std::array<size_t, 20>{{1,
//                                             1,
//                                             2,
//                                             6,
//                                             24,
//                                             120,
//                                             720,
//                                             5'040,
//                                             40'320,
//                                             362'880,
//                                             3'628'800,
//                                             39'916'800,
//                                             479'001'600,
//                                             6'227'020'800,
//                                             87'178'291'200,
//                                             1'307'674'368'000,
//                                             20'922'789'888'000,
//                                             355'687'428'096'000,
//                                             6'402'373'705'728'000,
//                                             121'645'100'408'832'000,
//                                             2'432'902'008'176'640'000}};

//     if (k >= fak.size()) { return 0; }
//     return ((alpha/b_0) / T(2.32206)) / fak[k];
// }

// template <class T>
// T xi_sq(size_t k, T alpha, T b_0)
// {
//     return (T(1) / T(2.32206)) * (b_0 / alpha) / (k * k);
// }

template <class T>
T xi_sq(size_t k, T alpha, T b_0)
{
    k = std::max<size_t>(k, 1);
    constexpr T pipi6 = 3.1415 * 3.1415 / 6;
    return (alpha / b_0) / (pipi6 * (k * k * (1 + log(k))));
}

template <class T>
struct Fixed
{
    T    alpha;
    bool test_next(const T& pvalue) const { return pvalue < alpha; }
};
template <class T>
struct AlphaSpending
{
    T alpha_i;

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

    bool test_next(const T& pvalue)
    {
        bool reject = pvalue < alpha_i;
        alpha_i     = std::max<T>(0, alpha_i - pvalue);
        return reject;
    }
};
template <class T>
struct OnlineBonferroni
{
    T      alpha, alpha_i;
    size_t k = 0;

    explicit OnlineBonferroni(T alpha_0 = 0.05) : alpha(alpha_0) {}

    bool test_next(const T& pvalue)
    {
        k       = k + 1;
        alpha_i = alpha * discount_gamma<T>(k);
        return pvalue < alpha_i;
    }
};
template <class T>
struct OnlineSidak
{
    T      alpha, alpha_i;
    size_t k = 0;

    explicit OnlineSidak(T alpha_0 = 0.05) : alpha(alpha_0) {}

    bool test_next(const T& pvalue)
    {
        using std::pow;
        k       = k + 1;
        alpha_i = T(1) - pow(T(1) - alpha, discount_gamma<T>(k));
        return pvalue < alpha_i;
    }
};
template <class T>
struct ADDIS
{
    using float_type  = T;
    float_type lambda = 0.5;
    float_type alpha  = 0.05;
    float_type tau    = 0.5;
    float_type w0     = 0.00625;

    float_type alpha_i;
    size_t     s_cumsum = 0, c_cumsum = 0;

    explicit ADDIS(const float_type alpha_0 = 0.05) : alpha(alpha_0) {}

    bool test_next(const float_type pvalue)
    {
        bool s_k = pvalue <= tau;
        bool c_k = pvalue <= lambda * tau;

        c_cumsum += c_k;
        s_cumsum += s_k;

        alpha_i =
            tau * (T(1) - lambda) * (alpha - w0) * discount_gamma<T>(s_cumsum - c_cumsum + 1);

        return pvalue <= alpha_i;
    }
};

template <class T>
struct AlphaInvest
{
    T      wealth_k, alpha_i, alpha_0, wealth_0, phi_fac; //, phi_k;
    size_t k = 0, last_reject = 0;
    bool   flag = false;

    AlphaInvest(T alpha_0 = 0.05, T start_factor = 0.5, T phi_fac = 0.25)
        : alpha_0(alpha_0), wealth_0(alpha_0 * start_factor), phi_fac(phi_fac)
    {
        wealth_k = wealth_0;
        // phi_k = phi_fac * wealth_k;
        alpha_i = wealth_k / 2;
    }
    bool test_next(const T& pvalue)
    {
        k = k + 1;

        if (wealth_k <= 0) { return pvalue == 0; }

        bool reject = pvalue < alpha_i;
        // wealth_k    = wealth_k - alpha_i + reject * b0;
        if (reject)
        {
            last_reject = k;
            flag        = true;
        }

        wealth_k = wealth_k - (!reject) * alpha_i / (1 - alpha_i) +
                   reject * (alpha_0 - flag * wealth_0);

        alpha_i = std::min(wealth_k / (1 + (k + 1) - last_reject),
                           ((1 - phi_fac) * wealth_k) / ((1 - phi_fac) * wealth_k + 1));

        flag = false;
        return reject;
    }
};

// template <class T>
// struct LORD2indep
// {
//     T                   wealth_0, wealth_k, wealth_last_reject, alpha_i, alpha_0, b0;
//     size_t              k = 1, last_reject = 1;
//     std::vector<size_t> last_rejections;

//     LORD2indep(T alpha_0 = 0.05, T start_factor = 0.5)
//         : wealth_0(alpha_0 * start_factor)
//         , wealth_k(alpha_0 * start_factor)
//         , wealth_last_reject(alpha_0 * start_factor)
//         , alpha_i(alpha_0)
//         , alpha_0(alpha_0)
//         , b0(alpha_0 - alpha_0 * start_factor)
//         , k(1)
//     {
//     }

//     bool test_next(const T& pvalue)
//     {
//         k = k + 1;

//         bool reject = pvalue < alpha_i;
//         if (reject) { last_rejections.push_back(k); }

//         wealth_k = wealth_k - alpha_i + reject * b0;
//         alpha_i  = discount_gamma<T>(k) * wealth_0;

//         if (!last_rejections.empty())
//         {
//             T sum_gammas = std::accumulate(
//                 last_rejections.begin(), last_rejections.end(), T(0), [&](auto acc, auto& i)
//                 {
//                     return acc += discount_gamma<T>(k - i + 1);
//                 });
//             alpha_i += b0 * sum_gammas;
//         }
//         return reject;
//     }
// };

template <class T>
struct LOND
{
    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 LOND(float_type alpha_0 = 0.05) : alpha(alpha_0), alpha_i(alpha_0) {}

    bool test_next(const float_type pvalue)
    {
        alpha_i = 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;
    }
    auto wealth() const { return alpha_i; }
};

template <class T>
struct LORD
{
    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 LORD(float_type alpha_0 = 0.05, float_type factor = 0.5) : alpha(alpha_0)
    {
        w0      = alpha * factor;
        b0      = alpha - w0;
        alpha_i = xi_sq<T>(1, alpha, b0) * w0;
        w_tau = w = w0 - alpha_i + b0;
    }

    bool test_next(float_type const pvalue)
    {
        bool reject = pvalue < alpha_i;
        if (reject)
        {
            tau   = i;
            w_tau = w;
        }
        ++i;
        alpha_i = xi_sq<T>(i, alpha, b0) * w_tau;
        w       = w - alpha_i + reject * b0;
        return reject;
    }
    auto wealth() const { return w; }
};

template <class T>
struct LORD3
{
    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 LORD3(float_type alpha_0 = 0.05, float_type factor = 0.5) : alpha(alpha_0)
    {
        w0      = alpha * factor;
        b0      = alpha - w0;
        alpha_i = xi_sq<T>(1, alpha, b0) * w0;
        w_tau = w = w0 - alpha_i + b0;
    }

    bool test_next(const T& pvalue)
    {
        bool reject = pvalue < alpha_i;
        if (reject)
        {
            tau   = i;
            w_tau = w;
        }

        ++i;
        alpha_i = xi_sq<T>(i - tau, alpha, b0) * w_tau;
        w       = w - alpha_i + reject * b0;
        return reject;
    }
    auto wealth() const { return w; }
};


namespace mk2 {

template <class T>
struct LOND
{
    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 LOND(float_type alpha_0 = 0.05) : alpha(alpha_0), alpha_i(alpha_0) {}

    bool test_next(const float_type pvalue)
    {
        alpha_i = 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;
    }
    auto wealth() const { return alpha_i; }
};

template <class T>
struct LORD
{
    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 LORD(float_type alpha_0 = 0.05, float_type factor = 0.5) : alpha(alpha_0)
    {
        w0      = alpha * factor;
        b0      = alpha - w0;
        alpha_i = xi_sq<T>(1, alpha, b0) * w0;
        w_tau = w = w0 - alpha_i + b0;
    }

    void reject() {
        
    }

    bool test_next(float_type const pvalue)
    {
        bool reject = pvalue < alpha_i;
        if (reject)
        {
            tau   = i;
            w_tau = w;
        }
        ++i;
        alpha_i = xi_sq<T>(i, alpha, b0) * w_tau;
        w       = w - alpha_i + reject * b0;
        return reject;
    }
    auto wealth() const { return w; }
};

template <class T>
struct LORD3
{
    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 LORD3(float_type alpha_0 = 0.05, float_type factor = 0.5) : alpha(alpha_0)
    {
        w0      = alpha * factor;
        b0      = alpha - w0;
        alpha_i = xi_sq<T>(1, alpha, b0) * w0;
        w_tau = w = w0 - alpha_i + b0;
    }

    bool test_next(const T& pvalue)
    {
        bool reject = pvalue < alpha_i;
        if (reject)
        {
            tau   = i;
            w_tau = w;
        }

        ++i;
        alpha_i = xi_sq<T>(i - tau, alpha, b0) * w_tau;
        w       = w - alpha_i + reject * b0;
        return reject;
    }
    auto wealth() const { return w; }
};


}
} // namespace sd::disc::fdr