from typing import TYPE_CHECKING, overload

from riskit.types import BatchCostFunction, Backend, RiskMetric, Risk, Costs, ArrayLike
from riskit.expected import ExpectedValue
from riskit.variance import MeanVariance
from riskit.sampler import MonteCarlo
from riskit.backend import infer

if TYPE_CHECKING:
    from riskit.backend.numpy import (
        BatchCostFunction as NumpyBatchCostFunction,
        Risk as NumpyRisk,
    )
    from riskit.backend.jax import (
        BatchCostFunction as JaxBatchCostFunction,
        Risk as JaxRisk,
    )


class risk:
    @overload
    @staticmethod
    def expected_value_of[TrajectoriesT, UncertaintySamplesT](
        function: "NumpyBatchCostFunction[TrajectoriesT, UncertaintySamplesT]",
    ) -> RiskMetric[TrajectoriesT, UncertaintySamplesT, "NumpyRisk"]:
        """Creates an Expected Value risk metric using the NumPy backend."""
        ...

    @overload
    @staticmethod
    def expected_value_of[TrajectoriesT, UncertaintySamplesT](
        function: "JaxBatchCostFunction[TrajectoriesT, UncertaintySamplesT]",
    ) -> RiskMetric[TrajectoriesT, UncertaintySamplesT, "JaxRisk"]:
        """Creates an Expected Value risk metric using the JAX backend."""
        ...

    @overload
    @staticmethod
    def expected_value_of[
        TrajectoriesT,
        UncertaintySamplesT,
        CostsT: Costs,
        RiskT: Risk,
        ArrayT: ArrayLike,
    ](
        function: BatchCostFunction[TrajectoriesT, UncertaintySamplesT, CostsT],
        backend: Backend[CostsT, RiskT, ArrayT],
    ) -> RiskMetric[TrajectoriesT, UncertaintySamplesT, RiskT]:
        """Creates an Expected Value risk metric using the specified backend."""
        ...

    @staticmethod
    def expected_value_of[
        TrajectoriesT,
        UncertaintySamplesT,
        CostsT: Costs,
        RiskT: Risk,
        ArrayT: ArrayLike,
    ](
        function: BatchCostFunction[TrajectoriesT, UncertaintySamplesT, CostsT],
        backend: Backend[CostsT, RiskT, ArrayT] | None = None,
    ) -> RiskMetric[TrajectoriesT, UncertaintySamplesT, RiskT]:
        return ExpectedValue(
            cost=function,
            backend=backend
            or infer.backend_from(function, type=Backend[CostsT, RiskT, ArrayT]),
            sampler=MonteCarlo(),
        )

    @overload
    @staticmethod
    def mean_variance_of[TrajectoriesT, UncertaintySamplesT](
        function: "NumpyBatchCostFunction[TrajectoriesT, UncertaintySamplesT]",
        *,
        gamma: float,
    ) -> RiskMetric[TrajectoriesT, UncertaintySamplesT, "NumpyRisk"]:
        """Creates a Mean-Variance risk metric using the NumPy backend."""
        ...

    @overload
    @staticmethod
    def mean_variance_of[
        TrajectoriesT,
        UncertaintySamplesT,
        CostsT: Costs,
        RiskT: Risk,
        ArrayT: ArrayLike,
    ](
        function: BatchCostFunction[TrajectoriesT, UncertaintySamplesT, CostsT],
        *,
        backend: Backend[CostsT, RiskT, ArrayT],
        gamma: float,
    ) -> RiskMetric[TrajectoriesT, UncertaintySamplesT, RiskT]:
        """Creates a Mean-Variance risk metric using the specified backend."""
        ...

    @staticmethod
    def mean_variance_of[
        TrajectoriesT,
        UncertaintySamplesT,
        CostsT: Costs,
        RiskT: Risk,
        ArrayT: ArrayLike,
    ](
        function: BatchCostFunction[TrajectoriesT, UncertaintySamplesT, CostsT],
        *,
        backend: Backend[CostsT, RiskT, ArrayT] | None = None,
        gamma: float,
    ) -> RiskMetric[TrajectoriesT, UncertaintySamplesT, RiskT]:
        return MeanVariance(
            cost=function,
            backend=backend
            or infer.backend_from(function, type=Backend[CostsT, RiskT, ArrayT]),
            sampler=MonteCarlo(),
            gamma=gamma,
        )
