from typing import Optional, Protocol, Self, TypeVar, runtime_checkable

from ..rng import RNG

@runtime_checkable
class Continuous(Protocol):
    def pdf(self, x: float) -> float: ...
    def ln_pdf(self, x: float) -> float: ...

@runtime_checkable
class ContinuousCDF(Continuous, Protocol):
    def cdf(self, x: float) -> float: ...
    def sf(self, x: float) -> float: ...
    def inverse_cdf(self, p: float) -> float: ...
    @property
    def lower(self) -> float: ...
    @property
    def upper(self) -> float: ...

@runtime_checkable
class Discrete(Protocol):
    def pmf(self, x: int) -> float: ...
    def ln_pmf(self, x: int) -> float: ...

@runtime_checkable
class DiscreteCDF(Discrete, Protocol):
    def cdf(self, x: int) -> float: ...
    def sf(self, x: int) -> float: ...
    def inverse_cdf(self, p: float) -> int: ...
    @property
    def lower(self) -> int: ...
    @property
    def upper(self) -> int: ...

@runtime_checkable
class Distribution(Protocol):
    def mean(self) -> Optional[float]: ...
    def median(self) -> Optional[float]: ...
    def variance(self) -> Optional[float]: ...
    def std_dev(self) -> Optional[float]: ...
    def entropy(self) -> Optional[float]: ...
    def skewness(self) -> Optional[float]: ...

@runtime_checkable
class Sample[T](Protocol):
    def sample(self, rng: RNG) -> T: ...

# Concrete distributions

class Beta(ContinuousCDF, Continuous, Distribution, Sample):
    def __init__(self, shape_a: float, shape_b: float): ...
    @property
    def shape_a(self) -> float: ...
    @property
    def shape_b(self) -> float: ...

class BetaError:
    InvalidAlpha: Self
    InvalidBeta: Self

class Exp(ContinuousCDF, Continuous, Distribution, Sample[float]):
    def __init__(self, rate: float): ...
    @property
    def rate(self) -> float: ...

class ExpError:
    RateInvalid: Self

class Gamma(ContinuousCDF, Continuous, Distribution, Sample[float]):
    def __init__(self, shape: float, rate: float): ...
    @property
    def shape(self) -> float: ...
    @property
    def rate(self) -> float: ...

class GammaError:
    ShapeInvalid: Self
    RateInvalid: Self
    ShapeAndRateInfinite: Self

class InverseGamma(ContinuousCDF, Continuous, Distribution, Sample[float]):
    def __init__(self, shape: float, rate: float): ...
    @property
    def shape(self) -> float: ...
    @property
    def rate(self) -> float: ...

class InverseGammaError:
    ShapeInvalid: Self
    RateInvalid: Self

class Laplace(ContinuousCDF, Continuous, Distribution, Sample[float]):
    def __init__(self, location: float, scale: float): ...
    @property
    def location(self) -> float: ...
    @property
    def scale(self) -> float: ...

class LaplaceError:
    LocationInvalid: Self
    ScaleInvalid: Self

class LogNormal(ContinuousCDF, Continuous, Distribution, Sample[float]):
    def __init__(self, location: float, scale: float): ...
    @property
    def location(self) -> float: ...
    @property
    def scale(self) -> float: ...

class LogNormalError:
    LocationInvalid: Self
    ScaleInvalid: Self

class Normal(ContinuousCDF, Continuous, Distribution, Sample[float]):
    def __init__(self, mean: float, std_dev: float): ...

class NormalError:
    MeanInvalid: Self
    StandardDeviationInvalid: Self

class Poisson(DiscreteCDF, Discrete, Distribution, Sample[int]):
    def __init__(self, lambda_: float): ...
    @property
    def lambda_(self) -> float: ...

class PoissonError:
    LambdaInvalid: Self

class Uniform(ContinuousCDF, Continuous, Distribution, Sample[float]):
    def __init__(self, min: float, max: float): ...
    @property
    def min(self) -> float: ...
    @property
    def max(self) -> float: ...

class UniformError:
    MinInvalid: Self
    MaxInvalid: Self
    MaxNotGreaterThanMin: Self
