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

// Complex.hpp
// Complex-number template class
//
// Features:
//   - Works with any numeric type T = double, float, Int, Float, Rational, ...
//   - Custom implementation distinct from std::complex<double> (multi-precision compatible)
//   - 3-multiplication complex product via Strassen's method
//   - Type constraints via C++23 concepts
//   - T = double/float: Smith's division, hypot-based abs, NaN/Inf propagation (2026-03-20)

#pragma once

#include <iosfwd>
#include <sstream>
#include <string>
#include <cmath>
#include <cstddef>
#include <complex>
#include <type_traits>
#include <concepts>
#include <limits>

namespace sangi {

// Constraints required on the element type of Complex<T>
template<typename T>
concept ComplexScalar = requires(T a, T b) {
    { a + b } -> std::convertible_to<T>;
    { a - b } -> std::convertible_to<T>;
    { a * b } -> std::convertible_to<T>;
    { a / b } -> std::convertible_to<T>;
    { -a } -> std::convertible_to<T>;
    { T(0) };
    { T(1) };
};

// ================================================================
// Complex<T> class
// ================================================================

template<ComplexScalar T>
class Complex {
public:
    T re;  // real part
    T im;  // imaginary part

    // ============================================================
    // Constructors
    // ============================================================

    // Default: zero-initialize
    constexpr Complex() noexcept(noexcept(T(0))) : re(T(0)), im(T(0)) {}

    // From a real number (imag = 0)
    constexpr Complex(const T& real) noexcept(noexcept(T(0)))
        : re(real), im(T(0)) {}

    // Real and imaginary parts
    constexpr Complex(const T& real, const T& imag) noexcept
        : re(real), im(imag) {}

    // Conversion from a different element type
    template<ComplexScalar U>
    constexpr explicit Complex(const Complex<U>& other)
        : re(T(other.re)), im(T(other.im)) {}

    // Conversion from std::complex (T = double/float/long double)
    template<typename U>
        requires std::is_floating_point_v<U> && std::is_convertible_v<U, T>
    constexpr Complex(const std::complex<U>& sc)
        : re(T(sc.real())), im(T(sc.imag())) {}

    // Conversion to std::complex (T = double/float/long double)
    template<typename U = T>
        requires std::is_floating_point_v<U>
    constexpr operator std::complex<U>() const {
        return std::complex<U>(static_cast<U>(re), static_cast<U>(im));
    }

    // ============================================================
    // Accessors
    // ============================================================

    [[nodiscard]] constexpr const T& real() const noexcept { return re; }
    [[nodiscard]] constexpr const T& imag() const noexcept { return im; }
    [[nodiscard]] constexpr T& real() noexcept { return re; }
    [[nodiscard]] constexpr T& imag() noexcept { return im; }
    constexpr void real(const T& val) { re = val; }
    constexpr void imag(const T& val) { im = val; }

    // ============================================================
    // Comparison operators
    // ============================================================

    [[nodiscard]] constexpr bool operator==(const Complex& rhs) const {
        return re == rhs.re && im == rhs.im;
    }
    [[nodiscard]] constexpr bool operator!=(const Complex& rhs) const {
        return !(*this == rhs);
    }

    // Comparison with a real number
    [[nodiscard]] constexpr bool operator==(const T& rhs) const {
        return re == rhs && im == T(0);
    }
    [[nodiscard]] constexpr bool operator!=(const T& rhs) const {
        return !(*this == rhs);
    }

    friend constexpr bool operator==(const T& lhs, const Complex& rhs) {
        return rhs == lhs;
    }
    friend constexpr bool operator!=(const T& lhs, const Complex& rhs) {
        return !(rhs == lhs);
    }

    // ============================================================
    // Unary operators
    // ============================================================

    [[nodiscard]] constexpr Complex operator+() const { return *this; }
    [[nodiscard]] constexpr Complex operator-() const { return Complex(-re, -im); }

    // ============================================================
    // Complex + complex
    // ============================================================

