"""S-Distribution functions.

Density, cumulative distribution, quantile functions and random number
generation for the S distribution.
"""

import numpy as np

from .gnorm import qgnorm


def ds(q, loc=0, scale=1, log=False):
    """S-distribution density.

    Density function: f(x) = 1/(4*scale^2) * exp(-sqrt(abs(loc - x)) / scale)

    Parameters
    ----------
    q : array_like
        Quantiles.
    loc : float
        Location parameter.
    scale : float
        Scale parameter.
    log : bool
        If True, return log-density.

    Returns
    -------
    array
        Density values.
    """
    if log:
        return -np.log(4) - 2 * np.log(scale) - np.sqrt(np.abs(loc - q)) / scale
    density = 1 / (4 * scale**2) * np.exp(-np.sqrt(np.abs(loc - q)) / scale)
    return density


def ps(q, loc=0, scale=1):
    """S-distribution CDF."""
    sign_val = np.sign(q - loc)
    sqrt_term = np.sqrt(np.abs(loc - q))
    return 0.5 + 0.5 * sign_val * (
        1 - 1 / scale * (sqrt_term + scale) * np.exp(-sqrt_term / scale)
    )


def qs(p, loc=0, scale=1):
    """S-distribution quantile function."""
    p = np.asarray(p)
    return qgnorm(p, loc=loc, scale=scale**2, shape=0.5)


def rs(n, loc=0, scale=1):
    """S-distribution random number generation."""
    return qs(np.random.uniform(0, 1, n), loc, scale)
