
#pragma once

// #include <boost/math/distributions.hpp> // chi-square
#include <cmath>
#include <algorithm>

namespace sd
{

// template <typename T>
// T chi_squared_pvalue_boost(const T& p, size_t dof)
// {
//     namespace bm = boost::math;
//     bm::chi_squared chi(dof);
//     return T(1) - bm::cdf(chi, p);
//     // return chi2_cdf<T>(p, dof);
// }

template <typename T, size_t max_iteration = 100>
auto incomplete_gamma(T s, const T& z)
{
    // static_assert(std::is_floating_point_v<T>, "");

    using std::pow, std::exp;

    if (z < 0.0)
    {
        return T(0);
    }

    T sc = 1.0 / s;
    sc *= pow(z, s);
    sc *= exp(-z);

    T sum   = 1.0;
    T nom   = 1.0;
    T denom = 1.0;

    for (size_t i = 0; i < max_iteration; i++)
    {
        nom *= z;
        s++;
        denom *= s;
        sum += nom / denom;
    }

    return sum * sc;
}

template <typename T>
T chi_squared_pvalue(T cv, size_t dof)
{
    // static_assert(std::is_floating_point_v<T>, "");

    using std::pow, std::exp, std::isnan, std::isinf, std::tgamma;


    if (cv < 0 || dof < 1)
    {
        return 0.0;
    }

    const auto x = cv * 0.5;
    const auto k = static_cast<T>(dof) * 0.5;

    if (dof == 2)
    {
        return exp(-1.0 * x);
    }

    auto pvalue = incomplete_gamma<T>(k, x);
    if (isnan(pvalue) || isinf(pvalue) || pvalue <= 1e-8)
    {
        return 1e-14;
    }

    // pvalue /= gamma(k);
    pvalue /= tgamma(k);

    return 1.0 - pvalue;
}


// Use Horner's method to evaluate a polynomial.
template <typename T, size_t N>
T horners_method(T x, const T (&poly)[N])
{
    using std::fma;
    T p = poly[N - 1];
    for (size_t i = 2; i <= N; i++)
    {
        p = fma(p, x, poly[N - i]);
    }
    return p;
}

static constexpr size_t kLargeDOF = 150;

// Returns the probability of a normal z-value.
//
// Adapted from the POZ function in:
//     Ibbetson D, Algorithm 209
//     Collected Algorithms of the CACM 1963 p. 616
//
template<class real_t>
real_t POZ(real_t z)
{
    static constexpr real_t kP1[] = {
        0.797884560593,
        -0.531923007300,
        0.319152932694,
        -0.151968751364,
        0.059054035642,
        -0.019198292004,
        0.005198775019,
        -0.001075204047,
        0.000124818987,
    };
    static constexpr real_t kP2[] = {
        0.999936657524,
        0.000535310849,
        -0.002141268741,
        0.005353579108,
        -0.009279453341,
        0.011630447319,
        -0.010557625006,
        0.006549791214,
        -0.002034254874,
        -0.000794620820,
        0.001390604284,
        -0.000676904986,
        -0.000019538132,
        0.000152529290,
        -0.000045255659,
    };

    const real_t kZMax = 6.0; // Maximum meaningful z-value.
    if (z == 0.0)
    {
        return 0.5;
    }
    real_t x;
    real_t y = 0.5 * std::fabs(z);
    if (y >= (kZMax * 0.5))
    {
        x = 1.0;
    }
    else if (y < 1.0)
    {
        real_t w = y * y;
        x        = horners_method(w, kP1) * y * 2.0;
    }
    else
    {
        y -= 2.0;
        x = horners_method(y, kP2);
    }
    return z > 0.0 ? ((x + 1.0) * 0.5) : ((1.0 - x) * 0.5);
}

// Approximates the survival function of the normal distribution.
//
// Algorithm 26.2.18, from:
// [Abramowitz and Stegun, Handbook of Mathematical Functions,p.932]
// http://people.math.sfu.ca/~cbm/aands/abramowitz_and_stegun.pdf
//
template<class real_t>
real_t normal_survival(real_t z)
{
    // Maybe replace with the alternate formulation.
    // 0.5 * erfc((x - mean)/(sqrt(2) * sigma))
    static constexpr real_t kR[] = {
        1.0,
        0.196854,
        0.115194,
        0.000344,
        0.019527,
    };
    real_t r = horners_method(z, kR);
    r *= r;
    return real_t(0.5) / (r * r);
}

// Calculates the p-value (probability) of a given chi-square value
// and degrees of freedom.
//
// Adapted from the POCHISQ function from:
//     Hill, I. D. and Pike, M. C.  Algorithm 299
//     Collected Algorithms of the CACM 1963 p. 243
//
template<typename real_t>
real_t pochisq(real_t chi_square, size_t dof)
{
    static constexpr real_t kLogSqrtPi     = 0.5723649429247000870717135; // Log[Sqrt[Pi]]
    static constexpr real_t kInverseSqrtPi = 0.5641895835477562869480795; // 1/(Sqrt[Pi])

    // For large degrees of freedom, use the normal approximation by
    //     Wilson, E. B. and Hilferty, M. M. (1931)
    // Via Wikipedia:
    //   By the Central Limit Theorem, because the chi-square distribution is the
    //   sum of k independent random variables with finite mean and variance, it
    //   converges to a normal distribution for large k.
    if (dof > kLargeDOF)
    {
        // Re-scale everything.
        const real_t chi_square_scaled = std::pow(chi_square / dof, 1.0 / 3);
        const real_t mean              = 1 - 2.0 / (9 * dof);
        const real_t variance          = 2.0 / (9 * dof);
        // If variance is 0, this method cannot be used.
        if (variance != 0)
        {
            const real_t z = (chi_square_scaled - mean) / std::sqrt(variance);
            if (z > 0)
            {
                return normal_survival(z);
            }
            else if (z < 0)
            {
                return 1.0 - normal_survival(-z);
            }
            else
            {
                return 0.5;
            }
        }
    }

    // The chi square function is >= 0 for any degrees of freedom.
    // In other words, probability that the chi square function >= 0 is 1.
    if (chi_square <= 0.0)
        return 1.0;

    // If the degrees of freedom is zero, the chi square function is always 0 by
    // definition. In other words, the probability that the chi square function
    // is > 0 is zero (chi square values <= 0 have been filtered above).
    if (dof < 1)
        return 0;

    auto                    capped_exp = [](real_t x) { return x < -20 ? real_t(0.0) : std::exp(x); };
    static constexpr real_t kBigX      = 20;

    real_t       a    = 0.5 * chi_square;
    const bool   even = !(dof & 1); // True if dof is an even number.
    const real_t y    = capped_exp(-a);
    real_t       s    = even ? y : (real_t(2.0) * POZ(-std::sqrt(chi_square)));

    if (dof <= 2)
    {
        return s;
    }

    chi_square = 0.5 * (real_t(dof) - 1);
    real_t z   = (even ? real_t(1.0) : real_t(0.5));
    if (a > kBigX)
    {
        real_t e = (even ? real_t(0.0) : kLogSqrtPi);
        real_t c = std::log(a);
        while (z <= chi_square)
        {
            e = std::log(z) + e;
            s += capped_exp(c * z - a - e);
            z += 1;
        }
        return s;
    }

    real_t e = (even ? real_t(1.0) : (kInverseSqrtPi / std::sqrt(a)));
    real_t c = 0;
    while (z <= chi_square)
    {
        e = e * (a / z);
        c = c + e;
        z += 1;
    }
    return c * y + s;
}

template<class real_t>
real_t critchi(real_t p, size_t df)
{
    using std::sqrt;
    real_t minchisq = 0.0;
    real_t maxchisq = 99999.0;

    if (p <= 0.0)
        return maxchisq;
    else if (p >= 1.0)
        return (0.0);

    real_t chisqval = df / sqrt(p); /* fair first value */
    while (maxchisq - minchisq > real_t(0.000001))
    {
        if (pochisq(chisqval, df) < p)
            maxchisq = chisqval;
        else
            minchisq = chisqval;
        chisqval = (maxchisq + minchisq) * 0.5;
    }
    return chisqval;
}

template <class T>
T regularized_lower_incomplete_gamma(T x, T alpha)
{
    using std::exp, std::log, std::lgamma, std::abs;

    const T eps = 1e-10;

    if (x <= 0.0 || alpha <= 0.0) return 0.0;

    T gamma_f = exp(alpha * log(x) - lgamma(alpha) - x);
    if (x < alpha + 1) {
        T i       = alpha;
        T tmp_sum = 1.0;
        T sum     = tmp_sum;
        while (tmp_sum / sum > eps) {
            i += 1;
            tmp_sum *= x / i;
            sum += tmp_sum;
        }
        return gamma_f * sum / alpha;
    }
    else
    {
        // Solve via evaluting continued fractions
        // Evaluation of continued fraction
        T a    = 1.0 - alpha;
        T b    = 1 + x + a;
        T pa1  = 1.0;
        T pb1  = x;
        T pa2  = x + 1.0;
        T pb2  = b * x;
        T func = pa2 / pb2;
        T pa, pb, ratio, tmp;

        for (size_t i = 0; i < 100; ++i) {
            a += 1;
            b += 2;
            pa = b * pa2 - a * i * pa1;
            pb = b * pb2 - a * i * pb1;
            if (pb) {
                ratio = pa / pb;
                tmp   = abs(func - ratio);
                if (tmp <= eps * ratio) break;
                func = ratio;
            }
            else
            {
                tmp = 1.0;
            }

            pa1 = pa2;
            pb1 = pb2;
            pa2 = pa;
            pb2 = pb;
        }
        return T(1.0) - func * gamma_f;
    }
}

template <class T>
T complemented_incomplete_gamma(T x, T alpha)
{
    using std::exp, std::log, std::lgamma, std::abs;

    const T eps = 1e-10;

    if ((x <= 0) || (alpha <= 0)) return 1.0;

    if ((x < 1.0) || (x < alpha)) return 1.0 - regularized_lower_incomplete_gamma<T>(x, alpha);

    T gamma_f = exp(alpha * log(x) - lgamma(alpha) - x);

    T y    = 1.0 - alpha;
    T z    = 1.0 + x + y;
    T c    = 0.0;
    T pkm2 = 1.0;
    T qkm2 = x;
    T pkm1 = x + 1.0;
    T qkm1 = z * x;
    T func = pkm1 / qkm1;
    T ratio, tmp, pk, qk;

    for (size_t i = 0; i < 100; ++i) {
        c += 1.0;
        y += 1.0;
        z += 2.0;
        pk = pkm1 * z - pkm2 * y * c;
        qk = qkm1 * z - qkm2 * y * c;
        if (qk != 0) {
            ratio = pk / qk;
            tmp   = abs((func - ratio) / ratio);
            if (tmp <= eps * ratio) break;
            func = ratio;
        }
        else
        {
            tmp = 1.0;
        }
        pkm2 = pkm1;
        pkm1 = pk;
        qkm2 = qkm1;
        qkm1 = qk;
        if (abs(pk) > 1e32) {
            pkm2 *= 1e-32;
            pkm1 *= 1e-32;
            qkm2 *= 1e-32;
            qkm1 *= 1e-32;
        }
    }
    return func * gamma_f;
}

template <class T>
T chi2_sf(T x, T k)
{
    return complemented_incomplete_gamma<T>(x / 2, k / 2);
}

template <class T>
T chi2_cdf(T x, T k)
{
    using std::exp;
    if (k == T(2)) {
        return T(1) - exp(T(-0.5) * x);
    }
    else
    {
        return regularized_lower_incomplete_gamma<T>(x / 2, k / 2);
    }
}

template <class T>
T chi2_pvalue(T x, size_t k)
{
    return T(1) - chi2_cdf(x, T(k));
}

template <class T>
T chi2_pdf(T x, T k)
{
    using std::exp, std::pow, std::tgamma;
    if (x < T(0)) return T(0);
    return pow(x, T(0.5) * k - T(1.0)) * exp(-T(0.5) * x) /
           (pow(T(2.0), T(0.5) * k) * tgamma(T(0.5) * k));
}

template <class T>
struct type_identity
{
    using type = T;
};
template <class T>
using type_identity_t = typename type_identity<T>::type;

template <typename T>
T ibeta(T x, type_identity_t<T> a, type_identity_t<T> b, const size_t max_iter = 2000)
{
    using std::lgamma, std::log, std::exp, std::fabs, std::max;

    constexpr static T tiny = 1.0e-30;
    constexpr static T stop = 1.0e-8;

    if (x < 0.0 || x > 1.0) return std::numeric_limits<T>::infinity();

    /*The continued fraction converges nicely for x < (a+1)/(a+b+2)*/
    if (x > (a + 1.0) / (a + b + 2.0)) {
        return (1.0 - ibeta(b, a, 1.0 - x)); /*Use the fact that beta is symmetrical.*/
    }

    /*Find the first part before the continued fraction.*/
    const T lbeta_ab = lgamma(a) + lgamma(b) - lgamma(a + b);
    const T front    = exp(log(x) * a + log(1.0 - x) * b - lbeta_ab) / a;

    /*Use Lentz's algorithm to evaluate the continued fraction.*/
    T f = 1.0, c = 1.0, d = 0.0;
    for (size_t i = 0; i <= max_iter; ++i) {
        size_t m = i / 2;
        // size_t m = static_cast<size_t>(T(i) / 2);

        double numerator;
        if (i == 0) {
            numerator = 1.0; /*First numerator is 1.0.*/
        }
        else if (i % 2 == 0)
        {
            numerator =
                (m * (b - m) * x) / ((a + 2.0 * m - 1.0) * (a + 2.0 * m)); /*Even term.*/
        }
        else
        {
            numerator = -((a + m) * (a + b + m) * x) /
                        ((a + 2.0 * m) * (a + 2.0 * m + 1.0)); /*Odd term.*/
        }

        /*Do an iteration of Lentz's algorithm.*/
        // d = 1.0 + numerator * d;
        d = T(1) / max<T>(fabs(1.0 + numerator * d), tiny);
        // if (fabs(d) < tiny) d = tiny;
        // d = 1.0 / d;

        // c = 1.0 + numerator / c;
        // if (fabs(c) < tiny) c = tiny;
        c = max<T>(fabs(1.0 + numerator / c), tiny);

        const T cd = c * d;
        f *= cd;

        if (fabs(1.0 - cd) < stop) {
            return front * (f - 1.0);
        }
    }

    return std::numeric_limits<T>::infinity();
}

template <typename T>
T reg_ibeta(T x, size_t a, size_t b)
{
    using std::beta;
    return ibeta(T(x), T(a), T(b)) / beta(T(a), T(b));
}

template <typename T>
T f_dist_cdf(T x, T a, T b)
{
    return reg_ibeta(a * x / (a * x + b), T(a) / 2, T(b) / 2);
}

template <typename T>
T binomial_cdf_beta(size_t s, T p, size_t n)
{
    if (s < n * p) {
        s = n - s;
        p = 1. - p;
    }
    //  else
    //  if (s > n * p)

    //     return f_dist_cdf<T>(p, s, n);
    // else
    // return 0;

    return f_dist_cdf<T>(T(T(1) - p) / p * T(s + 1) / (n - s), 2 * (n - s), 2 * (s + 1));
}

// template <typename T>
// T binomial_cdf_beta(size_t s, T p, size_t n)
// {
//     return reg_ibeta(1 - p, n - s, s + 1);
// }

}