dists.py

from abc import ABC, abstractmethod
from scipy.stats import beta, gamma, truncexpon
from scipy.integrate import quad
from scipy.stats import lognorm
import numpy as np
from scipy.stats import norm, truncnorm


class Distribution(ABC):
    @abstractmethod
    def sample(n):
        pass

class uniform1D(Distribution):
    def __init__(self, low=0, high=1):
        self.low = low
        self.high = high
        self.name = f"U({low}, {high})"

    def sample(self, n):
        return np.random.uniform(self.low, self.high, size=n)

    def rvs(self, size):
        return np.random.uniform(self.low, self.high, size=size)

    def mean(self):
        return (self.low + self.high) / 2

    def std(self):
        return (self.high - self.low) / np.sqrt(12)

    def median(self):
        return (self.low + self.high) / 2

    def get_d(self):
        return 1

    def get_med_std(self):
        return (self.high - self.low) / 4

    def get_bounds(self):
        return self.low, self.high

    def pdf(self, x):
        return np.where((x >= self.low) & (x <= self.high), 1 / (self.high - self.low), 0)

    def cdf(self, x):
        return np.where(x < self.low, 0, np.where(x > self.high, 1, (x - self.low) / (self.high - self.low)))

    def ppf(self, q):
        return self.low + q * (self.high - self.low)




class normal1D(Distribution):
    def __init__(self, mu=0, sigma=1):
        self.mu, self.sigma = mu, sigma
        self.name= f"N({mu}, {sigma})"

    def sample(self, n):
        return np.random.normal(self.mu, self.sigma, size=n)

    def rvs(self, size):
        return np.random.normal(self.mu, self.sigma, size=size)

    def mean(self):
        return self.mu

    def std(self):
        return self.sigma

    def median(self):
        return self.mu

    def get_d(self):
        return 1

    def get_med_std(self):
        return self.sigma * np.sqrt(np.pi / 2)

    def get_bounds(self):
        return -np.inf, np.inf


class TruncGauss1D(Distribution):
    def __init__(self, trunc_left, trunc_right, mu=0, sigma=1):
        """trunc_left and trunc_right should be number of std from the center, hence we normalized the requested bounds"""
        self.mu, self.sigma = mu, sigma
        self.a = (trunc_left - mu) / sigma
        self.b = (trunc_right - mu) / sigma
        self.name= f"N({mu}, {sigma}), R=[{trunc_left}, {trunc_right}]"

    def sample(self, n):
        return self.rvs(size=n)

    # 1 D output enforced
    def rvs(self, size):
        return truncnorm.rvs(self.a, self.b, loc=self.mu, scale=self.sigma, size=size)

    def mean(self):
        return truncnorm.stats(self.a, self.b, loc=self.mu, scale=self.sigma)[0]

    def std(self):
        return np.sqrt(
            truncnorm.stats(self.a, self.b, loc=self.mu, scale=self.sigma)[1]
        )

    def median(self):
        return truncnorm.median(self.a, self.b, loc=self.mu, scale=self.sigma)

    def get_d(self):
        return 1

    def get_med_std(self):
        return 1 / (2 * truncnorm.pdf(
            self.median(), self.a, self.b, loc=self.mu, scale=self.sigma))

    def get_bounds(self):
        return self.mu + self.a * self.sigma, self.mu + self.b * self.sigma

    def pdf(self, x):
        """Probability density function of the truncated normal distribution."""
        return truncnorm.pdf(x, self.a, self.b, loc=self.mu, scale=self.sigma)

    def cdf(self, x):
        """Cumulative distribution function of the truncated normal distribution."""
        return truncnorm.cdf(x, self.a, self.b, loc=self.mu, scale=self.sigma)

    def ppf(self, x):
        return truncnorm.ppf(x, self.a, self.b, loc=self.mu, scale=self.sigma)


