import numpy as np

from numpy.typing import ArrayLike

class MarchenkoPastur:
    """Definition of a Marchenko-Pastur distribution"""

    def __init__(self, ratio: float, sigma: float = 1.0):
        """
        Parameters
        ----------
        ratio : float
            The ratio between the number of variables (columns) and the size of the sample (rows) contained in the data matrix. For numerical stability, it should be less than 1.
        sigma : float
            The standard deviation of the distribution of values, by default, 1.0.

        Raises
        ------
        ValueError
            If ratio or sigma are not strictly positive.
        """
        if ratio <= 0.0:
            raise ValueError("The ratio must be strictly positive, but found %s <= 0.0!" % ratio)
        self.ratio = ratio
        if sigma <= 0.0:
            raise ValueError("The standard deviation must be strictly positive, but found %s <= 0.0!" % sigma)
        self.sigma = sigma

        # Compute the limits of the distribution
        self.l_bottom = sigma**2 * (1.0 - np.sqrt(self.ratio))**2
        self.l_upper = sigma**2 * (1.0 + np.sqrt(self.ratio))**2

    def pdf(self, x: float | ArrayLike) -> float | ArrayLike:
        """
        Return the value of the probability distribution function.

        Parameters
        ----------
        x : float | ArrayLike
            The value(s) at which to compute the PDF.

        Returns
        -------
        float | ArrayLike
            The value(s) of the PDF
        """
        if not np.isscalar(x):
            return np.vectorize(self.pdf, otypes=[float])(x)

        if x == 0.0:
          return 0.0

        num = np.sqrt(max(self.l_upper - x, 0.0) * max(x - self.l_bottom, 0.0))
        den = 2.0 * np.pi * self.sigma**2 * self.ratio * x

        return float(num / den)