from typing import Protocol, NamedTuple, Self, Annotated
from types import TracebackType

from numtypes import Array, Dim1, Dim2, Dim3, Dims, Shape, shape_of

import numpy as np


class NumpyBackend:
    @staticmethod
    def create() -> "NumpyBackend":
        return NumpyBackend()

    async def __aenter__(self) -> Self:
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: TracebackType | None,
    ) -> None:
        pass

    async def array_of(self, costs: "Costs") -> "Costs":
        return costs

    async def to_risk(self, array: Array) -> "Risk":
        assert shape_of(array, matches=(-1,))

        return array

    async def mean(self, data: Array, *, axis: int) -> Array:
        return np.mean(data, axis=axis)

    async def var(self, data: Array, *, axis: int) -> Array:
        return np.var(data, axis=axis)

    async def axpby(
        self, *, a: float = 1.0, x: Array, b: float = 1.0, y: Array
    ) -> Array:
        return a * x + b * y


type Inputs = Array[Dim3]
"""Batch of control inputs. The dimensions are (time steps, control dimensions, trajectories)."""

type States = Array[Dim3]
"""Batch of states. The dimensions are (time steps, state dimensions, trajectories)."""

type UncertaintySamples[ShapeT: Shape = Shape, N: int = int] = Array[Dims[*ShapeT, N]]
"""Batch of uncertainty samples. The dimensions are (...uncertainty dimensions, uncertainty samples)."""

type Costs = Annotated[Array[Dim2], NumpyBackend.create]
"""Batch of costs. The dimensions are (trajectories, uncertainty samples)."""

type Risk = Array[Dim1]
"""Batch of risk values. The dimensions are (trajectories,)."""


class Uncertainty[ShapeT: Shape](Protocol):
    async def sample[N: int](self, count: N) -> UncertaintySamples[ShapeT, N]:
        """Returns samples from the distribution of the uncertain variable(s)."""
        ...


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


class InputAndState(NamedTuple):
    u: Inputs
    x: States

    async def get(self) -> Self:
        return self

    @property
    def trajectory_count(self) -> int:
        return self.u.shape[2]
