"""
distributions.py — Probability distribution wrappers for Open FAIR modeling.

Each class wraps a scipy distribution and provides a consistent interface:
  .sample(n)       -> np.ndarray of n samples
  .mean()          -> float
  .std()           -> float
  .percentile(p)   -> float  (p in [0, 100])
  .to_dict()       -> dict   (serializable representation)

Convenience factory functions (pert, normal, lognormal, ...) and a
from_dict() deserializer round out the module.
"""

from __future__ import annotations

import warnings
from abc import ABC, abstractmethod
from typing import Any, Dict, Optional

import numpy as np
from scipy import stats


# ---------------------------------------------------------------------------
# Base class
# ---------------------------------------------------------------------------

class Distribution(ABC):
    """Abstract base class for all Open FAIR probability distributions."""

    @abstractmethod
    def sample(self, n: int) -> np.ndarray:
        """Draw *n* independent samples and return them as a 1-D array."""

    @abstractmethod
    def mean(self) -> float:
        """Return the theoretical mean of the distribution."""

    @abstractmethod
    def std(self) -> float:
        """Return the theoretical standard deviation of the distribution."""

    def percentile(self, p: float) -> float:
        """Return the *p*-th percentile (0 ≤ p ≤ 100)."""
        if not 0 <= p <= 100:
            raise ValueError(f"Percentile p must be in [0, 100]; got {p}")
        return float(self._ppf(p / 100.0))

    @abstractmethod
    def _ppf(self, q: float) -> float:
        """Percent-point function (inverse CDF) for quantile q ∈ [0, 1]."""

    @abstractmethod
    def to_dict(self) -> Dict[str, Any]:
        """Return a JSON-serializable dictionary describing this distribution."""

    def __repr__(self) -> str:  # pragma: no cover
        return f"{self.__class__.__name__}({self.to_dict()})"


# ---------------------------------------------------------------------------
# PERT Distribution
# ---------------------------------------------------------------------------

class PERTDistribution(Distribution):
    """Modified PERT (Program Evaluation and Review Technique) distribution.

    Internally uses a scaled Beta distribution.  The shape parameter *lam*
    (default 4.0) controls how strongly the mode is weighted relative to the
    endpoints; higher values produce a tighter distribution around the mode.

    Parameters
    ----------
    low:  Minimum (optimistic) value.
    mode: Most-likely value.
    high: Maximum (pessimistic) value.
    lam:  Weighting factor (default 4.0).
    """

    def __init__(self, low: float, mode: float, high: float, lam: float = 4.0) -> None:
        if low > mode or mode > high:
            raise ValueError(
                f"PERT requires low <= mode <= high; got ({low}, {mode}, {high})"
            )
        if lam <= 0:
            raise ValueError(f"lam must be positive; got {lam}")
        self.low = float(low)
        self.mode = float(mode)
        self.high = float(high)
        self.lam = float(lam)
        self._dist = self._build_dist()

    def _build_dist(self) -> stats.rv_continuous:
        low, mode, high, lam = self.low, self.mode, self.high, self.lam
        # Degenerate case
        if np.isclose(low, high):
            return None  # handled in sample/mean/std/ppf

        mean_val = (low + lam * mode + high) / (lam + 2.0)

        denom = (mode - mean_val) * (high - low)
        if np.isclose(denom, 0.0):
            # mode == mean: use symmetric beta with alpha = beta = lam/2 + 1
            alpha = lam / 2.0 + 1.0
            beta_param = alpha
        else:
            alpha = (mean_val - low) * (2.0 * mode - low - high) / denom
            beta_param = alpha * (high - mean_val) / (mean_val - low)

        # Guard against non-positive shape parameters
        alpha = max(alpha, 1e-6)
        beta_param = max(beta_param, 1e-6)

        scale = high - low
        return stats.beta(alpha, beta_param, loc=low, scale=scale)

    # -- Distribution interface --

    def sample(self, n: int) -> np.ndarray:
        if self._dist is None:
            return np.full(n, self.low)
        return self._dist.rvs(size=n)

    def mean(self) -> float:
        if self._dist is None:
            return self.low
        return float(self._dist.mean())

    def std(self) -> float:
        if self._dist is None:
            return 0.0
        return float(self._dist.std())

    def _ppf(self, q: float) -> float:
        if self._dist is None:
            return self.low
        return float(self._dist.ppf(q))

    def to_dict(self) -> Dict[str, Any]:
        return {
            "type": "pert",
            "low": self.low,
            "mode": self.mode,
            "high": self.high,
            "lam": self.lam,
        }


