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

// weierstrass.hpp
// Weierstrass elliptic functions ℘ / ζ_W / σ_W
//
// Functions provided (take the half-periods ω₁, ω₂ as input):
//   weierstrassP(z, omega1, omega2)      — Weierstrass ℘ function
//   weierstrassZeta(z, omega1, omega2)   — Weierstrass ζ_W function (≠ Riemann ζ)
//   weierstrassSigma(z, omega1, omega2)  — Weierstrass σ_W function
//
// Properties:
//   ℘ is an even function with double periods (2ω₁, 2ω₂), with a pole of order 2 at lattice points
//   (℘')² = 4℘³ - g₂℘ - g₃    (Weierstrass differential equation)
//   ζ_W'(z) = -℘(z), σ_W'(z)/σ_W(z) = ζ_W(z)
//
// Implementation:
//   Uses the theta-function representation from DLMF 23.6
//     ℘(z) = e₃ + (π/(2ω₁) · θ₂(0,q) θ₃(0,q) · θ₄(v,q)/θ₁(v,q))²
//   q is generally complex (the real part of τ = ω₂/ω₁ is nonzero); internally the theta series for complex q is expanded directly.
//
// Supported types (v0.1):
//   Complex<double>  — fully supported (any half-periods / any z)
//   double           — only for real z, ω₁ + purely imaginary ω₂ (rectangular lattice). Real-valued if poles are avoided
//
// TODO:
//   - Float / Complex<Float> arbitrary-precision version (in a future release)
//   - Initialization from the invariants (g₂, g₃) (solve the cubic 4w³ - g₂ w - g₃ = 0)

#ifndef SANGI_SPECIAL_WEIERSTRASS_HPP
#define SANGI_SPECIAL_WEIERSTRASS_HPP

#include <math/special/special_common.hpp>
#include <math/core/Complex.hpp>
#include <cmath>
#include <complex>
#include <limits>
#include <numbers>