class Trunc1DExpon(Distribution):
    """
    Truncated Exponential distribution on [0, b] with rate l.

    API matches TruncGauss1D:
      - __init__(b, l=1)
      - sample(n)
      - rvs(size)
      - mean(), std(), median()
      - get_d(), get_med_std(), get_bounds()
    """

    def __init__(self, trunc, l=1):
        if trunc <= 0 or l <= 0:
            raise ValueError("Both trunc and l must be > 0")
        self.b = trunc
        self.l = l
        # SciPy’s truncexpon uses a "shape" c and a scale parameter
        # to produce support [0, c*scale].  We want [0, b],
        # so set c = l*b and scale = 1/l.
        self.c = self.l * self.b
        self.scale = 1.0 / self.l
        self.name = f"Exp({l}), R=[0, {trunc}]"

    def sample(self, n):
        # return truncexpon.rvs(self.c, loc=0, scale=self.scale, size=n)
        return truncexpon.rvs(self.c, loc=0, scale=self.scale, size=n, random_state=np.random.default_rng())


    def pdf(self, x):
        return truncexpon.pdf(x, self.c, loc=0, scale=self.scale)

    def cdf(self, x):
        return truncexpon.cdf(x, self.c, loc=0, scale=self.scale)

    def ppf(self, q):
        return truncexpon.ppf(q, self.c, loc=0, scale=self.scale)

    def rvs(self, size):
        return truncexpon.rvs(self.c, loc=0, scale=self.scale, size=size)

    def mean(self):
        return truncexpon.mean(self.c, loc=0, scale=self.scale)

    def std(self):
        # stats(...)[1] returns the variance
        return np.sqrt(truncexpon.stats(self.c, loc=0, scale=self.scale)[1])

    def median(self):
        return truncexpon.median(self.c, loc=0, scale=self.scale)

    def get_d(self):
        return 1

    def get_med_std(self):
        m = self.median()
        return 1.0 / (2.0 * truncexpon.pdf(m, self.c, loc=0, scale=self.scale))

    def get_bounds(self):
        return 0.0, self.b

########## Noa 11/07: I need to check this, generated by ChatGPT - I've started and it's look ok ##########
class Trunc1DLogNorm:
    """
    Truncated Lognormal distribution on [0, b] with parameters (s, scale).

    API matches TruncGauss1D / Trunc1DExpon:
      - __init__(b, s=1, scale=1)
      - sample(n)
      - rvs(size)
      - mean(), std(), median()
      - get_d(), get_med_std(), get_bounds()
    """

    def __init__(self, trunc, s=1.0, scale=1.0):
        if trunc <= 0 or s <= 0 or scale <= 0:
            raise ValueError("All parameters must be > 0")
        self.b = trunc
        self.s = s
        self.scale = scale
        self.mu = np.log(scale)
        self.Fb = lognorm.cdf(self.b, s=self.s, scale=self.scale)  # normalization factor
        self.name = f"Lognormal(s={s}, scale={scale}), R=[0, {trunc}]"

    def sample(self, n):
        return self.rvs(n)

    def rvs(self, size):
        u = np.random.uniform(0.0, self.Fb, size)
        return lognorm.ppf(u, s=self.s, scale=self.scale)

    def pdf(self, x):
        mask = (x > 0) & (x <= self.b)
        result = np.zeros_like(x, dtype=float)
        result[mask] = lognorm.pdf(x[mask], s=self.s, scale=self.scale) / self.Fb
        return result

    def cdf(self, x):
        x = np.clip(x, 0, self.b)
        return lognorm.cdf(x, s=self.s, scale=self.scale) / self.Fb

    def ppf(self, q):
        q = np.clip(q, 0.0, 1.0)
        return lognorm.ppf(q * self.Fb, s=self.s, scale=self.scale)

    def _raw_moment(self, k):
        sigma = self.s
        mu = self.mu
        z = (np.log(self.b) - mu - k * sigma ** 2) / sigma

        numerator = np.exp(k * mu + 0.5 * k ** 2 * sigma ** 2) * norm.cdf(z)
        return numerator / self.Fb

    def mean(self):
        return self._raw_moment(1)

    def std(self):
        mu = self._raw_moment(1)
        ex2 = self._raw_moment(2)
        return np.sqrt(ex2 - mu ** 2)

    def median(self):
        return self.ppf(0.5)

    def get_d(self):
        return 1

    def get_med_std(self):
        m = self.median()
        return 1.0 / (2.0 * self.pdf(np.array([m]))[0])

    def get_bounds(self):
        return 0.0, self.b




def truncated_normal(lower, upper, mu=0, sigma=2):
    a, b = (lower - mu) / sigma, (upper - mu) / sigma
    return truncnorm(a, b, loc=mu, scale=sigma)


