from typing import cast
from dataclasses import dataclass

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

import jax.numpy as jnp
import jax.random as jrandom
from jax import Array

# TODO: Extract seeds to be configurable.


@dataclass(frozen=True)
class Gaussian(Uncertainty):
    mean: Array
    covariance: Array

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

        assert self.covariance.shape == (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(self, count: int) -> UncertaintySamples:
        V, D = self.mean.shape

        key = jrandom.PRNGKey(0)

        samples_list = []
        for v in range(V):
            key, subkey = jrandom.split(key)
            samples_v = jrandom.multivariate_normal(
                subkey, self.mean[v], self.covariance[v], shape=(count,)
            )
            samples_list.append(samples_v)

        samples = jnp.stack(samples_list, axis=0).transpose(0, 2, 1)

        assert samples.shape == (V, D, count), (
            f"Samples shape {samples.shape} does not match expected shape {(V, D, count)}."
        )

        return samples

    @property
    def shape(self) -> tuple[int, int]:
        return cast(tuple[int, int], self.mean.shape)


@dataclass(frozen=True)
class Uniform(Uncertainty):
    lower: Array
    upper: Array

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

        assert self.upper.shape == 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 jnp.all(self.upper >= self.lower), (
            "All upper bounds must be greater than lower bounds."
        )

    async def sample(self, count: int) -> UncertaintySamples:
        key = jrandom.PRNGKey(0)

        samples = self._roll_axes_left(
            jrandom.uniform(
                key,
                shape=(count, *(shape := self.lower.shape)),
                minval=self.lower,
                maxval=self.upper,
            )
        )

        assert samples.shape == (*shape, count), (
            f"Samples shape {samples.shape} does not match expected shape {(*shape, count)}."
        )

        return samples

    @property
    def shape(self) -> tuple[int, ...]:
        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(*, mean: Array, covariance: Array) -> Gaussian:
        return Gaussian(mean=mean, covariance=covariance)

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