"""
GPU-accelerated gamma functions for Apple MLX.

Implements lgamma, gamma, digamma, and beta using vectorized MLX operations.
All computations run on the Apple GPU via MLX's lazy evaluation.

Algorithms:
    lgamma  -- Lanczos approximation (g=7, 9 coefficients) + reflection
    gamma   -- exp(lgamma) with sign handling
    digamma -- asymptotic expansion + recurrence + reflection
    beta    -- exp(lgamma(a) + lgamma(b) - lgamma(a+b))

References:
    Lanczos, C. (1964). SIAM J. Numer. Anal. 1, 86-96.
    Pugh, G. R. (2004). PhD thesis, University of British Columbia.
"""

from __future__ import annotations

import math

import mlx.core as mx

# ---------------------------------------------------------------------------
# Lanczos coefficients (g = 7, n = 9)
# These are the standard Paul Godfrey coefficients used in Boost / SciPy.
# ---------------------------------------------------------------------------
_LANCZOS_G = 7.0
_LANCZOS_COEFFS: list[float] = [
    0.99999999999980993,
    676.5203681218851,
    -1259.1392167224028,
    771.32342877765313,
    -176.61502916214059,
    12.507343278686905,
    -0.13857109526572012,
    9.9843695780195716e-6,
    1.5056327351493116e-7,
]

_LOG_SQRT_2PI = 0.5 * math.log(2.0 * math.pi)


# ===================================================================
# lgamma -- log |Gamma(x)|
# ===================================================================

def lgamma(x: mx.array) -> mx.array:
    """Log of the absolute value of the Gamma function, evaluated on GPU.

    Uses the Lanczos approximation for x >= 0.5 and the reflection formula
    ``Gamma(x) Gamma(1-x) = pi / sin(pi x)`` for x < 0.5.

    Parameters
    ----------
    x : mlx.core.array
        Input tensor.  Non-positive integers yield ``inf``.

    Returns
    -------
    mlx.core.array
        ``ln |Gamma(x)|`` element-wise, same shape and dtype as *x*.
    """
    x = mx.array(x, dtype=mx.float32) if not isinstance(x, mx.array) else x
    # Promote to at least float32
    if x.dtype in (mx.int32, mx.int64, mx.int16, mx.int8, mx.uint8, mx.uint16, mx.uint32, mx.uint64):
        x = x.astype(mx.float32)

    return _lgamma_impl(x)


def _lgamma_positive(x: mx.array) -> mx.array:
    """Lanczos lgamma for x >= 0.5 (no reflection needed)."""
    z = x - 1.0
    # Build the Lanczos sum:  c[0] + c[1]/(z+1) + c[2]/(z+2) + ...
    ag = mx.full(z.shape, _LANCZOS_COEFFS[0], dtype=z.dtype)
    for i in range(1, len(_LANCZOS_COEFFS)):
        ag = ag + _LANCZOS_COEFFS[i] / (z + float(i))
    t = z + _LANCZOS_G + 0.5
    return _LOG_SQRT_2PI + (z + 0.5) * mx.log(t) - t + mx.log(ag)


def _lgamma_impl(x: mx.array) -> mx.array:
    """Full lgamma via Lanczos + reflection, vectorised."""
    # Reflection:  lgamma(x) = log(pi) - log|sin(pi x)| - lgamma(1-x)
    pos_result = _lgamma_positive(x)
    neg_result = (
        math.log(math.pi)
        - mx.log(mx.abs(mx.sin(math.pi * x)) + 1e-45)
        - _lgamma_positive(1.0 - x)
    )
    result = mx.where(x >= 0.5, pos_result, neg_result)

    # Non-positive integers -> inf
    is_nonpos_int = (x <= 0.0) & (mx.abs(x - mx.round(x)) < 1e-7)
    result = mx.where(is_nonpos_int, mx.full(result.shape, float("inf"), dtype=result.dtype), result)
    return result


# ===================================================================
# gamma -- Gamma(x) = sgn * exp(lgamma(x))
# ===================================================================

def gamma(x: mx.array) -> mx.array:
    """Gamma function, evaluated on GPU.

    Computed as ``sign(Gamma(x)) * exp(lgamma(x))``.  For negative
    non-integer *x* the sign alternates between the poles at 0, -1, -2, ...

    Parameters
    ----------
    x : mlx.core.array
        Input tensor.  Non-positive integers yield ``inf``.

    Returns
    -------
    mlx.core.array
        ``Gamma(x)`` element-wise.
    """
    x = mx.array(x, dtype=mx.float32) if not isinstance(x, mx.array) else x
    if x.dtype in (mx.int32, mx.int64, mx.int16, mx.int8, mx.uint8, mx.uint16, mx.uint32, mx.uint64):
        x = x.astype(mx.float32)

    lg = lgamma(x)

    # Sign for negative x: Gamma(x) is negative when floor(x) is odd.
    sign_pos = mx.ones(x.shape, dtype=x.dtype)
    floor_x = mx.floor(x)
    # (-1)^floor(x) for x < 0
    sign_neg = mx.where((floor_x % 2.0) == 0.0, mx.ones_like(x), -mx.ones_like(x))
    sign = mx.where(x >= 0.0, sign_pos, sign_neg)

    return sign * mx.exp(lg)