class Beta1D(Distribution):
    def __init__(self, alpha_param, beta_param):
        self.alpha_param = alpha_param
        self.beta_param = beta_param
        self.name = f"Beta(alpha={alpha_param}, beta={beta_param}), R=[0, 1]"

    def sample(self, n):
        return self.rvs(size=n)

    def rvs(self, size):
        return beta.rvs(self.alpha_param, self.beta_param, size=size)

    def mean(self):
        return beta.stats(self.alpha_param, self.beta_param, moments='m')

    def pdf(self, x):
        """Probability density function at x (can be scalar or array)."""
        return beta.pdf(x, self.alpha_param, self.beta_param)

    def std(self):
        return np.sqrt(beta.stats(self.alpha_param, self.beta_param, moments='v'))

    def median(self):
        return beta.median(self.alpha_param, self.beta_param)

    def get_d(self):
        return 1

    def get_med_std(self):
        # Using the approximate formula: 1 / (2 * pdf(median)) -using normal approximation
        med = self.median()
        pdf_med = beta.pdf(med, self.alpha_param, self.beta_param)
        return 1 / (2 * pdf_med) if pdf_med != 0 else np.inf

    def get_bounds(self):
        return 0.0, 1.0

    def ppf(self, q):
        return beta.ppf(q, self.alpha_param, self.beta_param)


########## Tomer 01/07: I need to check this, generated by ChatGPT ##########
class TruncGamma1D(Distribution):
    def __init__(self, trunc_right, shape, scale=1):
        self.trunc_right = trunc_right
        self.shape = shape  # a = shape parameter (k)
        self.scale = scale  # theta = scale parameter

        # Compute CDF at truncation to adjust for sampling within [0, trunc_right]
        self.cdf_max = gamma.cdf(self.trunc_right, self.shape, scale=self.scale)
        self.name = f"Gamma(shape={shape}, scale={scale}), R=[0, {trunc_right}]"

    def sample(self, n):
        return self.rvs(size=n)

    def rvs(self, size):
        # Inverse transform sampling within truncated range
        u = np.random.uniform(0, self.cdf_max, size=size)
        return gamma.ppf(u, self.shape, scale=self.scale)

    def mean(self):
        # Adjusted mean under truncation
        full_mean = gamma.mean(self.shape, scale=self.scale)
        if self.trunc_right == np.inf:
            return full_mean
        else:
            return self._trunc_moment(1)

    def std(self):
        # Adjusted std under truncation
        if self.trunc_right == np.inf:
            return gamma.std(self.shape, scale=self.scale)
        else:
            mu = self.mean()
            second_moment = self._trunc_moment(2)
            return np.sqrt(second_moment - mu**2)

    def median(self):
        if self.trunc_right == np.inf:
            return gamma.median(self.shape, scale=self.scale)
        else:
            return gamma.ppf(0.5 * self.cdf_max, self.shape, scale=self.scale)

    def get_d(self):
        return 1

    def get_med_std(self):
        med = self.median()
        pdf_med = gamma.pdf(med, self.shape, scale=self.scale) / self.cdf_max
        return 1 / (2 * pdf_med) if pdf_med != 0 else np.inf

    def get_bounds(self):
        return 0.0, self.trunc_right

    def _trunc_moment(self, order):
        """
        Compute the truncated raw moment of given order
        E[X^order | X <= trunc_right]
        """
        def integrand(x):
            return x**order * gamma.pdf(x, self.shape, scale=self.scale)

        integral, _ = quad(integrand, 0, self.trunc_right)
        return integral / self.cdf_max




# If you already have your own base class, keep using it.
class Distribution:
    pass

