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

// lambert_w.hpp
// Lambert W function templates
//
// Provided functions:
//   lambertW0(x)    — principal branch W₀(x),   x ≥ -1/e
//   lambertWm1(x)   — secondary branch W₋₁(x), -1/e ≤ x < 0
//
// The function satisfying W(x)·e^{W(x)} = x.
// Found via Halley iteration (cubic convergence).
//
// Supported types:
//   float, double, long double — Halley iteration
//   Float                      — delegates to the sangi:: implementation
//   Complex<double>            — Halley iteration (complex argument)
//   Complex<Float>             — arbitrary-precision Halley iteration

#ifndef SANGI_SPECIAL_LAMBERT_W_HPP
#define SANGI_SPECIAL_LAMBERT_W_HPP

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

namespace sangi {
namespace special {

// ================================================================
// Lambert W₀(x) — principal branch
// ================================================================
// Defined for x ≥ -1/e. W₀ ≥ -1.
// W₀(0) = 0, W₀(e) = 1, W₀(-1/e) = -1

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

    constexpr T inv_e = T(1) / std::numbers::e_v<T>;

    // x < -1/e → outside the domain
    if (x < -inv_e) return std::numeric_limits<T>::quiet_NaN();

    // x = 0 → W = 0
    if (x == T(0)) return T(0);

    // x = +∞ → W = +∞
    if (std::isinf(x) && x > T(0)) return x;

    // Initial estimate
    T w;
    if (x < T(-0.3)) {
        // Near -1/e: p = √(2(ex+1)), W ≈ -1 + p - p²/3 + 11p³/72
        T p = std::sqrt(T(2) * (std::numbers::e_v<T> * x + T(1)));
        w = T(-1) + p - p * p / T(3) + T(11) * p * p * p / T(72);
    } else if (x <= T(3)) {
        // Small x: initial value close to W ≈ x
        // x ∈ (-0.3, 3]: Fritsch approximation
        if (x <= T(0.5)) {
            w = x * (T(1) - x);
        } else {
            T lnx1 = std::log(x + T(1));
            w = T(0.665) * (T(1) + T(0.0195) * lnx1) * lnx1 + T(0.04);
        }
    } else {
        // Large x: W ≈ ln(x) - ln(ln(x))
        T lnx = std::log(x);
        w = lnx - std::log(lnx);
    }

    // Halley iteration (cubic convergence)
    // w_{n+1} = w - (w·e^w - x) / (e^w·(w+1) - (w+2)(w·e^w - x)/(2w+2))
    for (int i = 0; i < 30; i++) {
        T ew = std::exp(w);
        T wew = w * ew;
        T f = wew - x;

        if (std::abs(f) < std::numeric_limits<T>::epsilon() * std::abs(x + T(1)))
            break;

        T wp1 = w + T(1);
        T denom = ew * wp1 - (w + T(2)) * f / (T(2) * wp1);
        if (denom == T(0)) break;

        w -= f / denom;
    }

    return w;
}

// ================================================================
// Lambert W₋₁(x) — secondary branch
// ================================================================
// Defined for -1/e ≤ x < 0. W₋₁ ≤ -1.
// W₋₁(-1/e) = -1

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

    constexpr T inv_e = T(1) / std::numbers::e_v<T>;

    // x < -1/e or x >= 0 → outside the domain
    if (x < -inv_e || x >= T(0)) return std::numeric_limits<T>::quiet_NaN();

    // x = -1/e → W = -1
    if (std::abs(x + inv_e) < std::numeric_limits<T>::epsilon())
        return T(-1);

    // Initial estimate: W ≈ ln(-x) - ln(-ln(-x))
    T w;
    T lnmx = std::log(-x);
    T lnmlnmx = std::log(-lnmx);
    w = lnmx - lnmlnmx;

    // Correction when x is close to -1/e
    if (x > T(-0.1)) {
        T p = std::sqrt(T(2) * (std::numbers::e_v<T> * x + T(1)));
        w = T(-1) - p - p * p / T(3) - T(11) * p * p * p / T(72);
    }

    // Halley iteration (cubic convergence)
    for (int i = 0; i < 30; i++) {
        T ew = std::exp(w);
        T wew = w * ew;
        T f = wew - x;

        if (std::abs(f) < std::numeric_limits<T>::epsilon() * std::abs(x))
            break;

        T wp1 = w + T(1);
        T denom = ew * wp1 - (w + T(2)) * f / (T(2) * wp1);
        if (denom == T(0)) break;

        w -= f / denom;
    }

    return w;
}

// Float type: use sangi::lambertW0, sangi::lambertWm1 directly.

// ================================================================
// Shared template for the complex Halley iteration
// ================================================================
// Solves f(w) = w·e^w - z = 0 by Halley's method.
// Takes an initial estimate w0 and iterates until convergence.