    [[nodiscard]] constexpr Complex operator+(const Complex& rhs) const {
        return Complex(re + rhs.re, im + rhs.im);
    }
    [[nodiscard]] constexpr Complex operator-(const Complex& rhs) const {
        return Complex(re - rhs.re, im - rhs.im);
    }

    // Floating-point: ordinary 4 multiplications + 2 additions/subtractions (similar cost)
    // Multi-precision: Strassen's method, 3 multiplications + 5 additions/subtractions (faster because multiplies are expensive)
    [[nodiscard]] constexpr Complex operator*(const Complex& rhs) const {
        if constexpr (std::is_floating_point_v<T>) {
            return Complex(re * rhs.re - im * rhs.im,
                           re * rhs.im + im * rhs.re);
        } else {
            T ac = re * rhs.re;
            T bd = im * rhs.im;
            return Complex(ac - bd, (re + im) * (rhs.re + rhs.im) - ac - bd);
        }
    }

    // Division via Smith's method (overflow-safe for T = double/float)
    // For multi-precision types, keep the naive implementation
    [[nodiscard]] constexpr Complex operator/(const Complex& rhs) const {
        if constexpr (std::is_floating_point_v<T>) {
            return divSmith(rhs);
        } else {
            T denom = rhs.re * rhs.re + rhs.im * rhs.im;
            if (denom == T(0)) {
                return Complex(std::numeric_limits<T>::quiet_NaN(), std::numeric_limits<T>::quiet_NaN());
            }
            return Complex(
                (re * rhs.re + im * rhs.im) / denom,
                (im * rhs.re - re * rhs.im) / denom
            );
        }
    }

    // ============================================================
    // Complex + real / real + complex
    // ============================================================

    [[nodiscard]] constexpr Complex operator+(const T& rhs) const {
        return Complex(re + rhs, im);
    }
    [[nodiscard]] constexpr Complex operator-(const T& rhs) const {
        return Complex(re - rhs, im);
    }
    [[nodiscard]] constexpr Complex operator*(const T& rhs) const {
        return Complex(re * rhs, im * rhs);
    }
    [[nodiscard]] constexpr Complex operator/(const T& rhs) const {
        return Complex(re / rhs, im / rhs);
    }

    friend constexpr Complex operator+(const T& lhs, const Complex& rhs) {
        return Complex(lhs + rhs.re, rhs.im);
    }
    friend constexpr Complex operator-(const T& lhs, const Complex& rhs) {
        return Complex(lhs - rhs.re, -rhs.im);
    }
    friend constexpr Complex operator*(const T& lhs, const Complex& rhs) {
        return Complex(lhs * rhs.re, lhs * rhs.im);
    }
    // Real / complex (Smith's method)
    friend constexpr Complex operator/(const T& lhs, const Complex& rhs) {
        if constexpr (std::is_floating_point_v<T>) {
            return Complex(lhs, T(0)).divSmith(rhs);
        } else {
            T denom = rhs.re * rhs.re + rhs.im * rhs.im;
            if (denom == T(0)) {
                return Complex(std::numeric_limits<T>::quiet_NaN(), std::numeric_limits<T>::quiet_NaN());
            }
            return Complex(lhs * rhs.re / denom, -lhs * rhs.im / denom);
        }
    }

    // ============================================================
    // Compound assignment operators
    // ============================================================

    constexpr Complex& operator+=(const Complex& rhs) {
        re += rhs.re; im += rhs.im; return *this;
    }
    constexpr Complex& operator-=(const Complex& rhs) {
        re -= rhs.re; im -= rhs.im; return *this;
    }
    constexpr Complex& operator*=(const Complex& rhs) {
        *this = *this * rhs; return *this;
    }
    constexpr Complex& operator/=(const Complex& rhs) {
        *this = *this / rhs; return *this;
    }