class TruncGMM2_1D(Distribution):
    def __init__(self, mu1, sigma1, mu2, sigma2, weight, trunc_left=-np.inf, trunc_right=np.inf):
        """
        Truncated 1D 2-Gaussian mixture:
          X ~ w*N(mu1, sigma1^2) + (1-w)*N(mu2, sigma2^2), truncated to [L, R].
        """
        assert sigma1 > 0 and sigma2 > 0, "Sigmas must be positive."
        assert 0 < weight < 1, "weight must be in (0,1)."
        assert trunc_left < trunc_right, "Require trunc_left < trunc_right."

        self.mu1, self.sigma1 = float(mu1), float(sigma1)
        self.mu2, self.sigma2 = float(mu2), float(sigma2)
        self.w = float(weight)
        self.L, self.R = float(trunc_left), float(trunc_right)

        # Standardized bounds for each component
        self.a1 = (self.L - self.mu1) / self.sigma1 if np.isfinite(self.L) else -np.inf
        self.b1 = (self.R - self.mu1) / self.sigma1 if np.isfinite(self.R) else  np.inf
        self.a2 = (self.L - self.mu2) / self.sigma2 if np.isfinite(self.L) else -np.inf
        self.b2 = (self.R - self.mu2) / self.sigma2 if np.isfinite(self.R) else  np.inf

        self._Z = self._compute_Z()  # normalizing constant over [L,R]
        if self._Z <= 0:
            raise ValueError("Truncation interval has ~zero mixture mass.")

        # Posterior component weights after truncation
        self._w1_trunc = self.w * self._Phi_diff(self.a1, self.b1) / self._Z
        self._w2_trunc = (1 - self.w) * self._Phi_diff(self.a2, self.b2) / self._Z

        self.name = f"GMM2_N({mu1},{sigma1})_N({mu2},{sigma2}), R=[{self.L}, {self.R}]"

    # -------- helpers (stable) --------
    @staticmethod
    def _phi(z): return norm.pdf(z)
    @staticmethod
    def _Phi(z): return norm.cdf(z)
    @staticmethod
    def _sf(z):  return norm.sf(z)   # 1 - cdf, stable in tails

    @classmethod
    def _Phi_diff(cls, a, b):
        """
        Stable P(a<Z<b) for Z~N(0,1).
        Avoids catastrophic cancellation for large |a|,|b|.
        """
        if np.isneginf(a): return cls._Phi(b)
        if np.isposinf(b): return cls._sf(a)
        if a >= 0:         return cls._sf(a) - cls._sf(b)          # both >= 0
        if b <= 0:         return cls._Phi(b) - cls._Phi(a)        # both <= 0
        return cls._Phi(b) + cls._Phi(-a) - 1.0                    # a<0<b

    def _compute_Z(self):
        # Mixture probability mass inside [L,R]
        z1 = self._Phi_diff(self.a1, self.b1)
        z2 = self._Phi_diff(self.a2, self.b2)
        return self.w * z1 + (1 - self.w) * z2

    # -------- interface --------
    def sample(self, n):
        return self.rvs(size=n)

    @staticmethod
    def _rvs_truncnorm(a, b, loc, scale, size):
        """
        Inverse-CDF sampling of a truncated normal with standardized bounds a,b.
        Works with a=-inf or b=+inf. `size` can be an int or tuple.
        """
        low = norm.cdf(a)
        high = norm.cdf(b)
        # Keep draws strictly inside (0,1) to avoid +/-inf from ppf
        tiny = np.nextafter(0.0, 1.0)
        low = max(low, tiny)
        high = min(high, 1.0 - tiny)
        u = np.random.uniform(low, high, size=size)
        z = norm.ppf(u)
        return loc + scale * z

    def rvs(self, size):
        """
        Return random variates with the requested shape:
          - int -> (n,)
          - tuple -> given shape
          - None -> scalar
        """
        if size is None:
            shape = ()
        elif np.isscalar(size):
            shape = (int(size),)
        else:
            shape = tuple(size)

        N = 1 if shape == () else int(np.prod(shape))
        if N == 0:
            return np.empty(shape)

        # Component assignment under truncated mixture weights
        p1 = float(self._w1_trunc)
        mask = np.random.random(N) < p1
        n1 = int(mask.sum())
        n2 = N - n1

        s1 = self._rvs_truncnorm(self.a1, self.b1, self.mu1, self.sigma1, n1) if n1 > 0 else np.empty(0)
        s2 = self._rvs_truncnorm(self.a2, self.b2, self.mu2, self.sigma2, n2) if n2 > 0 else np.empty(0)

        out = np.empty(N, dtype=float)
        out[mask] = s1
        out[~mask] = s2
        return out.reshape(shape)

    def mean(self):
        m1 = truncnorm.stats(self.a1, self.b1, loc=self.mu1, scale=self.sigma1, moments='m')
        m2 = truncnorm.stats(self.a2, self.b2, loc=self.mu2, scale=self.sigma2, moments='m')
        return self._w1_trunc * m1 + self._w2_trunc * m2

    def std(self):
        m1 = truncnorm.stats(self.a1, self.b1, loc=self.mu1, scale=self.sigma1, moments='m')
        v1 = truncnorm.stats(self.a1, self.b1, loc=self.mu1, scale=self.sigma1, moments='v')
        m2 = truncnorm.stats(self.a2, self.b2, loc=self.mu2, scale=self.sigma2, moments='m')
        v2 = truncnorm.stats(self.a2, self.b2, loc=self.mu2, scale=self.sigma2, moments='v')

        EX = self._w1_trunc * m1 + self._w2_trunc * m2
        EX2 = self._w1_trunc * (v1 + m1**2) + self._w2_trunc * (v2 + m2**2)
        return np.sqrt(max(EX2 - EX**2, 0.0))

    def cdf(self, x):
        x = np.asarray(x)

        def _F_scalar(xi):
            if xi <= self.L: return 0.0
            if xi >= self.R: return 1.0
            t1 = (xi - self.mu1) / self.sigma1
            t2 = (xi - self.mu2) / self.sigma2
            num = self.w * self._Phi_diff(self.a1, t1) + (1 - self.w) * self._Phi_diff(self.a2, t2)
            return num / self._Z

        if np.isscalar(x):
            return _F_scalar(x)
        return np.vectorize(_F_scalar)(x)

    def median(self):
        # Solve F(x)=0.5 via bisection on [L,R] (or wide finite bracket if unbounded)
        L = self.L if np.isfinite(self.L) else min(self.mu1, self.mu2) - 10 * max(self.sigma1, self.sigma2)
        R = self.R if np.isfinite(self.R) else max(self.mu1, self.mu2) + 10 * max(self.sigma1, self.sigma2)

        def F(x):
            if x <= self.L: return 0.0
            if x >= self.R: return 1.0
            t1 = (x - self.mu1) / self.sigma1
            t2 = (x - self.mu2) / self.sigma2
            num = self.w * self._Phi_diff(self.a1, t1) + (1 - self.w) * self._Phi_diff(self.a2, t2)
            return num / self._Z

        lo, hi = L, R
        for _ in range(80):
            mid = 0.5 * (lo + hi)
            fmid = F(mid)
            if abs(fmid - 0.5) < 1e-12:
                return mid
            if fmid < 0.5:
                lo = mid
            else:
                hi = mid
        return 0.5 * (lo + hi)

    def get_d(self):
        return 1

    def get_med_std(self):
        med = self.median()
        pdf_med = self.pdf(med)
        return 1 / (2 * pdf_med) if pdf_med > 0 else np.inf

    def get_bounds(self):
        return self.L, self.R

    def pdf(self, x):
        x = np.asarray(x)
        # Untruncated mixture pdf
        f = self.w * norm.pdf(x, loc=self.mu1, scale=self.sigma1) + (1 - self.w) * norm.pdf(x, loc=self.mu2, scale=self.sigma2)
        # Zero outside [L,R], renormalize inside
        if np.isscalar(x):
            if x < self.L or x > self.R:
                return 0.0
            return f / self._Z
        else:
            return np.where((x >= self.L) & (x <= self.R), f / self._Z, 0.0)