namespace detail {

template<typename R>
Complex<R> lambertW_halley(const Complex<R>& z, Complex<R> w, R eps, int max_iter) {
    using C = Complex<R>;
    for (int i = 0; i < max_iter; i++) {
        C ew = sangi::exp(w);
        C wew = w * ew;
        C f = wew - z;

        R f_abs = sangi::abs(f);
        R z_abs = sangi::abs(z);
        if (f_abs < eps * (z_abs + R(1)))
            break;

        C wp1 = w + C(R(1));
        C denom = ew * wp1 - (w + C(R(2))) * f / (C(R(2)) * wp1);
        R denom_abs = sangi::abs(denom);
        if (denom_abs == R(0)) break;

        w = w - f / denom;
    }
    return w;
}

} // namespace detail

// ================================================================
// Complex<double> Lambert W₀(z) — principal branch
// ================================================================
// Defined over the entire complex plane. The branch cut is (-∞, -1/e].

[[nodiscard]] inline Complex<double> lambertW0(const Complex<double>& z) {
    using C = Complex<double>;
    constexpr double eps = std::numeric_limits<double>::epsilon();

    // z = 0 → W = 0
    if (sangi::abs(z) == 0.0) return C(0.0);

    // Initial estimate
    C w;
    double r = sangi::abs(z);
    if (r < 0.5) {
        // Small |z|: W ≈ z - z²
        w = z - z * z;
    } else if (r < 3.0) {
        // Intermediate region: build the initial estimate from the real Fritsch approximation
        // Near im=0 a real initial estimate is safe
        C lnz = sangi::log(z);
        double lnz_abs = sangi::abs(lnz);
        if (lnz_abs < 0.5) {
            // log(z) is small → approximation like W ≈ z·(1-z)/(1+z)
            w = z / (C(1.0) + z);
        } else {
            w = lnz - sangi::log(lnz);
        }
    } else {
        // Large |z|: W ≈ log(z) - log(log(z))
        C lnz = sangi::log(z);
        w = lnz - sangi::log(lnz);
    }

    return detail::lambertW_halley<double>(z, w, eps, 50);
}

// ================================================================
// Complex<double> Lambert W₋₁(z) — secondary branch
// ================================================================
// On the real axis, -1/e ≤ x < 0 with W₋₁ ≤ -1.
// In the complex plane this corresponds to W₋₁(z) = W_k(z) (k=-1).
// Initial estimate: log(z) - log(log(z)) - 2πi

[[nodiscard]] inline Complex<double> lambertWm1(const Complex<double>& z) {
    using C = Complex<double>;
    constexpr double eps = std::numeric_limits<double>::epsilon();
    constexpr double two_pi = 2.0 * std::numbers::pi;

    // Initial estimate: asymptotic expansion of W_{-1}
    C lnz = sangi::log(z);
    C shifted = C(lnz.re, lnz.im - two_pi);  // log(z) - 2πi
    C w = shifted - sangi::log(shifted);

    return detail::lambertW_halley<double>(z, w, eps, 50);
}

// ================================================================
// Complex<Float> Lambert W₀(z) — principal branch
// ================================================================

[[nodiscard]] inline Complex<Float> lambertW0(const Complex<Float>& z, int precision) {
    using C = Complex<Float>;
    int wp = precision + 25;
    // PrecisionGuard removed: setResultPrecision(wp) on zw / w propagates req=wp.
    C zw = z;
    zw.re.setResultPrecision(wp);
    zw.im.setResultPrecision(wp);

    // z = 0 → W = 0
    if (sangi::abs(zw).isZero()) {
        Float zero(0); zero.setPrecision(precision);
        return C(zero);
    }

    // Obtain a double-precision initial estimate
    Complex<double> zd(z.re.toDouble(), z.im.toDouble());
    Complex<double> wd = lambertW0(zd);
    C w(Float(wd.re), Float(wd.im));
    w.re.setResultPrecision(wp);
    w.im.setResultPrecision(wp);

    // Converge to arbitrary precision via Halley iteration
    C result = detail::lambertW_halley<Float>(zw, w, Float::epsilon(wp),
                                               10 * (wp / 50 + 1));
    result.re.setPrecision(precision);
    result.im.setPrecision(precision);
    return result;
}

// ================================================================
// Complex<Float> Lambert W₋₁(z) — secondary branch
// ================================================================

[[nodiscard]] inline Complex<Float> lambertWm1(const Complex<Float>& z, int precision) {
    using C = Complex<Float>;
    int wp = precision + 25;
    // PrecisionGuard removed: setResultPrecision(wp) on zw / w propagates req=wp.
    C zw = z;
    zw.re.setResultPrecision(wp);
    zw.im.setResultPrecision(wp);

    // Obtain a double-precision initial estimate
    Complex<double> zd(z.re.toDouble(), z.im.toDouble());
    Complex<double> wd = lambertWm1(zd);
    C w(Float(wd.re), Float(wd.im));
    w.re.setResultPrecision(wp);
    w.im.setResultPrecision(wp);

    // Converge to arbitrary precision via Halley iteration
    C result = detail::lambertW_halley<Float>(zw, w, Float::epsilon(wp),
                                               10 * (wp / 50 + 1));
    result.re.setPrecision(precision);
    result.im.setPrecision(precision);
    return result;
}

} // namespace special
} // namespace sangi

#endif // SANGI_SPECIAL_LAMBERT_W_HPP