# ===================================================================
# digamma -- psi(x) = d/dx ln Gamma(x)
# ===================================================================

# Bernoulli numbers B_{2k} / (2k) for k=1..6  (used in asymptotic series)
_DIGAMMA_ASYMP_COEFFS: list[float] = [
    -1.0 / 12.0,
    1.0 / 120.0,
    -1.0 / 252.0,
    1.0 / 240.0,
    -5.0 / 660.0,
    691.0 / 32760.0,
    -1.0 / 12.0,  # not used, padding
]


def digamma(x: mx.array) -> mx.array:
    """Digamma (psi) function, evaluated on GPU.

    For x >= 7, uses the asymptotic expansion::

        psi(x) ~ ln(x) - 1/(2x) - sum_{k=1}^{6} B_{2k} / (2k x^{2k})

    For 0 < x < 7, applies the recurrence ``psi(x) = psi(x+1) - 1/x``
    repeatedly to shift x into the asymptotic regime.

    For x < 0, uses the reflection formula::

        psi(1-x) = psi(x) + pi * cot(pi x)

    Parameters
    ----------
    x : mlx.core.array
        Input tensor.  Non-positive integers yield ``nan``.

    Returns
    -------
    mlx.core.array
        ``psi(x)`` element-wise.
    """
    x = mx.array(x, dtype=mx.float32) if not isinstance(x, mx.array) else x
    if x.dtype in (mx.int32, mx.int64, mx.int16, mx.int8, mx.uint8, mx.uint16, mx.uint32, mx.uint64):
        x = x.astype(mx.float32)

    return _digamma_impl(x)


def _digamma_positive(x: mx.array) -> mx.array:
    """Digamma for x > 0 via recurrence + asymptotic expansion."""
    # Shift x up until x >= 7 using recurrence psi(x) = psi(x+1) - 1/x
    # Maximum 7 iterations is enough since we only need x >= 7.
    result = mx.zeros(x.shape, dtype=x.dtype)
    z = x
    for _ in range(7):
        mask = z < 7.0
        result = mx.where(mask, result - 1.0 / z, result)
        z = mx.where(mask, z + 1.0, z)

    # Asymptotic expansion for z >= 7
    inv_z = 1.0 / z
    inv_z2 = inv_z * inv_z
    # psi(z) ~ ln(z) - 1/(2z) - 1/(12z^2) + 1/(120z^4) - 1/(252z^6) + ...
    asymp = mx.log(z) - 0.5 * inv_z
    zpower = inv_z2  # z^{-2}
    asymp = asymp + (-1.0 / 12.0) * zpower
    zpower = zpower * inv_z2  # z^{-4}
    asymp = asymp + (1.0 / 120.0) * zpower
    zpower = zpower * inv_z2  # z^{-6}
    asymp = asymp + (-1.0 / 252.0) * zpower
    zpower = zpower * inv_z2  # z^{-8}
    asymp = asymp + (1.0 / 240.0) * zpower
    zpower = zpower * inv_z2  # z^{-10}
    asymp = asymp + (-5.0 / 660.0) * zpower
    zpower = zpower * inv_z2  # z^{-12}
    asymp = asymp + (691.0 / 32760.0) * zpower

    result = result + asymp
    return result


def _digamma_impl(x: mx.array) -> mx.array:
    """Full digamma via reflection + recurrence + asymptotic."""
    # Reflection: psi(x) = psi(1-x) - pi * cot(pi * x)  for x < 0
    pos_result = _digamma_positive(mx.maximum(x, mx.array(1e-20)))

    x_refl = 1.0 - x
    refl_psi = _digamma_positive(mx.maximum(x_refl, mx.array(1e-20)))
    cot_val = mx.cos(math.pi * x) / (mx.sin(math.pi * x) + 1e-45)
    neg_result = refl_psi - math.pi * cot_val

    result = mx.where(x > 0.0, pos_result, neg_result)

    # Non-positive integers -> nan
    is_nonpos_int = (x <= 0.0) & (mx.abs(x - mx.round(x)) < 1e-7)
    result = mx.where(is_nonpos_int, mx.full(result.shape, float("nan"), dtype=result.dtype), result)
    return result


# ===================================================================
# beta -- B(a, b) = Gamma(a) Gamma(b) / Gamma(a+b)
# ===================================================================

def beta(a: mx.array, b: mx.array) -> mx.array:
    """Beta function B(a,b) = Gamma(a)*Gamma(b) / Gamma(a+b), on GPU.

    Computed in log-space for numerical stability.

    Parameters
    ----------
    a, b : mlx.core.array
        Input tensors (must be broadcastable). Typically a, b > 0.

    Returns
    -------
    mlx.core.array
        ``B(a, b)`` element-wise.
    """
    a = mx.array(a, dtype=mx.float32) if not isinstance(a, mx.array) else a
    b = mx.array(b, dtype=mx.float32) if not isinstance(b, mx.array) else b
    if a.dtype in (mx.int32, mx.int64):
        a = a.astype(mx.float32)
    if b.dtype in (mx.int32, mx.int64):
        b = b.astype(mx.float32)

    return mx.exp(lgamma(a) + lgamma(b) - lgamma(a + b))
