#pragma once

#include <cmath>
#include <math/nchoosek.hxx>

namespace sd::disc
{

template <typename T>
constexpr T binomial_logpmf(size_t k, T p, size_t n)
{
    return k * log(p) + T(n - k) * log(1.0 - p) + log_nchoosek<T>(n, k);
}
template <typename T>
constexpr T binomial_pmf(size_t k, T p, size_t n)
{
    return exp(binomial_logpmf(k, p, n));
}

template <typename T>
constexpr T standard_normal_cdf(T const& z)
{
    constexpr auto cdf = [mu = 0, sig = 1](auto const& x) -> T {
        constexpr T sqrt_2 = T(1.148698354997035);
        using std::erfc;
        return T(0.5) * erfc(-sqrt_2 * (x - mu) / sig);
    };
    return cdf(z);
}

// template <typename T>
// T erfinv(T x)
// {
//     T w, p;
//     w = -logf((1.0f - x) * (1.0f + x));
//     if (w < 5.000000f)
//     {
//         w = w - 2.500000f;
//         p = 2.81022636e-08f;
//         p = 3.43273939e-07f + p * w;
//         p = -3.5233877e-06f + p * w;
//         p = -4.39150654e-06f + p * w;
//         p = 0.00021858087f + p * w;
//         p = -0.00125372503f + p * w;
//         p = -0.00417768164f + p * w;
//         p = 0.246640727f + p * w;
//         p = 1.50140941f + p * w;
//     }
//     else
//     {
//         w = sqrtf(w) - 3.000000f;
//         p = -0.000200214257f;
//         p = 0.000100950558f + p * w;
//         p = 0.00134934322f + p * w;
//         p = -0.00367342844f + p * w;
//         p = 0.00573950773f + p * w;
//         p = -0.0076224613f + p * w;
//         p = 0.00943887047f + p * w;
//         p = 1.00167406f + p * w;
//         p = 2.83297682f + p * w;
//     }
//     return p * x;
// }

template <typename T>
T erfinv(const T& x)
{
    using std::log, std::abs, std::sqrt;

    constexpr T pi = T(3.14159265358979323846264338);
    constexpr T c  = T(0.15449436008930206298828125); // 0.147

    T lnx = log((1 - x) * (1 + x)); // 1 - x*x
    T tt1 = 2 / (pi * c) + (lnx / 2);
    T tt2 = 1 / c * lnx;

    return ((x < 0) ? -1 : 1) * sqrt(-tt1 + sqrt(tt1 * tt1 - tt2));
}

template <typename T>
constexpr T standard_normal_quantile(T const& p)
{
    constexpr auto cdf_inv = [mu = 0, sig = 1](auto const& x) -> T {
        constexpr T sqrt_2 = T(1.148698354997035);
        return sqrt_2 * erfinv(2 * x - 1) * sig + mu;
    };
    return cdf_inv(std::clamp<T>(p, T(1e-16), T(1)-T(1e-16)));
    // return cdf_inv(p);
}
template <typename T>
T gauss_ppf(T p, T mu = 0, T sig = 1)
{
    auto cdf_inv = [&](auto const& x) -> T {
        constexpr T sqrt_2 = T(1.148698354997035);
        return sqrt_2 * erfinv(2 * x - 1) * sig + mu;
    };
    return cdf_inv(std::clamp<T>(p, T(1e-16), T(1)-T(1e-16)));
    // return cdf_inv(p);
    // return p < T(1e-16) ? T(-5.998) : (p > (T(1) - T(1e-16)) ? T(5.59) : cdf_inv(p);
}

template <typename T>
T z_score(size_t s, T p, size_t n)
{
    T u = static_cast<T>(n) * p;
    T o = sqrt(u * (1 - p));
    T z = (T(s) - u) / o;
    return z;
}
template <typename T>
T cdf_z_lower(size_t s, T p, size_t n)
{
    return standard_normal_cdf(z_score(s, p, n));
}
template <typename T>
T cdf_z_upper(size_t s, T p, size_t n)
{
    return standard_normal_cdf(-z_score(s, p, n));
}
template <typename T>
T cdf_z_tt(size_t s, T p, size_t n)
{
    using std::abs;
    return 2 * standard_normal_cdf(-abs(z_score(s, p, n)));
}

template <typename T>
T kl_bernoulli(T const& q, T const& p)
{
    auto kl1 = [](const auto& q, const auto& p) {
        using std::log;
        return q < 1e-32 ? T(0) : T(q * log(q / p));
    };
    return kl1(q, p) + kl1(T(1) - q, T(1) - p);
}

template <typename T>
T binomial_log_cdf_chernoff_impl(size_t s, T p, size_t n)
{
    return T(n) * -kl_bernoulli(static_cast<T>(s) / static_cast<T>(n), p);
}

template <typename T>
T binomial_cdf_exact_lower(size_t s, T p, size_t n)
{
    using std::log, std::exp;
    const T lg_true  = log(p);
    const T lg_false = log(1.0 - p);
    T       b        = 0;
    T       cdf      = 0;
    for (size_t k = 0; k <= s; ++k)
    {
        if (k > 0) b += log(n - k + 1) - log(k);
        auto lg_pmf = k * lg_true + T(n - k) * lg_false;
        cdf += exp(lg_pmf + b);
    }
    return cdf;
}

template <typename T>
T binomial_cdf_exact_upper(size_t s, T p, size_t n)
{
    using std::log, std::exp;
    const T lg_true  = log(p);
    const T lg_false = log(1.0 - p);
    T       b        = 0;
    T       cdf      = 0;

    for (size_t k = 1; k < s; ++k) { b += log(n - k + 1) - log(k); }
    for (size_t k = s; k <= n; ++k)
    {
        if (k > 0) b += log(n - k + 1) - log(k);
        auto lg_pmf = k * lg_true + T(n - k) * lg_false;
        cdf += exp(lg_pmf + b);
    }
    return cdf;
}

template <typename T>
T binomial_cdf_z(size_t s, T p, size_t n)
{
    if (s < p * n)
        return cdf_z_lower(s, p, n);
    else if (s >= p * n)
        return cdf_z_upper(s, p, n);
    else
        return 1;
}
template <typename T>
T binomial_log_cdf_chernoff(size_t s, T p, size_t n)
{
    if (s < n * p)
        return binomial_log_cdf_chernoff_impl(n - s, T(1) - p, n);
    else if (s > n * p)
        return binomial_log_cdf_chernoff_impl(s, p, n);
    else
        return 0;
}
template <typename T>
T binomial_cdf_chernoff(size_t s, T p, size_t n)
{
    using std::exp;
    return exp(binomial_log_cdf_chernoff(s, p, n));
}
template <typename T>
T binomial_cdf_exact(size_t s, T p, size_t n)
{
    if (s < p * n)
        return binomial_cdf_exact_lower(s, p, n);
    else if (s > p * n)
        return binomial_cdf_exact_upper(s, p, n);
    else
        return 1;
}
template <class T>
T binomial_ci_z_gz(T p, size_t n, T z)
{
    using std::sqrt;
    return z * sqrt(p * (T(1) - p) / n);
}
template <class T>
T binomial_ci_agresti_coull_gz(size_t k, size_t n, T z)
{
    using std::sqrt;
    // auto z  = standard_normal_quantile(T(1) - alpha / 2);
    auto nt = n + z * z;
    auto pt = (k + z * z / 2) / nt;
    return z * sqrt(pt / nt * (1 - pt));
}

template <class T>
T binomial_ci_z(T p, size_t n, T alpha)
{
    return binomial_ci_z_gz(p, n, standard_normal_quantile(T(1) - alpha / 2));
}
template <class T>
T binomial_ci_agresti_coull(size_t k, size_t n, T alpha)
{
    return binomial_ci_agresti_coull_gz(k, n, standard_normal_quantile(T(1) - alpha / 2));
}
template <class T>
bool is_outside_of_ci(T q, T p, T ci)
{
    return q < (p - ci) || (p + ci) < q;
}
template <class T>
bool is_outside_of_ci(size_t k, T p, size_t n, T ci)
{
    auto q = T(k) / n;
    return is_outside_of_ci(q, p, ci);
}
template <class T>
T error_margin(T alpha, T p, T q, size_t n)
{
    return sqrt((p + q - (q - p) * (q - p)) / n) * standard_normal_quantile(alpha);
}

template <class T>
T bin_error_margin_1(T alpha, size_t k, size_t n)
{
    return sqrt(T(k) * T(n - k)) * standard_normal_quantile(1.0 - alpha / 2) / T(n * sqrt(n));
}

template <class T>
T bin_error_margin_1(T alpha, T p, size_t n)
{
    return sqrt(p * (1 - p)) * standard_normal_quantile(1.0 - alpha / 2);
}

} // namespace sd::disc
