from typing import cast
from dataclasses import dataclass

from riskit.backend.numpy.backend import Uncertainty, UncertaintySamples

import numpy as np
from numtypes import Array, shape_of, Dims, Shape


# NOTE: We explicitly implement the Uncertainty protocol to help the type checker.
# TODO: Add seeds for reproducibility.


@dataclass(frozen=True)
class Gaussian[V: int, D: int](Uncertainty[Dims[V, D]]):
    mean: Array[Dims[V, D]]
    covariance: Array[Dims[V, D, D]]

    def __post_init__(self) -> None:
        V, D = self.mean.shape

        assert shape_of(self.covariance, matches=(V, D, D)), (
            f"Covariance shape {self.covariance.shape} does not match mean shape {self.mean.shape}. "
            f"Expected covariance shape to be {(V, D, D)}, got {self.covariance.shape}."
        )

    async def sample[N: int](self, count: N) -> UncertaintySamples[Dims[V, D], N]:
        V, D = self.mean.shape

        samples = np.array(
            [
                np.random.multivariate_normal(
                    self.mean[v], self.covariance[v], size=count
                )
                for v in range(V)  # TODO: For loop is not the best.
            ]
        ).transpose(0, 2, 1)

        assert shape_of(samples, matches=(V, D, count), name="samples")

        return samples

    @property
    def shape(self) -> Dims[V, D]:
        return self.mean.shape


@dataclass(frozen=True)
class Uniform[ShapeT: Shape](Uncertainty[ShapeT]):
    lower: Array[ShapeT]
    upper: Array[ShapeT]

    def __post_init__(self) -> None:
        lower_shape = self.lower.shape

        assert shape_of(self.upper, matches=lower_shape), (
            f"Upper bound shape {self.upper.shape} does not match lower bound shape {self.lower.shape}. "
            f"Expected upper bound shape to be {lower_shape}, got {self.upper.shape}."
        )
        assert np.all(self.upper >= self.lower), (
            "All upper bounds must be greater than lower bounds."
        )

    async def sample[N: int](self, count: N) -> UncertaintySamples[ShapeT, N]:
        samples = self._roll_axes_left(
            np.random.uniform(
                low=self.lower,
                high=self.upper,
                size=(count, *(shape := self.lower.shape)),
            )
        )

        assert shape_of(samples, matches=(*shape, count), name="samples")

        return cast(UncertaintySamples[ShapeT, N], samples)

    @property
    def shape(self) -> ShapeT:
        return self.lower.shape

    def _roll_axes_left(self, array: Array) -> Array:
        rest = array.shape[1:]
        return array.transpose(*range(1, len(rest) + 1), 0)


class distribution:
    @staticmethod
    def gaussian[V: int, D: int](
        *, mean: Array[Dims[V, D]], covariance: Array[Dims[V, D, D]]
    ) -> Gaussian[V, D]:
        return Gaussian(mean=mean, covariance=covariance)

    @staticmethod
    def uniform[ShapeT: Shape](
        *, lower: Array[ShapeT], upper: Array[ShapeT]
    ) -> Uniform[ShapeT]:
        return Uniform(lower=lower, upper=upper)
