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

// lerch.hpp
// Lerch transcendent Φ(z, s, a) = Σ_{k=0}^∞ z^k / (k + a)^s
//
// Unified representation:
//   Li_s(z)   = z · Φ(z, s, 1)
//   ζ(s, a)   = Φ(1, s, a)            (Re(s) > 1)
//   ζ(s)      = Φ(1, s, 1)            (Riemann ζ)
//   η(s)      = Φ(-1, s, 1)           (Dirichlet η)
//
// Uses: analytic number theory, unified representation of polylog/Hurwitz, generating functions of statistical distributions.
//
// Functions provided:
//   lerchPhi(z, s, a)              — native float (float / double / long double)
//   lerchPhi(z, s, a, precision)   — Float arbitrary precision
//   lerchPhi(z, s, a)              — Float (default-precision shorthand)
//
// Algorithm (stocking phase initial drop):
//   - z = 0          : a^{-s}
//   - s = 0          : 1/(1-z)
//   - z = 1          : Hurwitz ζ(s, a) (finite only for s > 1)
//   - z = -1         : Hasse-accelerated alternating sum β(s, a)
//   - |z| < 1        : direct Taylor series (no convergence acceleration)
//   - |z| ≥ 1, z∉{±1}: NaN (analytic continuation is future work)
//
// Limitations:
//   - a must be a positive real. NaN when a is a non-positive integer.
//   - Complex z and analytic continuation (|z|>1) to be extended in a future release.
//
// References:
//   - DLMF §25.14
//   - mpmath.lerchphi (high-precision reference)

#ifndef SANGI_SPECIAL_LERCH_HPP
#define SANGI_SPECIAL_LERCH_HPP

#include <math/special/special_common.hpp>
#include <math/special/zeta.hpp>
#include <math/core/mp/Float.hpp>
#include <cmath>
#include <limits>

namespace sangi {

// Forward declaration: hurwitzZeta is declared in Float.hpp
[[nodiscard]] Float hurwitzZeta(const Float& s, const Float& a, int precision);

namespace special {

namespace detail_lerch {

    // Direct Taylor: Σ_{k=0}^∞ z^k / (k+a)^s   (native)
    template<IsNativeFloat T>
    [[nodiscard]] T direct_series_native(T z, T s, T a) {
        T sum = std::pow(a, -s);
        T zk = z;
        const T eps = std::numeric_limits<T>::epsilon();
        // High max_iter to handle |z| close to 1
        const T abs_z = std::abs(z);
        int max_iter = (abs_z > T(0.9)) ? 200000 : 20000;
        int small_count = 0;
        for (int k = 1; k < max_iter; ++k) {
            T term = zk * std::pow(a + T(k), -s);
            sum += term;
            if (k >= 10 && std::abs(term) < eps * std::abs(sum)) {
                if (++small_count >= 3) break;
            } else {
                small_count = 0;
            }
            zk *= z;
        }
        return sum;
    }

    // Hasse-accelerated alternating sum: Σ (-1)^k/(k+a)^s
    //   = Σ_{n=0}^N (1/2^{n+1}) Σ_{k=0}^n (-1)^k C(n,k) (k+a)^{-s}
    template<IsNativeFloat T>
    [[nodiscard]] T hasse_alternating_native(T s, T a, int N = 64) {
        T result = T(0);
        T half_pow = T(0.5);
        for (int n = 0; n <= N; ++n) {
            T inner = T(0);
            T binom = T(1);
            T sign = T(1);
            for (int k = 0; k <= n; ++k) {
                inner += sign * binom * std::pow(a + T(k), -s);
                sign = -sign;
                if (k < n) binom *= T(n - k) / T(k + 1);
            }
            result += half_pow * inner;
            half_pow *= T(0.5);
        }
        return result;
    }

}  // namespace detail_lerch

// ================================================================
// Native floating-point version
// ================================================================

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

    // a non-positive real → pole
    if (a <= T(0)) return std::numeric_limits<T>::quiet_NaN();

    // z = 0
    if (z == T(0)) return std::pow(a, -s);