namespace sangi {
namespace special {

namespace detail_weier {

// ------------------------------------------------------------------
// Internal theta functions (arbitrary complex v, complex q)
// theta.hpp handles only real q (or the Float version of Complex<R>),
// so this is a local implementation for the "general complex q" needed by the Weierstrass functions.
// ------------------------------------------------------------------

/// θ₁(v, q) = 2 Σ_{n=0}^∞ (-1)^n q^{(n+1/2)²} sin((2n+1)v)
template<typename R>
[[nodiscard]] Complex<R> theta1(const Complex<R>& v, const Complex<R>& q,
                                int max_iter = 200) {
    using C = Complex<R>;
    C sum(R(0));
    R sign = R(1);
    R eps = std::numeric_limits<R>::epsilon();

    for (int n = 0; n < max_iter; n++) {
        R k = R(n) + R(0.5);
        C qpow = pow(q, C(k * k));
        C term = qpow * sin(C(R(2 * n + 1)) * v);
        if (sign < R(0)) term = -term;
        sum = sum + term;
        if (n >= 4 && abs(term) < eps * (R(1) + abs(sum))) break;
        sign = -sign;
    }
    return C(R(2)) * sum;
}

/// θ₂(v, q) = 2 Σ_{n=0}^∞ q^{(n+1/2)²} cos((2n+1)v)
template<typename R>
[[nodiscard]] Complex<R> theta2(const Complex<R>& v, const Complex<R>& q,
                                int max_iter = 200) {
    using C = Complex<R>;
    C sum(R(0));
    R eps = std::numeric_limits<R>::epsilon();

    for (int n = 0; n < max_iter; n++) {
        R k = R(n) + R(0.5);
        C qpow = pow(q, C(k * k));
        C term = qpow * cos(C(R(2 * n + 1)) * v);
        sum = sum + term;
        if (n >= 4 && abs(term) < eps * (R(1) + abs(sum))) break;
    }
    return C(R(2)) * sum;
}

/// θ₃(v, q) = 1 + 2 Σ_{n=1}^∞ q^{n²} cos(2nv)
template<typename R>
[[nodiscard]] Complex<R> theta3(const Complex<R>& v, const Complex<R>& q,
                                int max_iter = 200) {
    using C = Complex<R>;
    C sum(R(0));
    R eps = std::numeric_limits<R>::epsilon();

    for (int n = 1; n < max_iter; n++) {
        C qpow = pow(q, C(R(n) * R(n)));
        C term = qpow * cos(C(R(2 * n)) * v);
        sum = sum + term;
        if (n >= 4 && abs(term) < eps * (R(1) + abs(sum))) break;
    }
    return C(R(1)) + C(R(2)) * sum;
}

/// θ₄(v, q) = 1 + 2 Σ_{n=1}^∞ (-1)^n q^{n²} cos(2nv)
template<typename R>
[[nodiscard]] Complex<R> theta4(const Complex<R>& v, const Complex<R>& q,
                                int max_iter = 200) {
    using C = Complex<R>;
    C sum(R(0));
    R sign = R(-1);
    R eps = std::numeric_limits<R>::epsilon();

    for (int n = 1; n < max_iter; n++) {
        C qpow = pow(q, C(R(n) * R(n)));
        C term = qpow * cos(C(R(2 * n)) * v);
        if (sign < R(0)) term = -term;
        sum = sum + term;
        if (n >= 4 && abs(term) < eps * (R(1) + abs(sum))) break;
        sign = -sign;
    }
    return C(R(1)) + C(R(2)) * sum;
}

// ∂θ₁/∂v at v=0  =  2 Σ (-1)^n q^{(n+1/2)²} (2n+1)
template<typename R>
[[nodiscard]] Complex<R> theta1_prime_at_zero(const Complex<R>& q,
                                              int max_iter = 200) {
    using C = Complex<R>;
    C sum(R(0));
    R sign = R(1);
    R eps = std::numeric_limits<R>::epsilon();

    for (int n = 0; n < max_iter; n++) {
        R k = R(n) + R(0.5);
        C qpow = pow(q, C(k * k));
        C term = qpow * C(R(2 * n + 1));
        if (sign < R(0)) term = -term;
        sum = sum + term;
        if (n >= 4 && abs(term) < eps * (R(1) + abs(sum))) break;
        sign = -sign;
    }
    return C(R(2)) * sum;
}

// ∂³θ₁/∂v³ at v=0  =  -2 Σ (-1)^n q^{(n+1/2)²} (2n+1)³
template<typename R>
[[nodiscard]] Complex<R> theta1_triple_prime_at_zero(const Complex<R>& q,
                                                    int max_iter = 200) {
    using C = Complex<R>;
    C sum(R(0));
    R sign = R(1);
    R eps = std::numeric_limits<R>::epsilon();

    for (int n = 0; n < max_iter; n++) {
        R k = R(n) + R(0.5);
        C qpow = pow(q, C(k * k));
        R cube = R(2 * n + 1) * R(2 * n + 1) * R(2 * n + 1);
        C term = qpow * C(cube);
        if (sign < R(0)) term = -term;
        sum = sum + term;
        if (n >= 4 && abs(term) < eps * (R(1) + abs(sum))) break;
        sign = -sign;
    }
    return C(R(-2)) * sum;
}

// ∂θ₁/∂v at general v  =  2 Σ (-1)^n q^{(n+1/2)²} (2n+1) cos((2n+1)v)
template<typename R>
[[nodiscard]] Complex<R> theta1_prime(const Complex<R>& v, const Complex<R>& q,
                                     int max_iter = 200) {
    using C = Complex<R>;
    C sum(R(0));
    R sign = R(1);
    R eps = std::numeric_limits<R>::epsilon();

    for (int n = 0; n < max_iter; n++) {
        R k = R(n) + R(0.5);
        C qpow = pow(q, C(k * k));
        C term = qpow * C(R(2 * n + 1)) * cos(C(R(2 * n + 1)) * v);
        if (sign < R(0)) term = -term;
        sum = sum + term;
        if (n >= 4 && abs(term) < eps * (R(1) + abs(sum))) break;
        sign = -sign;
    }
    return C(R(2)) * sum;
}

// Validity check and normalization of the half-periods (ω₁, ω₂):
// Invalid unless τ = ω₂/ω₁ is in the upper half-plane (Im(τ) > 0).
// When valid, returns q = exp(iπτ) (|q| < 1 is guaranteed).
struct LatticeData {
    Complex<double> alpha;  // π/(2ω₁)
    Complex<double> q;      // nome
    bool valid;
};

inline LatticeData prepare_lattice(const Complex<double>& omega1,
                                   const Complex<double>& omega2) {
    LatticeData out{};
    out.valid = false;
    if (omega1.re == 0.0 && omega1.im == 0.0) return out;
    Complex<double> tau = omega2 / omega1;
    if (!(tau.im > 0.0)) return out;
    constexpr double PI = std::numbers::pi_v<double>;
    Complex<double> i_pi_tau(-PI * tau.im, PI * tau.re);  // iπτ
    out.q = exp(i_pi_tau);
    out.alpha = Complex<double>(PI, 0.0) / (Complex<double>(2.0, 0.0) * omega1);
    out.valid = true;
    return out;
}

}  // namespace detail_weier

// ================================================================
// Weierstrass ℘ function (Complex<double>)
// ================================================================
// ℘(z) = e₃ + (π/(2ω₁) · θ₂(0,q) θ₃(0,q) · θ₄(v,q)/θ₁(v,q))²
// e₃ = -(π/(2ω₁))² · (θ₂⁴(0,q) + θ₃⁴(0,q)) / 3

[[nodiscard]] inline Complex<double> weierstrassP(const Complex<double>& z,
                                                  const Complex<double>& omega1,
                                                  const Complex<double>& omega2) {
    using C = Complex<double>;
    auto L = detail_weier::prepare_lattice(omega1, omega2);
    if (!L.valid) return C(std::numeric_limits<double>::quiet_NaN());

    C zero(0.0, 0.0);
    C v = L.alpha * z;

    C t2_0 = detail_weier::theta2<double>(zero, L.q);
    C t3_0 = detail_weier::theta3<double>(zero, L.q);
    C t1_v = detail_weier::theta1<double>(v, L.q);
    C t4_v = detail_weier::theta4<double>(v, L.q);

    // Lattice points (θ₁(v) ≈ 0) are poles
    if (abs(t1_v) < std::numeric_limits<double>::min() * 1e10) {
        return C(std::numeric_limits<double>::infinity());
    }

    C alpha_pi = L.alpha;  // π/(2ω₁)

    // e₃ = -(α)² · (θ₂⁴ + θ₃⁴) / 3
    C t2_4 = t2_0 * t2_0 * t2_0 * t2_0;
    C t3_4 = t3_0 * t3_0 * t3_0 * t3_0;
    C alpha_sq = alpha_pi * alpha_pi;
    C e3 = -alpha_sq * (t2_4 + t3_4) / C(3.0, 0.0);

    // ratio = α · θ₂(0) · θ₃(0) · θ₄(v) / θ₁(v)
    C ratio = alpha_pi * t2_0 * t3_0 * t4_v / t1_v;
    return e3 + ratio * ratio;
}

// ================================================================
// Weierstrass ζ_W function (not the Riemann ζ)
// ================================================================
// ζ_W(z) = η₁ z / ω₁ + (π/(2ω₁)) · θ₁'(v,q) / θ₁(v,q)
// η₁ = -π²/(12 ω₁) · θ₁'''(0,q) / θ₁'(0,q)

[[nodiscard]] inline Complex<double> weierstrassZeta(const Complex<double>& z,
                                                     const Complex<double>& omega1,
                                                     const Complex<double>& omega2) {
    using C = Complex<double>;
    auto L = detail_weier::prepare_lattice(omega1, omega2);
    if (!L.valid) return C(std::numeric_limits<double>::quiet_NaN());

    constexpr double PI = std::numbers::pi_v<double>;
    C v = L.alpha * z;

    C t1p_0 = detail_weier::theta1_prime_at_zero<double>(L.q);
    C t1ppp_0 = detail_weier::theta1_triple_prime_at_zero<double>(L.q);
    C t1_v = detail_weier::theta1<double>(v, L.q);
    C t1p_v = detail_weier::theta1_prime<double>(v, L.q);

    if (abs(t1_v) < std::numeric_limits<double>::min() * 1e10) {
        return C(std::numeric_limits<double>::infinity());
    }

    // η₁ = -π²/(12 ω₁) · θ₁'''(0)/θ₁'(0)
    C eta1 = -C(PI * PI, 0.0) / (C(12.0, 0.0) * omega1) * t1ppp_0 / t1p_0;

    return eta1 * z / omega1 + L.alpha * t1p_v / t1_v;
}

// ================================================================
// Weierstrass σ_W function
// ================================================================
// σ_W(z) = (2ω₁/π) · θ₁(v,q) / θ₁'(0,q) · exp(η₁ z² / (2ω₁))

[[nodiscard]] inline Complex<double> weierstrassSigma(const Complex<double>& z,
                                                      const Complex<double>& omega1,
                                                      const Complex<double>& omega2) {
    using C = Complex<double>;
    auto L = detail_weier::prepare_lattice(omega1, omega2);
    if (!L.valid) return C(std::numeric_limits<double>::quiet_NaN());

    constexpr double PI = std::numbers::pi_v<double>;
    C v = L.alpha * z;

    C t1p_0 = detail_weier::theta1_prime_at_zero<double>(L.q);
    C t1ppp_0 = detail_weier::theta1_triple_prime_at_zero<double>(L.q);
    C t1_v = detail_weier::theta1<double>(v, L.q);

    C eta1 = -C(PI * PI, 0.0) / (C(12.0, 0.0) * omega1) * t1ppp_0 / t1p_0;

    C prefactor = (C(2.0, 0.0) * omega1 / C(PI, 0.0)) * t1_v / t1p_0;
    C expfactor = exp(eta1 * z * z / (C(2.0, 0.0) * omega1));
    return prefactor * expfactor;
}

// ================================================================
// double overloads (real z, real half-period ω₁ + purely imaginary half-period ω₂)
// ----------------------------------------------------------------
// Note: for a rectangular lattice (ω₁ real, ω₂ purely imaginary), ℘(z) is real for real z.

[[nodiscard]] inline double weierstrassP(double z, double omega1, double omega2_imag) {
    auto r = weierstrassP(Complex<double>(z, 0.0),
                          Complex<double>(omega1, 0.0),
                          Complex<double>(0.0, omega2_imag));
    return r.re;
}

[[nodiscard]] inline double weierstrassZeta(double z, double omega1, double omega2_imag) {
    auto r = weierstrassZeta(Complex<double>(z, 0.0),
                             Complex<double>(omega1, 0.0),
                             Complex<double>(0.0, omega2_imag));
    return r.re;
}

[[nodiscard]] inline double weierstrassSigma(double z, double omega1, double omega2_imag) {
    auto r = weierstrassSigma(Complex<double>(z, 0.0),
                              Complex<double>(omega1, 0.0),
                              Complex<double>(0.0, omega2_imag));
    return r.re;
}

}  // namespace special
}  // namespace sangi

#endif  // SANGI_SPECIAL_WEIERSTRASS_HPP
