import numpy as np
import pytest
from scipy.stats import chi2

from gri_utils import constants
from gri_utils.ellipsoids import scale_chi2


def test_err():
    with pytest.raises(ValueError):
        assert scale_chi2([[1], [1]]) == 0
    with pytest.raises(ValueError):
        assert scale_chi2([[1], [1]], dof=1) == 0


def test_scalar():
    assert scale_chi2(1) == pytest.approx(constants.SIG_TO_95_1D**0.5)
    assert scale_chi2(1.0) == pytest.approx(constants.SIG_TO_95_1D**0.5)


def test_vector():
    vec = np.array([1, 2, 3])
    assert np.allclose(scale_chi2([1, 2, 3]), vec * constants.SIG_TO_95_2D**0.5)


def test_2d():
    vec = np.array([[1, 2], [3, 4]])
    assert np.allclose(scale_chi2([[1, 2], [3, 4]]), vec * constants.SIG_TO_95_2D)


def test_3d():
    vec = np.array([[1, 2, 3], [3, 4, 5], [5, 6, 7]])
    assert np.allclose(
        scale_chi2([[1, 2, 3], [3, 4, 5], [5, 6, 7]]),
        vec * constants.SIG_TO_95_3D,
    )


def test_overload():
    assert scale_chi2(1, dof=2) == pytest.approx(constants.SIG_TO_95_2D**0.5)
    assert scale_chi2(1, dof=2, power=1) == pytest.approx(constants.SIG_TO_95_2D)
    assert scale_chi2(1, conf=0.5, dof=2, power=1) == pytest.approx(chi2.ppf(0.5, 2))
