import pytest

from moscot.costs._utils import (

    _get_available_backends_and_costs,

    get_available_costs,

    get_cost,

)

class TestCostUtils:

    ALL_BACKENDS_N_COSTS = {

        "moscot": ("barcode_distance", "leaf_distance"),

        "ott": (

            "euclidean",

            "sq_euclidean",

            "cosine",

            "pnorm_p",

            "sq_pnorm",

        ),

    }

    @staticmethod

    def test_get_available_backends_n_costs():

        assert dict(_get_available_backends_and_costs()) == {

            k: list(v) for k, v in _get_available_backends_and_costs().items()

        }

    @staticmethod

    def test_get_available_costs():

        assert get_available_costs() == TestCostUtils.ALL_BACKENDS_N_COSTS

        assert get_available_costs("moscot") == {"moscot": (TestCostUtils.ALL_BACKENDS_N_COSTS["moscot"])}

        assert get_available_costs("ott") == {"ott": TestCostUtils.ALL_BACKENDS_N_COSTS["ott"]}

        with pytest.raises(KeyError):

            get_available_costs("foo")

    @staticmethod

    def test_get_cost_fails():

        invalid_cost = "foo"

        invalid_backend = "bar"

        with pytest.raises(

            ValueError, match=f"Cost `{invalid_cost!r}` is not available for backend `{invalid_backend!r}`."

        ):

            get_cost(invalid_cost, backend=invalid_backend)

        for backend in TestCostUtils.ALL_BACKENDS_N_COSTS:

            with pytest.raises(

                ValueError, match=f"Cost `{invalid_cost!r}` is not available for backend `{backend!r}`."

            ):

                get_cost(invalid_cost, backend=backend)