"""Tests for adapters for probability distribution objects, scipy facing."""

import numpy as np
import pandas as pd
import pytest
from skbase.testing import QuickTester

from skpro.distributions.adapters.scipy import _ScipyAdapter
from skpro.tests.test_all_estimators import BaseFixtureGenerator, PackageConfig

__author__ = ["fkiraly", "malikrafsan"]


def test_empirical_from_discrete():
    """Test empirical_from_discrete."""
    from scipy.stats import rv_discrete

    from skpro.distributions.adapters.scipy._empirical import empirical_from_discrete

    xk = np.arange(7)
    pk = (0.1, 0.2, 0.3, 0.1, 0.1, 0.0, 0.2)
    pk2 = (0.1, 0.1, 0.4, 0.0, 0.1, 0.2, 0.1)

    dist1 = rv_discrete(name="custm", values=(xk, pk))
    dist2 = rv_discrete(name="custm", values=(xk, pk2))

    emp = empirical_from_discrete([dist1, dist2])
    assert isinstance(emp.spl, pd.DataFrame)
    assert isinstance(emp.weights, pd.Series)
    assert emp.spl.shape == (14, 1)
    assert emp.weights.shape == (14,)
    expected_idx = pd.MultiIndex.from_arrays(
        [[0, 1, 2, 3, 4, 5, 6] * 2, [0] * 7 + [1] * 7]
    )
    assert np.all(emp.spl.index == expected_idx)
    assert np.all(emp.spl.columns == [0])

    emp2 = empirical_from_discrete(
        [dist1, dist2], index=pd.Index(["foo", "bar"]), columns=["abc"]
    )
    assert isinstance(emp2.spl, pd.DataFrame)
    assert isinstance(emp2.weights, pd.Series)
    assert emp2.spl.shape == (14, 1)
    assert emp2.weights.shape == (14,)
    expected_idx = pd.MultiIndex.from_arrays(
        [[0, 1, 2, 3, 4, 5, 6] * 2, ["foo"] * 7 + ["bar"] * 7]
    )
    assert np.all(emp2.spl.index == expected_idx)
    assert np.all(emp2.spl.columns == ["abc"])


class ScipyDistributionFixtureGenerator(BaseFixtureGenerator):
    """Fixture generator for scipy distributions adapter.

    Fixtures parameterized
    ----------------------
    object_class: object inheriting from BaseObject
        ranges over object classes not excluded by EXCLUDE_OBJECTS, EXCLUDED_TESTS
    object_instance: instance of object inheriting from BaseObject
        ranges over object classes not excluded by EXCLUDE_OBJECTS, EXCLUDED_TESTS
        instances are generated by create_test_instance class method
    """

    object_type_filter = _ScipyAdapter


class TestScipyAdapter(PackageConfig, ScipyDistributionFixtureGenerator, QuickTester):
    """Test the scipy adapter."""

    METHOD_TESTS = {
        "NO_PARAMS": [("mean", "mean"), ("var", "var")],
        "X_PARAMS": [("cdf", "cdf"), ("ppf", "ppf")],
        "CONTINUOUS": [("pdf", "pdf"), ("log_pdf", "logpdf")],
        "DISCRETE": [("pmf", "pmf"), ("log_pmf", "logpmf")],
    }

    X_VALUES = [0.1, 0.5, 0.99]

    @pytest.mark.parametrize("method,scipy_method", METHOD_TESTS["NO_PARAMS"])
    def test_method_no_params(self, object_instance, method, scipy_method):
        """Test method that doesn't need additional parameters."""
        res = getattr(object_instance, method)()
        params = object_instance._get_scipy_param()
        scipy_obj = object_instance._get_scipy_object()

        scipy_res = getattr(scipy_obj, scipy_method)(*params[0], **params[1])

        assert np.allclose(res, scipy_res)

    @pytest.mark.parametrize("method,scipy_method", METHOD_TESTS["X_PARAMS"])
    @pytest.mark.parametrize("x", X_VALUES)
    def test_method_with_x_params(self, object_instance, method, scipy_method, x):
        """Test method that needs x as parameter."""
        res = getattr(object_instance, method)(x)
        params = object_instance._get_scipy_param()
        scipy_obj = object_instance._get_scipy_object()

        scipy_res = getattr(scipy_obj, scipy_method)(x, *params[0], **params[1])

        assert np.allclose(res, scipy_res)

    @pytest.mark.parametrize("method,scipy_method", METHOD_TESTS["CONTINUOUS"])
    @pytest.mark.parametrize("x", X_VALUES)
    def test_method_continuous_dist(self, object_instance, method, scipy_method, x):
        """Test continuous distribution method."""
        res = getattr(object_instance, method)(x)
        if object_instance._tags["distr:measuretype"] != "continuous":
            return None  # in this case, scipy method is not defined
        else:
            params = object_instance._get_scipy_param()
            scipy_obj = object_instance._get_scipy_object()
            scipy_res = getattr(scipy_obj, scipy_method)(x, *params[0], **params[1])

        assert np.allclose(res, scipy_res)

    @pytest.mark.parametrize("method,scipy_method", METHOD_TESTS["DISCRETE"])
    @pytest.mark.parametrize("x", X_VALUES)
    def test_method_discrete_dist(self, object_instance, method, scipy_method, x):
        """Test discrete distribution method."""
        res = getattr(object_instance, method)(x)
        if object_instance._tags["distr:measuretype"] != "discrete":
            return None  # in this case, scipy method is not defined
        else:
            params = object_instance._get_scipy_param()
            scipy_obj = object_instance._get_scipy_object()
            scipy_res = getattr(scipy_obj, scipy_method)(x, *params[0], **params[1])

        assert np.allclose(res, scipy_res)
