"""Simple Bayesian pooling utilities.

This module implements conjugate posteriors used in the article:

- Beta--Binomial for availability probabilities.
- Dirichlet--Multinomial for blocked-state probabilities.

For the bias and focus models we use linear-Gaussian shrinkage in
``sprint_forecast.models``.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np


@dataclass(frozen=True)
class BetaBinomialPosterior:
    """Conjugate posterior for a Binomial likelihood with Beta prior."""

    a: float
    b: float

    def __post_init__(self) -> None:
        if self.a <= 0 or self.b <= 0:
            raise ValueError("a and b must be positive")

    @classmethod
    def from_prior_and_counts(
        cls, a0: float, b0: float, *, successes: int, trials: int
    ) -> "BetaBinomialPosterior":
        if a0 <= 0 or b0 <= 0:
            raise ValueError("a0 and b0 must be positive")
        if trials < 0:
            raise ValueError("trials must be non-negative")
        if not (0 <= successes <= trials):
            raise ValueError("successes must be in [0, trials]")
        return cls(a=a0 + successes, b=b0 + trials - successes)

    @property
    def mean(self) -> float:
        return float(self.a / (self.a + self.b))

    def sample(self, rng: np.random.Generator, size: int = 1) -> np.ndarray:
        if size <= 0:
            raise ValueError("size must be positive")
        return rng.beta(self.a, self.b, size=size)


@dataclass(frozen=True)
class DirichletMultinomialPosterior:
    """Conjugate posterior for a Multinomial likelihood with Dirichlet prior."""

    alpha: np.ndarray

    def __post_init__(self) -> None:
        a = np.asarray(self.alpha, dtype=float).reshape(-1)
        if np.any(a <= 0):
            raise ValueError("All Dirichlet parameters must be positive")
        object.__setattr__(self, "alpha", a)

    @classmethod
    def from_prior_and_counts(
        cls,
        prior_alpha: np.ndarray,
        *,
        counts: np.ndarray,
    ) -> "DirichletMultinomialPosterior":
        prior = np.asarray(prior_alpha, dtype=float).reshape(-1)
        c = np.asarray(counts, dtype=float).reshape(-1)
        if prior.shape != c.shape:
            raise ValueError("prior_alpha and counts must have the same shape")
        if np.any(prior <= 0):
            raise ValueError("prior_alpha must be positive")
        if np.any(c < 0):
            raise ValueError("counts must be non-negative")
        return cls(alpha=prior + c)

    @property
    def mean(self) -> np.ndarray:
        return self.alpha / float(np.sum(self.alpha))

    def sample(self, rng: np.random.Generator, size: int = 1) -> np.ndarray:
        if size <= 0:
            raise ValueError("size must be positive")
        return rng.dirichlet(self.alpha, size=size)