# ---------------------------------------------------------------------------
# Normal Distribution
# ---------------------------------------------------------------------------

class NormalDistribution(Distribution):
    """Gaussian (Normal) distribution.

    Parameters
    ----------
    mean: Mean (mu).
    std:  Standard deviation (sigma).  Must be > 0.
    """

    def __init__(self, mean: float, std: float) -> None:
        if std < 0:
            raise ValueError(f"std must be non-negative; got {std}")
        self._mean = float(mean)
        self._std = float(std)
        self._dist = stats.norm(loc=self._mean, scale=max(self._std, 1e-300))

    def sample(self, n: int) -> np.ndarray:
        return self._dist.rvs(size=n)

    def mean(self) -> float:
        return self._mean

    def std(self) -> float:
        return self._std

    def _ppf(self, q: float) -> float:
        return float(self._dist.ppf(q))

    def to_dict(self) -> Dict[str, Any]:
        return {"type": "normal", "mean": self._mean, "std": self._std}


# ---------------------------------------------------------------------------
# LogNormal Distribution
# ---------------------------------------------------------------------------

class LogNormalDistribution(Distribution):
    """Log-normal distribution parameterised by the *real-space* mean and std.

    Internally converts to log-space parameters mu (meanlog) and sigma
    (sdlog) using the standard moment-matching formulas.

    Parameters
    ----------
    meanlog: Log-space mean (mu).
    sdlog:   Log-space standard deviation (sigma).  Must be > 0.

    Use the class method :meth:`from_moments` to construct from the
    desired real-space mean and standard deviation.
    """

    def __init__(self, meanlog: float, sdlog: float) -> None:
        if sdlog <= 0:
            raise ValueError(f"sdlog must be positive; got {sdlog}")
        self.meanlog = float(meanlog)
        self.sdlog = float(sdlog)
        # scipy lognorm: shape=s=sigma, scale=exp(mu)
        self._dist = stats.lognorm(s=self.sdlog, scale=np.exp(self.meanlog))

    @classmethod
    def from_moments(cls, mean: float, std: float) -> "LogNormalDistribution":
        """Construct a LogNormal from real-space mean and std.

        Parameters
        ----------
        mean: Desired real-space mean (must be > 0).
        std:  Desired real-space standard deviation (must be > 0).
        """
        if mean <= 0:
            raise ValueError(f"real-space mean must be positive; got {mean}")
        if std <= 0:
            raise ValueError(f"real-space std must be positive; got {std}")
        cv2 = (std / mean) ** 2
        sdlog = float(np.sqrt(np.log1p(cv2)))
        meanlog = float(np.log(mean) - 0.5 * sdlog ** 2)
        return cls(meanlog=meanlog, sdlog=sdlog)

    def sample(self, n: int) -> np.ndarray:
        return self._dist.rvs(size=n)

    def mean(self) -> float:
        return float(self._dist.mean())

    def std(self) -> float:
        return float(self._dist.std())

    def _ppf(self, q: float) -> float:
        return float(self._dist.ppf(q))

    def to_dict(self) -> Dict[str, Any]:
        return {"type": "lognormal", "meanlog": self.meanlog, "sdlog": self.sdlog}


# ---------------------------------------------------------------------------
# Uniform Distribution
# ---------------------------------------------------------------------------

