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


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

    The probability density function is::

        f(x; nu, b) = 2/Gamma(nu) * b^((nu+1)/2) * x^((nu-1)/2) * K_{nu-1}(2*sqrt(b*x))

    for x >= 0, where K_{nu-1} is the modified Bessel function of the second
    kind of order nu-1, nu > 0 is the shape parameter and b > 0 is the rate
    parameter.
    """

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


    def _argcheck(self, nu, b):
        return (nu > 0) & (b > 0)

    def _pdf(self, x, nu, b):
        # k.pdf(x, nu, b) = 2/Gamma(nu) * b^((nu+1)/2) * x^((nu-1)/2) * K_{nu-1}(2*sqrt(b*x))
        return np.exp(self._logpdf(x, nu, b))

    def _logpdf(self, x, nu, b):
        return (
            np.log(2.0)
            - gammaln(nu)
            + (nu + 1) / 2.0 * np.log(b)
            + (nu - 1) / 2.0 * np.log(x)
            + np.log(kv(nu - 1, 2.0 * np.sqrt(b * x)))
        )


k_dist = k_gen(a=0.0, name="k_distribution", shapes="nu, b")