import numpy as np
from scipy.stats import norm, truncnorm


class TruncGMM1D(Distribution):
    """
    Truncated 1D Gaussian mixture on [L, R].

    Example for 3 components:
        dist = TruncGMM1D(
            mus=[mu1, mu2, mu3],
            sigmas=[sigma1, sigma2, sigma3],
            weights=[w1, w2, w3],
            trunc_left=L,
            trunc_right=R
        )

    The untruncated density is:
        sum_k weights[k] * N(mus[k], sigmas[k]^2)

    Then it is truncated and renormalized to [L, R].
    """

    def __init__(self, mus, sigmas, weights, trunc_left=-np.inf, trunc_right=np.inf):
        mus = np.asarray(mus, dtype=float)
        sigmas = np.asarray(sigmas, dtype=float)
        weights = np.asarray(weights, dtype=float)

        if mus.ndim != 1 or sigmas.ndim != 1 or weights.ndim != 1:
            raise ValueError("mus, sigmas, and weights must be 1D arrays/lists.")

        if not (len(mus) == len(sigmas) == len(weights)):
            raise ValueError("mus, sigmas, and weights must have the same length.")

        if len(mus) < 1:
            raise ValueError("Need at least one Gaussian component.")

        if np.any(sigmas <= 0):
            raise ValueError("All sigmas must be positive.")

        if np.any(weights < 0):
            raise ValueError("All weights must be nonnegative.")

        if np.sum(weights) <= 0:
            raise ValueError("At least one weight must be positive.")

        if trunc_left >= trunc_right:
            raise ValueError("Require trunc_left < trunc_right.")

        weights = weights / np.sum(weights)

        self.mus = mus
        self.sigmas = sigmas
        self.weights = weights
        self.K = len(mus)

        self.L = float(trunc_left)
        self.R = float(trunc_right)

        self.a = np.where(np.isfinite(self.L), (self.L - self.mus) / self.sigmas, -np.inf)
        self.b = np.where(np.isfinite(self.R), (self.R - self.mus) / self.sigmas, np.inf)

        self._component_masses = np.array([self._Phi_diff(self.a[k], self.b[k]) for k in range(self.K)])
        self._Z = float(np.sum(self.weights * self._component_masses))

        if self._Z <= 0:
            raise ValueError("Truncation interval has ~zero mixture mass.")

        self._weights_trunc = self.weights * self._component_masses / self._Z

        parts = ", ".join(f"N({self.mus[k]}, {self.sigmas[k]})" for k in range(self.K))
        self.name = f"TruncGMM{self.K}_1D({parts}), R=[{self.L}, {self.R}]"

    # -------- helpers --------

    @staticmethod
    def _Phi(z):
        return norm.cdf(z)

    @staticmethod
    def _sf(z):
        return norm.sf(z)

    @classmethod
    def _Phi_diff(cls, a, b):
        """
        Stable P(a < Z < b), Z ~ N(0,1).
        """
        if np.isneginf(a): return cls._Phi(b)
        if np.isposinf(b):  return cls._sf(a)
        if a >= 0:  return cls._sf(a) - cls._sf(b)
        if b <= 0: return cls._Phi(b) - cls._Phi(a)
        return cls._Phi(b) + cls._Phi(-a) - 1.0

    @staticmethod
    def _rvs_truncnorm(a, b, loc, scale, size):
        low = norm.cdf(a)
        high = norm.cdf(b)

        tiny = np.nextafter(0.0, 1.0)
        low = max(low, tiny)
        high = min(high, 1.0 - tiny)

        u = np.random.uniform(low, high, size=size)
        z = norm.ppf(u)
        return loc + scale * z

    # -------- interface --------

    def sample(self, n):
        return self.rvs(size=n)

    def rvs(self, size=None):
        if size is None: shape = ()
        elif np.isscalar(size):  shape = (int(size),)
        else: shape = tuple(size)

        N = 1 if shape == () else int(np.prod(shape))

        if N == 0: return np.empty(shape)

        comp = np.random.choice(self.K, size=N, p=self._weights_trunc)
        out = np.empty(N, dtype=float)

        for k in range(self.K):
            mask = comp == k
            nk = int(mask.sum())

            if nk > 0:
                out[mask] = self._rvs_truncnorm(self.a[k], self.b[k], self.mus[k],self.sigmas[k],nk)

        out = out.reshape(shape)

        if shape == ():
            return float(out)

        return out

    def pdf(self, x):
        scalar_input = np.ndim(x) == 0
        x = np.asarray(x, dtype=float)
        f = np.zeros_like(x, dtype=float)
        inside = (x >= self.L) & (x <= self.R)

        for k in range(self.K):
            f += self.weights[k] * norm.pdf(x, loc=self.mus[k], scale=self.sigmas[k])

        result = np.where(inside, f / self._Z, 0.0)

        if scalar_input:
            return float(result)

        return result

    def cdf(self, x):
        scalar_input = np.ndim(x) == 0
        x = np.asarray(x, dtype=float)

        result = np.zeros_like(x, dtype=float)

        result[x >= self.R] = 1.0

        inside = (x > self.L) & (x < self.R)

        if np.any(inside):
            xi = x[inside]
            num = np.zeros_like(xi, dtype=float)

            for k in range(self.K):
                t = (xi - self.mus[k]) / self.sigmas[k]
                vals = np.array([self._Phi_diff(self.a[k], tj) for tj in t])
                num += self.weights[k] * vals

            result[inside] = num / self._Z

        if scalar_input:
            return float(result)

        return result

    def mean(self):
        means = np.array([truncnorm.stats(self.a[k],self.b[k],loc=self.mus[k],scale=self.sigmas[k],moments="m") for k in range(self.K)])

        return float(np.sum(self._weights_trunc * means))

    def std(self):
        means = np.zeros(self.K)
        variances = np.zeros(self.K)

        for k in range(self.K):
            means[k] = truncnorm.stats(self.a[k],self.b[k],loc=self.mus[k],scale=self.sigmas[k],moments="m")
            variances[k] = truncnorm.stats(self.a[k], self.b[k],loc=self.mus[k],scale=self.sigmas[k],moments="v")

        EX = np.sum(self._weights_trunc * means)
        EX2 = np.sum(self._weights_trunc * (variances + means**2))

        return float(np.sqrt(max(EX2 - EX**2, 0.0)))

    def median(self):
        lo = self.L
        hi = self.R

        if not np.isfinite(lo):
            lo = np.min(self.mus) - 12 * np.max(self.sigmas)

        if not np.isfinite(hi):
            hi = np.max(self.mus) + 12 * np.max(self.sigmas)

        for _ in range(100):
            mid = 0.5 * (lo + hi)
            fmid = self.cdf(mid)

            if abs(fmid - 0.5) < 1e-12:
                return mid

            if fmid < 0.5:
                lo = mid
            else:
                hi = mid

        return 0.5 * (lo + hi)

    def get_d(self):
        return 1

    def get_med_std(self):
        med = self.median()
        pdf_med = self.pdf(med)
        return 1.0 / (2.0 * pdf_med) if pdf_med > 0 else np.inf

    def get_bounds(self):
        return self.L, self.R