    // s = 0
    if (s == T(0)) {
        if (std::abs(z) >= T(1)) return std::numeric_limits<T>::quiet_NaN();
        return T(1) / (T(1) - z);
    }

    // z = 1: Hurwitz ζ
    if (z == T(1)) {
        if (s <= T(1)) return std::numeric_limits<T>::quiet_NaN();
        return hurwitzZeta(s, a);
    }

    // z = -1: Hasse-accelerated alternating sum
    if (z == T(-1)) {
        if (s <= T(0)) return std::numeric_limits<T>::quiet_NaN();
        return detail_lerch::hasse_alternating_native(s, a);
    }

    // |z| > 1: analytic continuation not supported
    if (std::abs(z) > T(1))
        return std::numeric_limits<T>::quiet_NaN();

    return detail_lerch::direct_series_native(z, s, a);
}

// ================================================================
// Float arbitrary-precision version
// ================================================================

[[nodiscard]] inline Float lerchPhi(const Float& z, const Float& s, const Float& a, int precision) {
    if (z.isNaN() || s.isNaN() || a.isNaN()) return Float::nan();
    if (a.isZero() || a.isNegative()) return Float::nan();

    int wp = precision + 30;
    Float zw = z; zw.truncateToApprox(wp);
    Float sw = s; sw.truncateToApprox(wp);
    Float aw = a; aw.truncateToApprox(wp);

    Float one = Float::one(wp);
    Float zero = Float::zero(wp);

    if (zw.isZero()) {
        Float r = exp(-sw * log(aw, wp), wp);
        r.setPrecision(precision);
        return r;
    }

    if (sw.isZero()) {
        Float r = one / (one - zw);
        r.setPrecision(precision);
        return r;
    }

    if (zw == one) {
        return hurwitzZeta(sw, aw, precision);
    }

    Float neg_one = -one;
    if (zw == neg_one) {
        // Hasse-accelerated alternating sum (error ~2^{-N}, 1 bit/term). BUGFIX (2026-05-30):
        // N = wp·1.05 misused decimal digits as a bit-term count, capping 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;
        Float result = zero;
        Float half_pow = ldexp(one, -1);
        for (int n = 0; n <= N; ++n) {
            Float inner = zero;
            Float binom = Float::one(wp);
            int sgn = 1;
            for (int k = 0; k <= n; ++k) {
                Float ka = aw + Float(k);
                Float term = exp(-sw * log(ka, wp), wp);
                if (sgn > 0) inner = inner + binom * term;
                else         inner = inner - binom * term;
                sgn = -sgn;
                inner.truncateToApprox(wp);
                if (k < n) {
                    binom = binom * Float(n - k) / Float(k + 1);
                    binom.truncateToApprox(wp);
                }
            }
            result = result + half_pow * inner;
            result.truncateToApprox(wp);
            half_pow = ldexp(half_pow, -1);
        }
        result.setPrecision(precision);
        return result;
    }

    // |z| > 1: not supported
    Float abs_z = abs(zw);
    if (abs_z > one) return Float::nan();
    if (abs_z == one) return Float::nan();

    // Direct series
    Float sum = exp(-sw * log(aw, wp), wp);
    Float zk = zw;
    int max_iter = wp * 8 + 500;
    int small_count = 0;
    for (int k = 1; k < max_iter; ++k) {
        Float ka = aw + Float(k);
        Float term = zk * exp(-sw * log(ka, wp), wp);
        term.truncateToApprox(wp);
        sum = sum + term;
        sum.truncateToApprox(wp);
        Float at = abs(term);
        Float as = abs(sum);
        Float thresh = as * Float::epsilon(wp);
        if (k >= 10 && at < thresh) {
            if (++small_count >= 3) break;
        } else {
            small_count = 0;
        }
        zk = zk * zw;
        zk.truncateToApprox(wp);
    }
    sum.setPrecision(precision);
    return sum;
}

[[nodiscard]] inline Float lerchPhi(const Float& z, const Float& s, const Float& a) {
    return lerchPhi(z, s, a, Float::defaultPrecision());
}

}  // namespace special
}  // namespace sangi

#endif  // SANGI_SPECIAL_LERCH_HPP