"""
Tests for mlx_gamma against scipy.special reference values.

Run:  pytest tests/test_gamma.py -v
"""

import math

import mlx.core as mx
import numpy as np
import pytest
from scipy import special as sp

from mlx_gamma import beta, digamma, gamma, lgamma


# -------------------------------------------------------------------
# Helpers
# -------------------------------------------------------------------

def _to_np(arr: mx.array) -> np.ndarray:
    return np.array(arr.tolist())


# -------------------------------------------------------------------
# lgamma
# -------------------------------------------------------------------

class TestLgamma:
    """Test lgamma against scipy.special.gammaln / loggamma."""

    def test_positive_integers(self):
        """lgamma(n) = log((n-1)!) for positive integers."""
        x = mx.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 10.0, 20.0])
        result = _to_np(lgamma(x))
        expected = sp.gammaln(np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 10.0, 20.0]))
        np.testing.assert_allclose(result, expected, rtol=1e-5, atol=1e-6)

    def test_positive_reals(self):
        """lgamma on non-integer positive reals."""
        x_np = np.array([0.5, 1.5, 2.5, 3.7, 10.3, 50.1, 100.0])
        x = mx.array(x_np.tolist())
        result = _to_np(lgamma(x))
        expected = sp.gammaln(x_np)
        np.testing.assert_allclose(result, expected, rtol=1e-5)

    def test_large_values(self):
        """lgamma for large arguments (Stirling regime)."""
        x_np = np.array([100.0, 500.0, 1000.0])
        x = mx.array(x_np.tolist())
        result = _to_np(lgamma(x))
        expected = sp.gammaln(x_np)
        np.testing.assert_allclose(result, expected, rtol=1e-4)

    def test_negative_non_integers(self):
        """lgamma for negative non-integers uses reflection."""
        x_np = np.array([-0.5, -1.5, -2.5, -3.7])
        x = mx.array(x_np.tolist())
        result = _to_np(lgamma(x))
        expected = sp.gammaln(x_np)
        # Reflection introduces slightly more error
        np.testing.assert_allclose(result, expected, rtol=1e-4, atol=1e-4)

    def test_nonpositive_integers_are_inf(self):
        """lgamma(0), lgamma(-1), lgamma(-2) should be inf."""
        x = mx.array([0.0, -1.0, -2.0, -3.0])
        result = _to_np(lgamma(x))
        assert all(math.isinf(v) for v in result), f"Expected inf, got {result}"

    def test_half_integer(self):
        """lgamma(0.5) = log(sqrt(pi))."""
        result = lgamma(mx.array([0.5])).tolist()[0]
        expected = 0.5 * math.log(math.pi)
        assert abs(result - expected) < 1e-5

    def test_vectorized_large(self):
        """lgamma on a large array (10k elements) matches scipy."""
        x_np = np.linspace(0.01, 50.0, 10000)
        x = mx.array(x_np.tolist())
        result = _to_np(lgamma(x))
        expected = sp.gammaln(x_np)
        # atol needed because some lgamma values are near zero (e.g. near x=1,2)
        # where relative error is meaningless but absolute error < 1e-5
        np.testing.assert_allclose(result, expected, rtol=1e-4, atol=1e-5)


# -------------------------------------------------------------------
# gamma
# -------------------------------------------------------------------

class TestGamma:
    """Test gamma function."""

    def test_positive_integers(self):
        """Gamma(n) = (n-1)! for positive integers."""
        x = mx.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
        result = _to_np(gamma(x))
        expected = [1.0, 1.0, 2.0, 6.0, 24.0, 120.0]
        np.testing.assert_allclose(result, expected, rtol=1e-5)

    def test_half(self):
        """Gamma(0.5) = sqrt(pi)."""
        result = gamma(mx.array([0.5])).tolist()[0]
        expected = math.sqrt(math.pi)
        assert abs(result - expected) / expected < 1e-5

    def test_negative_noninteger(self):
        """Gamma for negative non-integers."""
        x_np = np.array([-0.5, -1.5, -2.5])
        x = mx.array(x_np.tolist())
        result = _to_np(gamma(x))
        expected = sp.gamma(x_np)
        np.testing.assert_allclose(result, expected, rtol=1e-3, atol=1e-5)


