from dataclasses import dataclass

from riskit.types import (
    Costs,
    Risk,
    ArrayLike,
    BatchCostFunction,
    TrajectoriesProvider,
    Uncertainties,
    Backend,
    Sampler,
)


@dataclass(frozen=True)
class ExpectedValue[
    TrajectoriesT,
    UncertaintySamplesT,
    CostsT: Costs,
    RiskT: Risk,
    ArrayT: ArrayLike,
]:
    cost: BatchCostFunction[TrajectoriesT, UncertaintySamplesT, CostsT]
    backend: Backend[CostsT, RiskT, ArrayT]
    sampler: Sampler[UncertaintySamplesT]

    async def compute(
        self,
        *,
        trajectories: TrajectoriesProvider[TrajectoriesT],
        uncertainties: Uncertainties[UncertaintySamplesT],
    ) -> RiskT:
        M = trajectories.trajectory_count

        samples, N = await self.sampler.sample_from(uncertainties)
        costs = self.cost(trajectories=await trajectories.get(), uncertainties=samples)

        assert costs.shape == (M, N), (
            f"Costs shape {costs.shape} does not match expected shape {(M, N)}."
        )

        async with self.backend as op:
            return await op.to_risk(await op.mean(await op.array_of(costs), axis=1))

    def sampled_with(
        self, sampler: Sampler[UncertaintySamplesT]
    ) -> "ExpectedValue[TrajectoriesT, UncertaintySamplesT, CostsT, RiskT, ArrayT]":
        return ExpectedValue(cost=self.cost, backend=self.backend, sampler=sampler)