class UniformDistribution(Distribution):
    """Continuous uniform distribution over [low, high].

    Parameters
    ----------
    low:  Lower bound.
    high: Upper bound.  Must be >= low.
    """

    def __init__(self, low: float, high: float) -> None:
        if high < low:
            raise ValueError(f"high must be >= low; got low={low}, high={high}")
        self.low = float(low)
        self.high = float(high)
        self._dist = stats.uniform(loc=self.low, scale=max(self.high - self.low, 1e-300))

    def sample(self, n: int) -> np.ndarray:
        return self._dist.rvs(size=n)

    def mean(self) -> float:
        return (self.low + self.high) / 2.0

    def std(self) -> float:
        return (self.high - self.low) / np.sqrt(12.0)

    def _ppf(self, q: float) -> float:
        return float(self._dist.ppf(q))

    def to_dict(self) -> Dict[str, Any]:
        return {"type": "uniform", "low": self.low, "high": self.high}


# ---------------------------------------------------------------------------
# Triangular Distribution
# ---------------------------------------------------------------------------

class TriangularDistribution(Distribution):
    """Triangular distribution.

    Parameters
    ----------
    low:  Lower bound.
    mode: Peak of the triangle.
    high: Upper bound.  Must satisfy low <= mode <= high.
    """

    def __init__(self, low: float, mode: float, high: float) -> None:
        if not (low <= mode <= high):
            raise ValueError(
                f"Triangular requires low <= mode <= high; got ({low}, {mode}, {high})"
            )
        self.low = float(low)
        self.mode = float(mode)
        self.high = float(high)
        scale = max(self.high - self.low, 1e-300)
        c = (self.mode - self.low) / scale
        self._dist = stats.triang(c=c, loc=self.low, scale=scale)

    def sample(self, n: int) -> np.ndarray:
        return self._dist.rvs(size=n)

    def mean(self) -> float:
        return (self.low + self.mode + self.high) / 3.0

    def std(self) -> float:
        return float(self._dist.std())

    def _ppf(self, q: float) -> float:
        return float(self._dist.ppf(q))

    def to_dict(self) -> Dict[str, Any]:
        return {"type": "triangular", "low": self.low, "mode": self.mode, "high": self.high}


# ---------------------------------------------------------------------------
# Constant Distribution
# ---------------------------------------------------------------------------

class ConstantDistribution(Distribution):
    """Degenerate distribution that always returns a single constant value.

    Useful for fixing a factor to a known value during sensitivity analysis
    or when expert estimates are highly confident.

    Parameters
    ----------
    value: The constant to return on every sample.
    """

    def __init__(self, value: float) -> None:
        self.value = float(value)

    def sample(self, n: int) -> np.ndarray:
        return np.full(n, self.value)

    def mean(self) -> float:
        return self.value

    def std(self) -> float:
        return 0.0

    def _ppf(self, q: float) -> float:
        return self.value

    def to_dict(self) -> Dict[str, Any]:
        return {"type": "constant", "value": self.value}


# ---------------------------------------------------------------------------
# Beta Distribution (scaled)
# ---------------------------------------------------------------------------

class BetaDistribution(Distribution):
    """Scaled Beta distribution over the interval [low, high].

    Parameters
    ----------
    alpha: Shape parameter alpha (> 0).
    beta:  Shape parameter beta (> 0).
    low:   Lower bound (default 0).
    high:  Upper bound (default 1).
    """

    def __init__(
        self,
        alpha: float,
        beta: float,
        low: float = 0.0,
        high: float = 1.0,
    ) -> None:
        if alpha <= 0 or beta <= 0:
            raise ValueError(f"alpha and beta must be positive; got ({alpha}, {beta})")
        if high <= low:
            raise ValueError(f"high must be > low; got low={low}, high={high}")
        self.alpha = float(alpha)
        self.beta = float(beta)
        self.low = float(low)
        self.high = float(high)
        scale = self.high - self.low
        self._dist = stats.beta(self.alpha, self.beta, loc=self.low, scale=scale)

    def sample(self, n: int) -> np.ndarray:
        return self._dist.rvs(size=n)

    def mean(self) -> float:
        return float(self._dist.mean())

    def std(self) -> float:
        return float(self._dist.std())

    def _ppf(self, q: float) -> float:
        return float(self._dist.ppf(q))

    def to_dict(self) -> Dict[str, Any]:
        return {
            "type": "beta",
            "alpha": self.alpha,
            "beta": self.beta,
            "low": self.low,
            "high": self.high,
        }