# ------ theoretical and emprical variance of median -------
def theoretical_variance_of_median(dist, n=None):
    """theoretical variance of the median is approx 1/(4*n*f(median)^2)"""
    f_median = dist.pdf(dist.median())
    if n is not None:
        return 1/(4*n*f_median**2)
    else:
        return 1/(4*f_median**2)

def plot_median_of_sample_size_m_and_size_n():
    gmm = TruncGMM1D(mus=[-2.0, 1.5, 4.0], sigmas=[1.0, 0.15, 1.2], weights=[0.4, 0.1, 0.5], trunc_left=-5.0,
                     trunc_right=8.0)

    # gmm = TruncGMM1D(mus=[0, 5], sigmas=[0.0001, 3], weights=[0.505, 0.495], trunc_left=-10, trunc_right=10)

    m= 100
    n=1000

    m_medians = [np.median(gmm.sample(m)) for _ in range(5000)]
    n_medians = [np.median(gmm.sample(n)) for _ in range(5000)]
    # print the fraction of median that <1
    print(f"Fraction of medians < 1 for m={m}: {(np.array(m_medians) < 1).mean():.4f}")


    #plot
    import matplotlib.pyplot as plt
    all_medians = np.concatenate([m_medians, n_medians])
    bins = np.linspace(all_medians.min(), all_medians.max(), 50)

    plt.hist(m_medians, bins=bins, alpha=0.5, label=f'm={m}')
    plt.hist(n_medians, bins=bins, alpha=0.5, label=f'n={n}')
    plt.axvline(gmm.median(), color='red', linestyle='--', label='True Median')
    plt.legend()
    plt.title('Distribution of Sample Medians')
    plt.xlabel('Median Value')
    plt.ylabel('Frequency')
    plt.show()





if __name__ == '__main__':
    # expriment: sample data of size 1000 from unif[0,1], and for 1000 times, plot the mean+lap noise with scale 0.1. then, plot it.
    import matplotlib.pyplot as plt
    import numpy as np

    scale =0.3
    noise1 = np.random.laplace(scale=scale, size=1000)
    noise2 = np.random.laplace(loc=-1,scale=scale, size=1000)

    #plot hist of medians, shared bins, with alpha 0.5
    bins = np.linspace(min(noise1.min(), noise2.min()), max(noise1.max(), noise2.max()), 50)
    plt.hist(noise1, bins=bins, alpha=0.5,)
    plt.hist(noise2, bins=bins, alpha=0.5)
    plt.title('Distribution of Noisy Medians')
    plt.xlabel('Noisy Median Value')
    plt.ylabel('Frequency')
    plt.show()