// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

// zeta.hpp
// Template wrappers for the zeta function family
//
// Provided functions:
//   riemannZeta(s)      — Riemann zeta function ζ(s)
//   hurwitzZeta(s, a)   — Hurwitz zeta function ζ(s,a)
//   dirichletEta(s)     — Dirichlet η(s) = (1-2^{1-s})·ζ(s)
//
// Supported types:
//   float, double, long double — C++17 std:: delegation / direct Euler-Maclaurin implementation
//   Float                      — delegates to the sangi:: implementation

#ifndef SANGI_SPECIAL_ZETA_HPP
#define SANGI_SPECIAL_ZETA_HPP

#include <math/special/special_common.hpp>
#include <math/special/gamma.hpp>
#include <math/core/mp/Float.hpp>
#include <cmath>
#include <limits>
#include <vector>
#include <algorithm>

namespace sangi {
namespace special {

// ================================================================
// Riemann zeta function ζ(s)
// ================================================================

/// Native floating-point types: delegate to std::riemann_zeta
template<IsNativeFloat T>
[[nodiscard]] T riemannZeta(T s) {
    if (std::isnan(s)) return std::numeric_limits<T>::quiet_NaN();
    return std::riemann_zeta(s);
}

// Float type: use sangi::zeta(s, precision) directly.

// ================================================================
// Hurwitz zeta function ζ(s, a) = Σ_{k=0}^{∞} 1/(a+k)^s
// ================================================================
//
// Euler-Maclaurin formula:
// ζ(s,a) ≈ Σ_{k=0}^{N-1} (a+k)^{-s} + (a+N)^{1-s}/(s-1) + (a+N)^{-s}/2
//         + Σ_{j=1}^{M} B_{2j}/(2j)! · s(s+1)···(s+2j-2) · (a+N)^{-(s+2j-1)}

template<IsNativeFloat T>
[[nodiscard]] T hurwitzZeta(T s, T a) {
    if (std::isnan(s) || std::isnan(a)) return std::numeric_limits<T>::quiet_NaN();
    if (a <= T(0)) return std::numeric_limits<T>::quiet_NaN();

    // s = 1 is a pole
    if (s == T(1)) return std::numeric_limits<T>::quiet_NaN();

    // a = 1 reduces to the Riemann zeta
    if (a == T(1)) return riemannZeta(s);

    // N: number of terms in the Euler-Maclaurin direct sum (larger -> higher precision, more work)
    //    15 is an empirically optimal value that secures ~12 digits of precision in double
    const int N = 15;
    // M: number of Bernoulli correction terms (uses B_2 ~ B_{2M})
    //    M=8 is near the limit of double precision; M>8 degrades precision due to divergence of the Bernoulli numbers
    const int M = 8;

    // B_{2j} for j = 1..8
    static constexpr double B2j[] = {
         1.0/6.0,            // B_2
        -1.0/30.0,           // B_4
         1.0/42.0,           // B_6
        -1.0/30.0,           // B_8
         5.0/66.0,           // B_10
        -691.0/2730.0,       // B_12
         7.0/6.0,            // B_14
        -3617.0/510.0,       // B_16
    };

    T sum = T(0);
    for (int k = 0; k < N; k++) {
        sum += std::pow(a + T(k), -s);
    }

    T aN = a + T(N);
    T aN_inv = T(1) / aN;

    // Integral term: (a+N)^{1-s} / (s-1)
    sum += std::pow(aN, T(1) - s) / (s - T(1));

    // Midpoint correction: (1/2) · (a+N)^{-s}
    sum += T(0.5) * std::pow(aN, -s);

    // Bernoulli correction terms
    T rising = s;  // rising factorial s(s+1)···(s+2j-2); for j=1 just s
    T aN_pow = std::pow(aN, -s - T(1));  // (a+N)^{-(s+1)}

    for (int j = 0; j < M; j++) {
        // (2j+2)! = (2(j+1))!
        T fact = T(1);
        for (int i = 1; i <= 2*(j+1); i++) fact *= T(i);

        sum += static_cast<T>(B2j[j]) / fact * rising * aN_pow;

        // Update rising factorial: append (s+2j-1)(s+2j) → next j has 2(j+1)-1=2j+1 factors
        rising *= (s + T(2*j+1)) * (s + T(2*(j+1)));
        // (a+N)^{-(s+2(j+1)+1)} = (a+N)^{-(s+2j+1)} · (a+N)^{-2}
        aN_pow *= aN_inv * aN_inv;
    }

    return sum;
}

// Float type: use sangi::hurwitzZeta(s, a, precision) directly.

// ----------------------------------------------------------------
// Complex Hurwitz ζ(s, a) = Σ_{k=0}^∞ (a+k)^{-s}  (Euler-Maclaurin, complex s, a)
// ----------------------------------------------------------------
// Valid as an analytic continuation for all s ≠ 1. The argument shift N makes Re(a+N) sufficiently large.
// The main use is the general Jonquière inversion of polylog for |z|>1 (Re(a) ∈ (0,1] is the primary target).
template<IsNativeFloat R>
[[nodiscard]] Complex<R> hurwitzZeta(Complex<R> s, Complex<R> a) {
    using C = Complex<R>;
    if (s.re == R(1) && s.im == R(0))
        return C(std::numeric_limits<R>::quiet_NaN());

    // B_2, B_4, ..., B_16
    static constexpr R B2j[] = {
         R(1)/R(6), -R(1)/R(30), R(1)/R(42), -R(1)/R(30),
         R(5)/R(66), -R(691)/R(2730), R(7)/R(6), -R(3617)/R(510),
    };
    constexpr int M = 8;
    constexpr int N = 20;

    C sum(R(0));
    for (int k = 0; k < N; k++)
        sum = sum + exp(-s * log(a + C(R(k))));   // (a+k)^{-s}

    C aN = a + C(R(N));
    C ln_aN = log(aN);
    sum = sum + exp((C(R(1)) - s) * ln_aN) / (s - C(R(1)));  // (a+N)^{1-s}/(s-1)
    sum = sum + C(R(0.5)) * exp(-s * ln_aN);                 // (1/2)(a+N)^{-s}

    // Σ B_{2j}/(2j)! · (s)_{2j-1} · (a+N)^{-(s+2j-1)}
    C rising = s;                          // (s)_1 = s
    C aN_inv = C(R(1)) / aN;
    C aN_pow = exp(-s * ln_aN) * aN_inv;   // (a+N)^{-(s+1)}
    for (int j = 0; j < M; j++) {
        R fact = R(1);
        for (int i = 1; i <= 2*(j+1); i++) fact *= R(i);
        sum = sum + C(B2j[j] / fact) * rising * aN_pow;
        rising = rising * (s + C(R(2*j+1))) * (s + C(R(2*(j+1))));
        aN_pow = aN_pow * aN_inv * aN_inv;
    }
    return sum;
}

// ================================================================
// Dirichlet η function: η(s) = (1 - 2^{1-s}) · ζ(s)
// ================================================================
//
// Alternating zeta function: η(s) = Σ_{k=1}^{∞} (-1)^{k-1} / k^s
// η(1) = ln(2), η(0) = 1/2

template<IsNativeFloat T>
[[nodiscard]] T dirichletEta(T s) {
    if (std::isnan(s)) return std::numeric_limits<T>::quiet_NaN();

    // η(1) = ln(2) (avoids the pole of ζ(1))
    if (s == T(1)) return std::log(T(2));

    T factor = T(1) - std::pow(T(2), T(1) - s);
    return factor * riemannZeta(s);
}

// Float type: use sangi::dirichletEta(s, precision) directly.

// ================================================================
// Complex<R> overloads
// ================================================================

// ----------------------------------------------------------------
// η(s) — Hasse series (Euler acceleration)
// ----------------------------------------------------------------
// η(s) = Σ_{n=0}^{N} (1/2^{n+1}) Σ_{k=0}^{n} (-1)^k C(n,k) (k+1)^{-s}
// O(N) exp + O(N²) multiply-adds. N=50 → 2^{-50} ≈ 10^{-15} precision.

template<IsNativeFloat R>
[[nodiscard]] Complex<R> dirichletEta(Complex<R> s) {
    using C = Complex<R>;
    constexpr int N = 50;

    // Precompute (k+1)^{-s}
    C ks[N + 1];
    for (int k = 0; k <= N; k++) {
        ks[k] = exp(-s * C(std::log(R(k + 1))));
    }

    C eta(R(0));
    R half_pow = R(0.5);  // (1/2)^{n+1}

    for (int n = 0; n <= N; n++) {
        // Σ_{k=0}^{n} (-1)^k C(n,k) (k+1)^{-s}
        C inner(R(0));
        R binom = R(1);   // C(n, 0) = 1
        R sign = R(1);

        for (int k = 0; k <= n; k++) {
            inner = inner + C(sign * binom) * ks[k];
            sign = -sign;
            if (k < n) binom *= R(n - k) / R(k + 1);
        }

        eta = eta + C(half_pow) * inner;
        half_pow *= R(0.5);
    }

    return eta;
}

// ----------------------------------------------------------------
// ζ(s) = η(s) / (1 - 2^{1-s})
// ----------------------------------------------------------------

template<IsNativeFloat R>
[[nodiscard]] Complex<R> riemannZeta(Complex<R> s) {
    using C = Complex<R>;

    // s = 1 is a pole
    if (s.re == R(1) && s.im == R(0))
        return C(std::numeric_limits<R>::quiet_NaN());

    C eta = dirichletEta(s);
    C factor = C(R(1)) - exp((C(R(1)) - s) * C(std::log(R(2))));
    return eta / factor;
}

// ================================================================
// Complex<Float> dedicated overloads (Euler acceleration + reflection formula)
// ================================================================
// The Hasse series (N=50, ~15 digits) is kept as IsNativeFloat-only.
// Complex<Float> gets arbitrary precision via Euler transform + reflection formula.

// ----------------------------------------------------------------
// η(s) — Euler acceleration (Complex<Float>)
// ----------------------------------------------------------------
// Euler transform of η(s) = Σ (-1)^{k-1}/k^s
// E_n = Σ_{k=0}^{n} C(n,k)/2^n · S_{k+1}

[[nodiscard]] inline Complex<Float> dirichletEta(Complex<Float> s, int precision) {
    using C = Complex<Float>;
    int wp = precision + 20;
    // PrecisionGuard removed: setResultPrecision(wp) on sw propagates req=wp.

    // η(1) = ln(2)
    if (s.im.isZero() && s.re == Float(1))
        return C(log(Float(2), precision));

    // Term count: the Euler transform error is ~2^{-n} (1 bit/term). BUGFIX (2026-05-30):
    // n = wp·1.05 misused decimal digits as the bit-term count and plateaued at ~0.3·P digits
    // → wp_bits·1.05 (AUDIT_FLOAT_UNIT_MIXING pattern C).
    int n = static_cast<int>(Float::precisionToBits(wp) * 1.05) + 5;

    C sw = s;
    sw.re.setResultPrecision(wp);  // pad input to working precision (also sets eff)
    sw.im.setResultPrecision(wp);

    // Euler transform
    C partial_sum;
    C weighted_sum;

    // w_0 = 2^{-n}
    Float w = pow(Float(2), Float(-n), wp);

    for (int k = 0; k <= n; k++) {
        // (k+1)^{-s} = exp(-s · ln(k+1))
        Float ln_kp1 = log(Float(k + 1), wp);
        C term_k = exp(-sw * C(ln_kp1));

        // Update partial sum: S_{k+1} = S_k + (-1)^k / (k+1)^s
        if (k % 2 == 0)
            partial_sum = partial_sum + term_k;
        else
            partial_sum = partial_sum - term_k;

        // Weighted sum
        weighted_sum = weighted_sum + C(w) * partial_sum;

        // Update weight: w_{k+1} = w_k · (n-k)/(k+1)
        if (k < n) {
            w = w * Float(n - k) / Float(k + 1);
            w.setPrecision(wp);
        }
    }

    weighted_sum.re.setPrecision(precision);
    weighted_sum.im.setPrecision(precision);
    return weighted_sum;
}

[[nodiscard]] inline Complex<Float> dirichletEta(Complex<Float> s) {
    return dirichletEta(std::move(s), Float::defaultPrecision());
}

// ----------------------------------------------------------------
// ζ(s) = η(s) / (1 - 2^{1-s}) (Complex<Float>)
// ----------------------------------------------------------------
// Re(s) < 0: reflection formula ζ(s) = 2^s·π^{s-1}·sin(πs/2)·Γ(1-s)·ζ(1-s)

[[nodiscard]] inline Complex<Float> riemannZeta(Complex<Float> s, int precision) {
    using C = Complex<Float>;
    int wp = precision + 20;
    // PrecisionGuard removed: the internal dirichletEta / lnGamma calls produce
    // req=wp on their results via the precision argument, and Complex arithmetic propagates wp through mergeRequested.

    // Re(s) < 0: reflection formula
    if (s.re.isNegative()) {
        C one_minus_s = C(Float(1)) - s;

        // Trivial zeros: s a negative even integer
        if (s.im.isZero() && s.re.isInteger()) {
            int si = static_cast<int>(s.re.toDouble());
            if (si % 2 == 0 && si < 0)
                return C();
        }

        C z_1ms = riemannZeta(one_minus_s, wp);
        C g_1ms = gamma(one_minus_s, wp);

        Float pi_val = Float::pi(wp);
        C ln2(log(Float(2), wp));
        C lnpi(log(pi_val, wp));

        C two_pow_s = exp(s * ln2);
        C pi_pow_sm1 = exp((s - C(Float(1))) * lnpi);
        C sin_term = sin(C(ldexp(pi_val, -1)) * s);

        C result = two_pow_s * pi_pow_sm1 * sin_term * g_1ms * z_1ms;
        result.re.setPrecision(precision);
        result.im.setPrecision(precision);
        return result;
    }

    // Re(s) >= 0, s ≠ 1: via η
    C eta = dirichletEta(s, wp);
    C ln2(log(Float(2), wp));
    C factor = C(Float(1)) - exp((C(Float(1)) - s) * ln2);
    C result = eta / factor;

    result.re.setPrecision(precision);
    result.im.setPrecision(precision);
    return result;
}

[[nodiscard]] inline Complex<Float> riemannZeta(Complex<Float> s) {
    return riemannZeta(std::move(s), Float::defaultPrecision());
}

// ----------------------------------------------------------------
// Complex Hurwitz ζ(s, a) (Complex<Float>, arbitrary-precision Euler-Maclaurin)
// ----------------------------------------------------------------
// Same formula as the native version. N ≈ wp/3 shift + M ≈ wp/5 Bernoulli correction (cut off by convergence test).
// Reuses the Bernoulli numbers B_{2k} from detail::computeBernoulliFloat (gamma.hpp).
[[nodiscard]] inline Complex<Float> hurwitzZeta(Complex<Float> s, Complex<Float> a,
                                                int precision) {
    using C = Complex<Float>;
    if (s.im.isZero() && s.re == Float(1))
        return C(Float::nan());

    int wp = precision + 40;
    C sw = s; sw.re.setResultPrecision(wp); sw.im.setResultPrecision(wp);
    C aw = a; aw.re.setResultPrecision(wp); aw.im.setResultPrecision(wp);

    // BUGFIX (2026-05-30): fixed the same unit-mixing as in the native version. a+N ≳ 0.366·wp is
    // needed for the minimum achievable error 10^{-wp} → wp·0.4.
    int target_aN = static_cast<int>(wp * 0.4) + 5;
    int N = std::max(1, target_aN - static_cast<int>(aw.re.toDouble()));

    C sum(Float::zero(wp));
    for (int k = 0; k < N; k++) {
        C term = exp(-sw * log(aw + C(Float(k))));   // (a+k)^{-s}
        sum = sum + term;
        sum.re.truncateToApprox(wp); sum.im.truncateToApprox(wp);
    }

    C aN = aw + C(Float(N));
    C ln_aN = log(aN);
    const C one_c(Float::one(wp));
    sum = sum + exp((one_c - sw) * ln_aN) / (sw - one_c);   // (a+N)^{1-s}/(s-1)
    sum = sum + C(Float::one(wp) / Float(2)) * exp(-sw * ln_aN);  // (1/2)(a+N)^{-s}
    sum.re.truncateToApprox(wp); sum.im.truncateToApprox(wp);

    // BUGFIX (2026-05-30): the old wp/5+5 truncated far short of the optimal cutoff M*≈π·(a+N)
    // and plateaued at ~0.3·P digits → π·target_aN (cut off early by relative-eps test).
    constexpr double PI_VAL = 3.141592653589793238462643383279502884;
    int M = static_cast<int>(PI_VAL * static_cast<double>(target_aN)) + 10;
    int bp = wp + 50 + 3 * M;
    std::vector<Float> Bev = detail::computeBernoulliFloat(M, bp);  // Bev[k] = B_{2k}

    C rising = sw;                          // (s)_1
    C aN_inv = one_c / aN;
    C aN_pow = exp(-sw * ln_aN) * aN_inv;   // (a+N)^{-(s+1)}
    for (int j = 0; j < M; j++) {
        Float fact = Float::one(wp);
        for (int i = 1; i <= 2*(j+1); i++) fact = fact * Float(i);
        fact.truncateToApprox(wp);
        Float coeff = Bev[j + 1] / fact;    // B_{2(j+1)} / (2(j+1))!
        coeff.truncateToApprox(wp);
        C correction = C(coeff) * rising * aN_pow;
        sum = sum + correction;
        sum.re.truncateToApprox(wp); sum.im.truncateToApprox(wp);

        if (j >= 3) {
            Float ac = abs(correction);
            if (!ac.isZero() && ac < abs(sum) * Float::epsilon(wp)) break;
        }
        rising = rising * (sw + C(Float(2*j+1))) * (sw + C(Float(2*(j+1))));
        rising.re.truncateToApprox(wp); rising.im.truncateToApprox(wp);
        aN_pow = aN_pow * aN_inv * aN_inv;
        aN_pow.re.truncateToApprox(wp); aN_pow.im.truncateToApprox(wp);
    }

    sum.re.setPrecision(precision); sum.im.setPrecision(precision);
    return sum;
}

[[nodiscard]] inline Complex<Float> hurwitzZeta(Complex<Float> s, Complex<Float> a) {
    return hurwitzZeta(std::move(s), std::move(a), Float::defaultPrecision());
}

} // namespace special
} // namespace sangi

#endif // SANGI_SPECIAL_ZETA_HPP