    constexpr Complex& operator+=(const T& rhs) { re += rhs; return *this; }
    constexpr Complex& operator-=(const T& rhs) { re -= rhs; return *this; }
    constexpr Complex& operator*=(const T& rhs) { re *= rhs; im *= rhs; return *this; }
    constexpr Complex& operator/=(const T& rhs) { re /= rhs; im /= rhs; return *this; }

    // ============================================================
    // Member functions
    // ============================================================

    // |z|^2 = re^2 + im^2
    [[nodiscard]] constexpr T normSq() const { return re * re + im * im; }

    // String conversion
    [[nodiscard]] std::string toString() const {
        std::ostringstream osRe, osIm;
        osRe << ((re == T(0)) ? T(0) : re);
        osIm << ((im == T(0)) ? T(0) : im);
        std::string sRe = osRe.str();
        std::string sIm = osIm.str();

        if (sIm == "0" || sIm == "0.")
            return sRe;

        std::string s;
        if (sRe == "0" || sRe == "0.") sRe = "";

        if (sIm == "1")       s = sRe + "+i";
        else if (sIm == "-1") s = sRe + "-i";
        else if (!sIm.empty() && sIm[0] == '-') s = sRe + sIm + "i";
        else s = sRe + "+" + sIm + "i";

        return s;
    }