# ---------------------------------------------------------------------------
# Poisson Distribution
# ---------------------------------------------------------------------------

class PoissonDistribution(Distribution):
    """Poisson distribution for modeling discrete event counts.

    Commonly used to model threat event frequency when the expected annual
    count (rate) is known or estimated.

    Parameters
    ----------
    lam: Expected rate (lambda, > 0).
    """

    def __init__(self, lam: float) -> None:
        if lam <= 0:
            raise ValueError(f"lam must be positive; got {lam}")
        self.lam = float(lam)
        self._dist = stats.poisson(mu=self.lam)

    def sample(self, n: int) -> np.ndarray:
        return self._dist.rvs(size=n).astype(float)

    def mean(self) -> float:
        return self.lam

    def std(self) -> float:
        return float(np.sqrt(self.lam))

    def _ppf(self, q: float) -> float:
        return float(self._dist.ppf(q))

    def to_dict(self) -> Dict[str, Any]:
        return {"type": "poisson", "lam": self.lam}


# ---------------------------------------------------------------------------
# Convenience factory functions
# ---------------------------------------------------------------------------

def pert(low: float, mode: float, high: float, lam: float = 4.0) -> PERTDistribution:
    """Return a :class:`PERTDistribution` with the given parameters."""
    return PERTDistribution(low=low, mode=mode, high=high, lam=lam)


def normal(mean: float, std: float) -> NormalDistribution:
    """Return a :class:`NormalDistribution` with the given mean and std."""
    return NormalDistribution(mean=mean, std=std)


def lognormal(meanlog: float, sdlog: float) -> LogNormalDistribution:
    """Return a :class:`LogNormalDistribution` with the given log-space parameters."""
    return LogNormalDistribution(meanlog=meanlog, sdlog=sdlog)


def uniform(low: float, high: float) -> UniformDistribution:
    """Return a :class:`UniformDistribution` over [low, high]."""
    return UniformDistribution(low=low, high=high)


def triangular(low: float, mode: float, high: float) -> TriangularDistribution:
    """Return a :class:`TriangularDistribution` with the given parameters."""
    return TriangularDistribution(low=low, mode=mode, high=high)


def constant(value: float) -> ConstantDistribution:
    """Return a :class:`ConstantDistribution` fixed at *value*."""
    return ConstantDistribution(value=value)


def beta(
    alpha: float, beta_param: float, low: float = 0.0, high: float = 1.0
) -> BetaDistribution:
    """Return a :class:`BetaDistribution` with the given shape parameters."""
    return BetaDistribution(alpha=alpha, beta=beta_param, low=low, high=high)


def poisson(lam: float) -> PoissonDistribution:
    """Return a :class:`PoissonDistribution` with rate *lam*."""
    return PoissonDistribution(lam=lam)


# ---------------------------------------------------------------------------
# Deserializer
# ---------------------------------------------------------------------------

_REGISTRY: Dict[str, type] = {
    "pert": PERTDistribution,
    "normal": NormalDistribution,
    "lognormal": LogNormalDistribution,
    "uniform": UniformDistribution,
    "triangular": TriangularDistribution,
    "constant": ConstantDistribution,
    "beta": BetaDistribution,
    "poisson": PoissonDistribution,
}


def from_dict(d: Dict[str, Any]) -> Distribution:
    """Reconstruct a :class:`Distribution` from a dictionary produced by :meth:`to_dict`.

    Parameters
    ----------
    d: Dictionary with at least a ``"type"`` key matching one of the known
       distribution type strings.

    Raises
    ------
    KeyError: If the type string is not recognised.
    """
    d = dict(d)
    dist_type = d.pop("type").lower()
    if dist_type not in _REGISTRY:
        raise KeyError(f"Unknown distribution type '{dist_type}'. Known: {list(_REGISTRY)}")
    cls = _REGISTRY[dist_type]
    # BetaDistribution uses 'beta' as second param name internally
    if dist_type == "beta" and "beta" in d:
        d["beta_param"] = d.pop("beta")
    return cls(**d)
