from scipy.stats import rv_continuous
from scipy.special import kv, gammaln
from scipy.stats._distn_infrastructure import _ShapeInfo
import numpy as np


class k_gen(rv_continuous):
    """Generalized K distribution for radar clutter modeling.

    The probability density function is::

        f(x; mu, alpha, beta) = 2 / (Gamma(alpha) * Gamma(beta))
                                * (alpha*beta/mu)^((alpha+beta)/2)
                                * x^((alpha+beta)/2 - 1)
                                * K_{alpha-beta}(2*sqrt(alpha*beta*x/mu))

    for x > 0, where K_{alpha-beta} is the modified Bessel function of the
    second kind of order alpha-beta, mu > 0 is the mean, and alpha > 0,
    beta > 0 are shape parameters.

    The standard (2-parameter) K distribution is the special case beta=1,
    with nu=alpha and b=alpha/mu.
    """

    def _shape_info(self):
        return [
            _ShapeInfo("mu", domain=(0, np.inf), inclusive=(False, True)),
            _ShapeInfo("alpha", domain=(0, np.inf), inclusive=(True, True)),
            _ShapeInfo("beta", domain=(0, np.inf), inclusive=(True, True)),
        ]

    def _argcheck(self, mu, alpha, beta):
        return (mu > 0) & (alpha > 0) & (beta > 0)

    def _pdf(self, x, mu, alpha, beta):
        return np.exp(self._logpdf(x, mu, alpha, beta))

    def _logpdf(self, x, mu, alpha, beta):
        z = 2.0 * np.sqrt(alpha * beta * x / mu)
        return (
            np.log(2.0)
            - gammaln(alpha)
            - gammaln(beta)
            + (alpha + beta) / 2.0 * np.log(alpha * beta / mu)
            + ((alpha + beta) / 2.0 - 1.0) * np.log(x)
            + np.log(kv(alpha - beta, z))
        )


k_dist = k_gen(a=0.0, name="k_distribution", shapes="mu, alpha, beta")