    // Stream output
    friend std::ostream& operator<<(std::ostream& os, const Complex& z) {
        return os << z.toString();
    }

private:
    // Smith's method: always divide by the larger-magnitude component to avoid
    // intermediate overflow. Selecting the pivot component (rather than two
    // symmetric branches that each divide by one component) keeps the divisor
    // from ever being the smaller component, so a purely-real or purely-imaginary
    // divisor cannot constant-fold into a division by zero in a not-taken branch
    // (which would otherwise raise a spurious MSVC C4723). Numerically identical
    // to the two-branch form, term for term.
    [[nodiscard]] constexpr Complex divSmith(const Complex& rhs) const {
        using std::abs;
        T are = abs(rhs.re);
        T aim = abs(rhs.im);
        if (are == T(0) && aim == T(0)) {
            return Complex(std::numeric_limits<T>::quiet_NaN(),
                           std::numeric_limits<T>::quiet_NaN());
        }
        const bool pivotRe = (are >= aim);
        const T piv = pivotRe ? rhs.re : rhs.im;   // larger magnitude, guaranteed nonzero
        const T oth = pivotRe ? rhs.im : rhs.re;
        const T r = oth / piv;
        const T denom = piv + oth * r;
        const T reN = pivotRe ? (re + im * r) : (re * r + im);
        const T imN = pivotRe ? (im - re * r) : (im * r - re);
        return Complex(reN / denom, imN / denom);
    }
};

// ================================================================
// Memory-layout guarantees — SIMD / reinterpret_cast compatible
// Complex<T> stores re and im contiguously and shares the same layout as std::complex<T>.
// ================================================================
static_assert(sizeof(Complex<float>) == 2 * sizeof(float),
              "Complex<float> must be tightly packed");
static_assert(sizeof(Complex<double>) == 2 * sizeof(double),
              "Complex<double> must be tightly packed");
static_assert(offsetof(Complex<float>, re) == 0,
              "Complex<float>::re must be at offset 0");
static_assert(offsetof(Complex<float>, im) == sizeof(float),
              "Complex<float>::im must follow re immediately");
static_assert(offsetof(Complex<double>, re) == 0,
              "Complex<double>::re must be at offset 0");
static_assert(offsetof(Complex<double>, im) == sizeof(double),
              "Complex<double>::im must follow re immediately");

// ================================================================
// Free functions (math functions)
// ================================================================

// --- Basic functions ---

template<ComplexScalar T>
[[nodiscard]] constexpr const T& real(const Complex<T>& z) { return z.re; }

template<ComplexScalar T>
[[nodiscard]] constexpr const T& imag(const Complex<T>& z) { return z.im; }

template<ComplexScalar T>
[[nodiscard]] constexpr Complex<T> conj(const Complex<T>& z) {
    return Complex<T>(z.re, -z.im);
}

// |z|^2 (STL-compatible: norm)
template<ComplexScalar T>
[[nodiscard]] constexpr T normSq(const Complex<T>& z) {
    return z.re * z.re + z.im * z.im;
}

// |z|^2 (STL-compatible name)
template<ComplexScalar T>
[[nodiscard]] constexpr T norm(const Complex<T>& z) {
    return z.re * z.re + z.im * z.im;
}

// |z| (overflow-safe via std::hypot when T = double/float)
template<ComplexScalar T>
[[nodiscard]] T abs(const Complex<T>& z) {
    if constexpr (std::is_floating_point_v<T>) {
        return std::hypot(z.re, z.im);
    } else {
        using std::sqrt;
        return sqrt(z.re * z.re + z.im * z.im);
    }
}

// arg(z)
template<ComplexScalar T>
[[nodiscard]] T arg(const Complex<T>& z) {
    using std::atan2;
    return atan2(z.im, z.re);
}

// --- Polar form ---

template<ComplexScalar T>
[[nodiscard]] Complex<T> polar(const T& r, const T& theta) {
    using std::cos; using std::sin;
    return Complex<T>(r * cos(theta), r * sin(theta));
}

// e^(i*omega)
template<ComplexScalar T>
[[nodiscard]] Complex<T> expI(const T& omega) {
    using std::cos; using std::sin;
    return Complex<T>(cos(omega), sin(omega));
}

// --- Exponential / logarithm ---

template<ComplexScalar T>
[[nodiscard]] Complex<T> exp(const Complex<T>& z) {
    using std::exp; using std::cos; using std::sin;
    T r = exp(z.re);
    return Complex<T>(r * cos(z.im), r * sin(z.im));
}

template<ComplexScalar T>
[[nodiscard]] Complex<T> log(const Complex<T>& z) {
    using std::log;
    return Complex<T>(log(abs(z)), arg(z));
}

// log10(z) = log(z) * log10(e)  where log10(e) = 1/ln(10)
template<ComplexScalar T>
[[nodiscard]] Complex<T> log10(const Complex<T>& z) {
    if constexpr (std::is_floating_point_v<T>) {
        static const T inv_ln10 = T(1) / std::log(T(10));
        return sangi::log(z) * inv_ln10;
    } else {
        // For Float etc.: use the cached constant T::log10e() if available
        return sangi::log(z) * T::log10e();
    }
}

// --- Exponentiation ---

// z^n (integer power)
template<ComplexScalar T>
[[nodiscard]] Complex<T> pow(const Complex<T>& z, int n) {
    if (n == 0) return Complex<T>(T(1));
    if (n < 0) return pow(Complex<T>(T(1)) / z, -n);

    Complex<T> result(T(1));
    Complex<T> base = z;
    unsigned int m = static_cast<unsigned int>(n);
    while (m > 0) {
        if (m & 1) result *= base;
        base *= base;
        m >>= 1;
    }
    return result;
}

// z^w (complex power)
template<ComplexScalar T>
[[nodiscard]] Complex<T> pow(const Complex<T>& z, const Complex<T>& w) {
    if (z == T(0)) return Complex<T>(T(0));
    return exp(w * log(z));
}

// z^a (real power)
template<ComplexScalar T>
[[nodiscard]] Complex<T> pow(const Complex<T>& z, const T& a) {
    if (z == T(0)) return Complex<T>(T(0));
    return exp(a * log(z));
}

// --- Square root ---
// Source: libg++ 2.7.1

template<ComplexScalar T>
[[nodiscard]] Complex<T> sqrt(const Complex<T>& z) {
    using std::sqrt;
    T r = abs(z);
    if (r == T(0)) return Complex<T>(T(0), T(0));

    T x, y;
    if (z.re > T(0)) {
        x = sqrt((r + z.re) / T(2));
        y = z.im / (T(2) * x);
    } else {
        y = sqrt((r - z.re) / T(2));
        if (z.im < T(0)) y = -y;
        x = z.im / (T(2) * y);
    }
    return Complex<T>(x, y);
}

// --- Trigonometric functions ---

template<ComplexScalar T>
[[nodiscard]] Complex<T> sin(const Complex<T>& z) {
    // sin(a+bi) = sin(a)cosh(b) + i*cos(a)sinh(b)
    using std::sin; using std::cos; using std::sinh; using std::cosh;
    return Complex<T>(sin(z.re) * cosh(z.im), cos(z.re) * sinh(z.im));
}

template<ComplexScalar T>
[[nodiscard]] Complex<T> cos(const Complex<T>& z) {
    // cos(a+bi) = cos(a)cosh(b) - i*sin(a)sinh(b)
    using std::sin; using std::cos; using std::sinh; using std::cosh;
    return Complex<T>(cos(z.re) * cosh(z.im), -sin(z.re) * sinh(z.im));
}

template<ComplexScalar T>
[[nodiscard]] Complex<T> tan(const Complex<T>& z) {
    return sin(z) / cos(z);
}

// --- Hyperbolic functions ---

template<ComplexScalar T>
[[nodiscard]] Complex<T> sinh(const Complex<T>& z) {
    // sinh(z) = -i * sin(i*z)
    using std::sinh; using std::cosh; using std::sin; using std::cos;
    return Complex<T>(sinh(z.re) * cos(z.im), cosh(z.re) * sin(z.im));
}

template<ComplexScalar T>
[[nodiscard]] Complex<T> cosh(const Complex<T>& z) {
    // cosh(z) = cos(i*z)
    using std::sinh; using std::cosh; using std::sin; using std::cos;
    return Complex<T>(cosh(z.re) * cos(z.im), sinh(z.re) * sin(z.im));
}

template<ComplexScalar T>
[[nodiscard]] Complex<T> tanh(const Complex<T>& z) {
    return sinh(z) / cosh(z);
}

// --- Inverse trigonometric functions ---

template<ComplexScalar T>
[[nodiscard]] Complex<T> asin(const Complex<T>& z) {
    // asin(z) = -i * log(i*z + sqrt(1 - z^2))
    Complex<T> iz(-z.im, z.re);  // i*z
    Complex<T> w = sqrt(Complex<T>(T(1)) - z * z);
    Complex<T> lv = log(iz + w);
    return Complex<T>(lv.im, -lv.re);  // -i * lv
}

template<ComplexScalar T>
[[nodiscard]] Complex<T> acos(const Complex<T>& z) {
    // acos(z) = -i * log(z + i*sqrt(1 - z^2))
    Complex<T> w = sqrt(Complex<T>(T(1)) - z * z);
    Complex<T> iw(-w.im, w.re);  // i*w
    Complex<T> lv = log(z + iw);
    return Complex<T>(lv.im, -lv.re);  // -i * lv
}

template<ComplexScalar T>
[[nodiscard]] Complex<T> atan(const Complex<T>& z) {
    // atan(z) = (1/2i) * log((1+iz)/(1-iz))
    Complex<T> iz(-z.im, z.re);  // i*z
    Complex<T> lv = log((Complex<T>(T(1)) + iz) / (Complex<T>(T(1)) - iz));
    // (1/2i) * lv = lv / (2i) = lv * (-i/2)
    return Complex<T>(lv.im / T(2), -lv.re / T(2));
}

// --- Inverse hyperbolic functions ---

template<ComplexScalar T>
[[nodiscard]] Complex<T> asinh(const Complex<T>& z) {
    return log(z + sqrt(z * z + Complex<T>(T(1))));
}

template<ComplexScalar T>
[[nodiscard]] Complex<T> acosh(const Complex<T>& z) {
    return log(z + sqrt(z * z - Complex<T>(T(1))));
}

template<ComplexScalar T>
[[nodiscard]] Complex<T> atanh(const Complex<T>& z) {
    Complex<T> one(T(1));
    return log((one + z) / (one - z)) / T(2);
}

// --- proj (Riemann-sphere projection) ---

template<ComplexScalar T>
[[nodiscard]] Complex<T> proj(const Complex<T>& z) {
    if constexpr (std::is_floating_point_v<T>) {
        if (std::isinf(z.re) || std::isinf(z.im)) {
            return Complex<T>(std::numeric_limits<T>::infinity(),
                              std::copysign(T(0), z.im));
        }
    }
    return z;
}

// ================================================================
// User-defined literals
// ================================================================

namespace literals {

[[nodiscard]] constexpr Complex<double> operator""_i(long double val) {
    return Complex<double>(0.0, static_cast<double>(val));
}

[[nodiscard]] constexpr Complex<double> operator""_i(unsigned long long val) {
    return Complex<double>(0.0, static_cast<double>(val));
}

[[nodiscard]] constexpr Complex<float> operator""_if(long double val) {
    return Complex<float>(0.0f, static_cast<float>(val));
}

[[nodiscard]] constexpr Complex<float> operator""_if(unsigned long long val) {
    return Complex<float>(0.0f, static_cast<float>(val));
}

} // namespace literals

// ================================================================
// Type traits
// ================================================================

// Traits that test whether a type is a Complex
template<typename T>
struct is_complex : std::false_type {};

template<ComplexScalar T>
struct is_complex<Complex<T>> : std::true_type {};

template<typename T>
inline constexpr bool is_complex_v = is_complex<T>::value;

// ================================================================
// numeric_traits specialization
// ================================================================

// Forward declaration
struct complex_tag;
template<typename T> struct numeric_traits;

template<ComplexScalar T>
struct numeric_traits<Complex<T>> {
    using value_type = Complex<T>;
    using real_type = T;
    using category = complex_tag;