# -------------------------------------------------------------------
# digamma
# -------------------------------------------------------------------

class TestDigamma:
    """Test digamma (psi) function."""

    def test_positive_integers(self):
        """psi(n) = -gamma + sum_{k=1}^{n-1} 1/k."""
        x_np = np.array([1.0, 2.0, 3.0, 5.0, 10.0])
        x = mx.array(x_np.tolist())
        result = _to_np(digamma(x))
        expected = sp.digamma(x_np)
        np.testing.assert_allclose(result, expected, rtol=1e-4, atol=1e-5)

    def test_positive_reals(self):
        """digamma on positive reals."""
        x_np = np.array([0.5, 1.5, 3.7, 10.3, 50.0])
        x = mx.array(x_np.tolist())
        result = _to_np(digamma(x))
        expected = sp.digamma(x_np)
        np.testing.assert_allclose(result, expected, rtol=1e-4, atol=1e-5)

    def test_psi_1(self):
        """psi(1) = -euler_gamma."""
        result = digamma(mx.array([1.0])).tolist()[0]
        expected = -0.5772156649015329  # Euler-Mascheroni
        assert abs(result - expected) < 1e-4

    def test_negative_noninteger(self):
        """digamma for negative non-integers (reflection)."""
        x_np = np.array([-0.5, -1.5, -2.5])
        x = mx.array(x_np.tolist())
        result = _to_np(digamma(x))
        expected = sp.digamma(x_np)
        np.testing.assert_allclose(result, expected, rtol=1e-3, atol=1e-3)


# -------------------------------------------------------------------
# beta
# -------------------------------------------------------------------

class TestBeta:
    """Test Beta function."""

    def test_basic(self):
        """Beta(a, b) = Gamma(a)Gamma(b)/Gamma(a+b)."""
        a = mx.array([1.0, 2.0, 3.0, 0.5])
        b = mx.array([1.0, 3.0, 4.0, 0.5])
        result = _to_np(beta(a, b))
        expected = sp.beta(
            np.array([1.0, 2.0, 3.0, 0.5]),
            np.array([1.0, 3.0, 4.0, 0.5]),
        )
        np.testing.assert_allclose(result, expected, rtol=1e-4)

    def test_symmetry(self):
        """B(a, b) = B(b, a)."""
        a = mx.array([2.0, 5.0, 0.3])
        b = mx.array([3.0, 1.0, 7.0])
        r1 = _to_np(beta(a, b))
        r2 = _to_np(beta(b, a))
        np.testing.assert_allclose(r1, r2, rtol=1e-6)

    def test_beta_half_half(self):
        """B(0.5, 0.5) = pi."""
        result = beta(mx.array([0.5]), mx.array([0.5])).tolist()[0]
        assert abs(result - math.pi) / math.pi < 1e-4


# -------------------------------------------------------------------
# Edge cases and dtypes
# -------------------------------------------------------------------

class TestEdgeCases:
    """Edge-case and dtype tests."""

    def test_integer_input(self):
        """Integer array should be promoted to float."""
        x = mx.array([1, 2, 3, 4, 5])
        result = _to_np(lgamma(x))
        expected = sp.gammaln(np.array([1.0, 2.0, 3.0, 4.0, 5.0]))
        np.testing.assert_allclose(result, expected, rtol=1e-5, atol=1e-6)

    def test_scalar(self):
        """Scalar input should work."""
        result = lgamma(mx.array(5.0)).tolist()
        expected = sp.gammaln(5.0)
        assert abs(result - expected) < 1e-5

    def test_single_element(self):
        """Single-element array."""
        result = gamma(mx.array([1.0])).tolist()[0]
        assert abs(result - 1.0) < 1e-6
