from typing import Protocol, NamedTuple
from types import TracebackType


class ArrayLike(Protocol):
    @property
    def shape(self) -> tuple[int, ...]:
        """Shape of the array like object."""
        ...


type Costs = ArrayLike
"""A costs array of shape (trajectories, uncertainty samples)."""

type Risk = ArrayLike
"""A risk array of shape (trajectories,)."""


class BatchCostFunction[TrajectoriesT, UncertaintySamplesT, CostsT: Costs](Protocol):
    def __call__(
        self, *, trajectories: TrajectoriesT, uncertainties: UncertaintySamplesT
    ) -> CostsT:
        """Describes the cost function that should be used for evaluating risk.

        Returns:
            An array like object of shape (M, N) where M is the number of trajectories
            and N is the number of uncertainty samples.
        """
        ...


class TrajectoriesProvider[TrajectoriesT](Protocol):
    async def get(self) -> TrajectoriesT:
        """Provides the batch of trajectories for which the risk metric should be computed."""
        ...

    @property
    def trajectory_count(self) -> int:
        """Number of trajectories provided."""
        ...


class Uncertainties[UncertaintySamplesT](Protocol):
    async def sample(self, count: int) -> UncertaintySamplesT:
        """Returns samples from the distribution(s) of the uncertain variable(s)."""
        ...


class Computation[CostsT: Costs, RiskT: Risk, ArrayT: ArrayLike](Protocol):
    async def array_of(self, costs: CostsT) -> ArrayT:
        """Converts the given costs to the backend's array type."""
        ...

    async def to_risk(self, array: ArrayT) -> RiskT:
        """Converts the given array to the backend's risk type."""
        ...

    async def mean(self, data: ArrayT, *, axis: int) -> ArrayT:
        """Computes the mean of the given data along the specified axis."""
        ...

    async def var(self, data: ArrayT, *, axis: int) -> ArrayT:
        """Computes the variance of the given data along the specified axis."""
        ...

    async def axpby(
        self, *, a: float = 1.0, x: ArrayT, b: float = 1.0, y: ArrayT
    ) -> ArrayT:
        """Computes the operation a * x + b * y."""
        ...


class Backend[CostsT: Costs, RiskT: Risk, ArrayT: ArrayLike](Protocol):
    async def __aenter__(self) -> Computation[CostsT, RiskT, ArrayT]:
        """Enters the backend context for computations."""
        ...

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: TracebackType | None,
    ) -> None:
        """Exits the backend context."""
        ...


class SamplingResult[UncertaintySamplesT](NamedTuple):
    samples: UncertaintySamplesT
    sample_count: int


class Sampler[UncertaintySamplesT](Protocol):
    async def sample_from(
        self, uncertainties: Uncertainties[UncertaintySamplesT]
    ) -> SamplingResult[UncertaintySamplesT]:
        """Samples from the given uncertainties."""
        ...


class RiskMetric[TrajectoriesT, UncertaintySamplesT, RiskT: Risk](Protocol):
    async def compute(
        self,
        *,
        trajectories: TrajectoriesProvider[TrajectoriesT],
        uncertainties: Uncertainties[UncertaintySamplesT],
    ) -> RiskT:
        """Computes the risk metric for the given trajectories and uncertainties.

        Returns:
            An array like object of shape (trajectories,) representing the risk
            values for each trajectory.
        """
        ...

    def sampled_with(
        self, sampler: Sampler[UncertaintySamplesT]
    ) -> "RiskMetric[TrajectoriesT, UncertaintySamplesT, RiskT]":
        """Returns a new risk metric that uses the given sampler to sample from uncertainties."""
        ...