    static constexpr bool is_supported = true;
    static constexpr bool is_complex = true;
    static constexpr bool is_integer = false;
    static constexpr bool is_floating_point = false;

    static Complex<T> zero() { return Complex<T>(T(0), T(0)); }
    static Complex<T> one() { return Complex<T>(T(1), T(0)); }
    static T epsilon() { return numeric_traits<T>::epsilon(); }

    static T abs(const Complex<T>& value) {
        return sangi::abs(value);
    }

    static Complex<T> conj(const Complex<T>& value) {
        return sangi::conj(value);
    }

    static T norm(const Complex<T>& value) {
        return sangi::normSq(value);
    }

    static bool pivotBetter(const Complex<T>& a, const Complex<T>& b) {
        return abs(a) > abs(b);
    }

    static bool isNaN(const Complex<T>& value) {
        if constexpr (std::is_floating_point_v<T>) {
            return std::isnan(value.re) || std::isnan(value.im);
        } else {
            return false;
        }
    }

    static bool isInfinite(const Complex<T>& value) {
        if constexpr (std::is_floating_point_v<T>) {
            return std::isinf(value.re) || std::isinf(value.im);
        } else {
            return false;
        }
    }

    static bool isFinite(const Complex<T>& value) {
        if constexpr (std::is_floating_point_v<T>) {
            return std::isfinite(value.re) && std::isfinite(value.im);
        } else {
            return true;
        }
    }

    static int getSign(const Complex<T>& value) {
        if (value.re != T(0)) {
            // Prefer the sign of the real part
            if constexpr (std::is_arithmetic_v<T>) {
                return value.re < T(0) ? -1 : 1;
            } else {
                return 1; // assume non-arithmetic types are positive
            }
        }
        if (value.im != T(0)) {
            if constexpr (std::is_arithmetic_v<T>) {
                return value.im < T(0) ? -1 : 1;
            } else {
                return 1;
            }
        }
        return 0;
    }
};

} // namespace sangi

// Add Complex<T> abs/sqrt overloads to the std namespace (for ADL fallback)
namespace std {
    template<typename T>
    auto abs(const sangi::Complex<T>& z) { return sangi::abs(z); }

    template<typename T>
    sangi::Complex<T> sqrt(const sangi::Complex<T>& z) { return sangi::sqrt(z); }
}