"""Distribution helpers used by the forecasting model."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable

import numpy as np


@dataclass(frozen=True)
class QuantileLevels:
    """Quantile levels for (optimistic, most-likely, pessimistic) points."""

    p_optimistic: float = 0.10
    p_most_likely: float = 0.50
    p_pessimistic: float = 0.90

    def __post_init__(self) -> None:
        levels = (self.p_optimistic, self.p_most_likely, self.p_pessimistic)
        if not (0.0 < levels[0] < levels[1] < levels[2] < 1.0):
            raise ValueError("Quantile levels must satisfy 0 < p_a < p_m < p_b < 1")

    @property
    def as_tuple(self) -> tuple[float, float, float]:
        return (self.p_optimistic, self.p_most_likely, self.p_pessimistic)


@dataclass(frozen=True)
class LogNormalParams:
    """Parameters of a lognormal distribution: log X ~ Normal(mu, sigma^2)."""

    mu: float
    sigma: float

    def __post_init__(self) -> None:
        if self.sigma <= 0.0:
            raise ValueError("sigma must be positive")

    def sample(self, rng: np.random.Generator, size: int = 1) -> np.ndarray:
        """Draw samples."""
        if size <= 0:
            raise ValueError("size must be positive")
        return np.exp(rng.normal(loc=self.mu, scale=self.sigma, size=size))


class WeightedLogNormalQuantileFitter:
    """Fit a lognormal distribution to three quantile targets.

    We interpret (a, m, b) as approximate quantiles at (p_a, p_m, p_b).
    On the log scale: log q_p = mu + sigma z_p.

    Fitting is done by weighted least squares for numerical robustness.
    """

    def __init__(
        self,
        levels: QuantileLevels | None = None,
        weights: tuple[float, float, float] = (1.0, 2.0, 1.0),
        *,
        min_sigma: float = 1e-6,
    ):
        self.levels = levels or QuantileLevels()
        if any(w <= 0 for w in weights):
            raise ValueError("All weights must be positive")
        self.weights = weights
        self.min_sigma = min_sigma

    @staticmethod
    def _z_values(levels: QuantileLevels) -> tuple[float, float, float]:
        # Use inverse CDF of standard normal. Avoid scipy dependency.
        from statistics import NormalDist

        nd = NormalDist(0.0, 1.0)
        pa, pm, pb = levels.as_tuple
        return (nd.inv_cdf(pa), nd.inv_cdf(pm), nd.inv_cdf(pb))

    def fit(self, a: float, m: float, b: float) -> LogNormalParams:
        """Return lognormal parameters matching the triad as quantiles."""
        if a <= 0 or m <= 0 or b <= 0:
            raise ValueError("Triad values must be positive")
        if not (a < m < b):
            raise ValueError("Triad must satisfy a < m < b")

        y = np.log(np.array([a, m, b], dtype=float))
        z = np.array(self._z_values(self.levels), dtype=float)
        w = np.array(self.weights, dtype=float)

        w_sum = float(np.sum(w))
        z_bar = float(np.sum(w * z) / w_sum)
        y_bar = float(np.sum(w * y) / w_sum)

        z_cent = z - z_bar
        y_cent = y - y_bar

        var_z = float(np.sum(w * z_cent * z_cent) / w_sum)
        if var_z <= 0.0:
            raise ValueError("Degenerate quantile levels: variance of z is zero")

        cov_zy = float(np.sum(w * z_cent * y_cent) / w_sum)
        sigma = cov_zy / var_z
        if sigma <= self.min_sigma:
            # Inconsistent triad or extreme rounding can yield sigma <= 0.
            # Treat as a modelling/data error rather than silently flipping sign.
            raise ValueError(
                "Inconsistent triad: fitted sigma is non-positive. "
                "Check that (a,m,b) are coherent quantiles."
            )

        mu = y_bar - sigma * z_bar
        return LogNormalParams(mu=float(mu), sigma=float(sigma))


def ensure_positive(values: Iterable[float], *, name: str) -> None:
    """Validate that all values are strictly positive."""
    for v in values:
        if v <= 0:
            raise ValueError(f"{name} must be strictly positive")
