"""gradient/Jacobian of normal and t loglikelihood

use chain rule

normal derivative wrt mu, sigma and beta

new version: loc-scale distributions, derivative wrt loc, scale

also includes "standardized" t distribution (for use in GARCH)

TODO:
* use sympy for derivative of loglike wrt shape parameters
  it works for df of t distribution dlog(gamma(a))da = polygamma(0,a) check
  polygamma is available in scipy.special
* get loc-scale example to work with mean = X*b
* write some full unit test examples

A: josef-pktd

"""

import numpy as np
from scipy import special
from scipy.special import gammaln


def norm_lls(y, params):
    """normal loglikelihood given observations and mean mu and variance sigma2

    Parameters
    ----------
    y : ndarray, 1d
        normally distributed random variable
    params : ndarray, (nobs, 2)
        array of mean, variance (mu, sigma2) with observations in rows

    Returns
    -------
    lls : ndarray
        contribution to loglikelihood for each observation
    """

    mu, sigma2 = params.T
    lls = -0.5 * (np.log(2 * np.pi) + np.log(sigma2) + (y - mu) ** 2 / sigma2)
    return lls


def norm_lls_grad(y, params):
    """Jacobian of normal loglikelihood wrt mean mu and variance sigma2

    Parameters
    ----------
    y : ndarray, 1d
        normally distributed random variable
    params : ndarray, (nobs, 2)
        array of mean, variance (mu, sigma2) with observations in rows

    Returns
    -------
    grad : array (nobs, 2)
        derivative of loglikelihood for each observation wrt mean in first
        column, and wrt variance in second column

    Notes
    -----
    this is actually the derivative wrt sigma not sigma**2, but evaluated
    with parameter sigma2 = sigma**2

    """
    mu, sigma2 = params.T
    dllsdmu = (y - mu) / sigma2
    dllsdsigma2 = ((y - mu) ** 2 / sigma2 - 1) / np.sqrt(sigma2)
    return np.column_stack((dllsdmu, dllsdsigma2))


def mean_grad(x, beta):
    """gradient/Jacobian for d (x*beta)/ d beta"""
    return x


def normgrad(y, x, params):
    """Jacobian of normal loglikelihood wrt mean mu and variance sigma2

    Parameters
    ----------
    y : ndarray, 1d
        normally distributed random variable with mean x*beta, and variance sigma2
    x : ndarray, 2d
        explanatory variables, observation in rows, variables in columns
    params : array_like, (nvars + 1)
        array of coefficients and variance (beta, sigma2)

    Returns
    -------
    grad : array (nobs, 2)
        derivative of loglikelihood for each observation wrt mean in first
        column, and wrt scale (sigma) in second column
    assume params = (beta, sigma2)

    Notes
    -----
    TODO: for heteroscedasticity need sigma to be a 1d array

    """
    beta = params[:-1]
    sigma2 = params[-1] * np.ones((len(y), 1))
    dmudbeta = mean_grad(x, beta)
    mu = np.dot(x, beta)
    # print(beta, sigma2)
    params2 = np.column_stack((mu, sigma2))
    dllsdms = norm_lls_grad(y, params2)
    grad = np.column_stack((dllsdms[:, :1] * dmudbeta, dllsdms[:, :1]))
    return grad


def tstd_lls(y, params, df):
    """t loglikelihood given observations and mean mu and variance sigma2 = 1

    Parameters
    ----------
    y : ndarray, 1d
        normally distributed random variable
    params : ndarray, (nobs, 2)
        array of mean, variance (mu, sigma2) with observations in rows
    df : int
        degrees of freedom of the t distribution

    Returns
    -------
    lls : ndarray
        contribution to loglikelihood for each observation

    Notes
    -----
    parametrized for garch
    """

    mu, sigma2 = params.T
    df = df * 1.0
    # lls = gammaln((df+1)/2.) - gammaln(df/2.) - 0.5*np.log((df-2)*np.pi)
    # lls -= (df+1)/2. * np.log(1. + (y-mu)**2/(df-2.)/sigma2) + 0.5 * np.log(sigma2)
    lls = gammaln((df + 1) / 2.0) - gammaln(df / 2.0) - 0.5 * np.log((df - 2) * np.pi)
    lls -= (df + 1) / 2.0 * np.log(
        1.0 + (y - mu) ** 2 / (df - 2) / sigma2
    ) + 0.5 * np.log(sigma2)

    return lls


def norm_dlldy(y):
    """derivative of log pdf of standard normal with respect to y"""
    return -y


def tstd_pdf(x, df):
    """pdf for standardized (not standard) t distribution, variance is one"""

    r = np.array(df * 1.0)
    Px = np.exp(special.gammaln((r + 1) / 2.0) - special.gammaln(r / 2.0)) / np.sqrt(
        (r - 2) * np.pi
    )
    Px /= (1 + (x**2) / (r - 2)) ** ((r + 1) / 2.0)
    return Px


def ts_lls(y, params, df):
    """t loglikelihood given observations and mean mu and variance sigma2 = 1

    Parameters
    ----------
    y : ndarray, 1d
        normally distributed random variable
    params : ndarray, (nobs, 2)
        array of mean, variance (mu, sigma2) with observations in rows
    df : int
        degrees of freedom of the t distribution

    Returns
    -------
    lls : ndarray
        contribution to loglikelihood for each observation

    Notes
    -----
    parametrized for garch
    normalized/rescaled so that sigma2 is the variance

    >>> df = 10; sigma = 1.
    >>> stats.t.stats(df, loc=0., scale=sigma.*np.sqrt((df-2.)/df))
    (array(0.0), array(1.0))
    >>> sigma = np.sqrt(2.)
    >>> stats.t.stats(df, loc=0., scale=sigma*np.sqrt((df-2.)/df))
    (array(0.0), array(2.0))
    """
    print(y, params, df)
    mu, sigma2 = params.T
    df = df * 1.0
    # lls = gammaln((df+1)/2.) - gammaln(df/2.) - 0.5*np.log((df-2)*np.pi)
    # lls -= (df+1)/2. * np.log(1. + (y-mu)**2/(df-2.)/sigma2) + 0.5 * np.log(sigma2)
    lls = gammaln((df + 1) / 2.0) - gammaln(df / 2.0) - 0.5 * np.log((df) * np.pi)
    lls -= (df + 1.0) / 2.0 * np.log(
        1.0 + (y - mu) ** 2 / (df) / sigma2
    ) + 0.5 * np.log(sigma2)
    return lls


def ts_dlldy(y, df):
    """derivative of log pdf of standard t with respect to y

    Parameters
    ----------
    y : array_like
        data points of random variable at which loglike is evaluated
    df : array_like
        degrees of freedom,shape parameters of log-likelihood function
        of t distribution

    Returns
    -------
    dlldy : ndarray
        derivative of loglikelihood wrt random variable y evaluated at the
        points given in y

    Notes
    -----
    with mean 0 and scale 1, but variance is df/(df-2)

    """
    df = df * 1.0
    # (df+1)/2. / (1 + y**2/(df-2.)) * 2.*y/(df-2.)
    # return -(df+1)/(df-2.) / (1 + y**2/(df-2.)) * y
    return -(df + 1) / (df) / (1 + y**2 / (df)) * y


def tstd_dlldy(y, df):
    """derivative of log pdf of standardized t with respect to y

        Parameters
        ----------
    y : array_like
        data points of random variable at which loglike is evaluated
    df : array_like
        degrees of freedom,shape parameters of log-likelihood function
        of t distribution

    Returns
    -------
    dlldy : ndarray
        derivative of loglikelihood wrt random variable y evaluated at the
        points given in y


    Notes
    -----
    parametrized for garch, standardized to variance=1
    """
    # (df+1)/2. / (1 + y**2/(df-2.)) * 2.*y/(df-2.)
    return -(df + 1) / (df - 2.0) / (1 + y**2 / (df - 2.0)) * y
    # return (df+1)/(df) / (1 + y**2/(df)) * y


def locscale_grad(y, loc, scale, dlldy, *args):
    """derivative of log-likelihood with respect to location and scale

    Parameters
    ----------
    y : array_like
        data points of random variable at which loglike is evaluated
    loc : float
        location parameter of distribution
    scale : float
        scale parameter of distribution
    dlldy : function
        derivative of loglikelihood fuction wrt. random variable x
    args : array_like
        shape parameters of log-likelihood function

    Returns
    -------
    dlldloc : ndarray
        derivative of loglikelihood wrt location evaluated at the
        points given in y
    dlldscale : ndarray
        derivative of loglikelihood wrt scale evaluated at the
        points given in y

    """
    yst = (y - loc) / scale  # ystandardized
    dlldloc = -dlldy(yst, *args) / scale
    dlldscale = -1.0 / scale - dlldy(yst, *args) * (y - loc) / scale**2
    return dlldloc, dlldscale


if __name__ == "__main__":
    verbose = 0
    if verbose:
        sig = 0.1
        beta = np.ones(2)
        rvs = np.random.randn(10, 3)
        x = rvs[:, 1:]
        y = np.dot(x, beta) + sig * rvs[:, 0]

        params = [1, 1, 1]
        print(normgrad(y, x, params))

        dllfdbeta = (y - np.dot(x, beta))[:, None] * x  # for sigma = 1
        print(dllfdbeta)

        print(locscale_grad(y, np.dot(x, beta), 1, norm_dlldy))
        print(y - np.dot(x, beta))

    from scipy import misc, stats

    def llt(y, loc, scale, df):
        return np.log(stats.t.pdf(y, df, loc=loc, scale=scale))

    def lltloc(loc, y, scale, df):
        return np.log(stats.t.pdf(y, df, loc=loc, scale=scale))

    def lltscale(scale, y, loc, df):
        return np.log(stats.t.pdf(y, df, loc=loc, scale=scale))

    def llnorm(y, loc, scale):
        return np.log(stats.norm.pdf(y, loc=loc, scale=scale))

    def llnormloc(loc, y, scale):
        return np.log(stats.norm.pdf(y, loc=loc, scale=scale))

    def llnormscale(scale, y, loc):
        return np.log(stats.norm.pdf(y, loc=loc, scale=scale))

    if verbose:
        print("\ngradient of t")
        print(misc.derivative(llt, 1, dx=1e-6, n=1, args=(0, 1, 10), order=3))
        print("t ", locscale_grad(1, 0, 1, tstd_dlldy, 10))
        print("ts", locscale_grad(1, 0, 1, ts_dlldy, 10))
        print(
            misc.derivative(llt, 1.5, dx=1e-10, n=1, args=(0, 1, 20), order=3),
        )
        print("ts", locscale_grad(1.5, 0, 1, ts_dlldy, 20))
        print(
            misc.derivative(llt, 1.5, dx=1e-10, n=1, args=(0, 2, 20), order=3),
        )
        print("ts", locscale_grad(1.5, 0, 2, ts_dlldy, 20))
        print(
            misc.derivative(llt, 1.5, dx=1e-10, n=1, args=(1, 2, 20), order=3),
        )
        print("ts", locscale_grad(1.5, 1, 2, ts_dlldy, 20))
        print(
            misc.derivative(lltloc, 1, dx=1e-10, n=1, args=(1.5, 2, 20), order=3),
        )
        print(misc.derivative(lltscale, 2, dx=1e-10, n=1, args=(1.5, 1, 20), order=3))
        y, loc, scale, df = 1.5, 1, 2, 20
        print("ts", locscale_grad(y, loc, scale, ts_dlldy, 20))
        print(
            misc.derivative(lltloc, loc, dx=1e-10, n=1, args=(y, scale, df), order=3),
        )
        print(
            misc.derivative(lltscale, scale, dx=1e-10, n=1, args=(y, loc, df), order=3)
        )

        print("\ngradient of norm")
        print(misc.derivative(llnorm, 1, dx=1e-6, n=1, args=(0, 1), order=3))
        print(locscale_grad(1, 0, 1, norm_dlldy))
        y, loc, scale = 1.5, 1, 2
        print("ts", locscale_grad(y, loc, scale, norm_dlldy))
        print(
            misc.derivative(llnormloc, loc, dx=1e-10, n=1, args=(y, scale), order=3),
        )
        print(
            misc.derivative(llnormscale, scale, dx=1e-10, n=1, args=(y, loc), order=3)
        )
        y, loc, scale = 1.5, 0, 1
        print("ts", locscale_grad(y, loc, scale, norm_dlldy))
        print(
            misc.derivative(llnormloc, loc, dx=1e-10, n=1, args=(y, scale), order=3),
        )
        print(
            misc.derivative(llnormscale, scale, dx=1e-10, n=1, args=(y, loc), order=3)
        )
        # print('still something wrong with handling of scale and variance'
        # looks ok now
        print("\nloglike of t")
        print(
            tstd_lls(1, np.array([0, 1]), 100),
            llt(1, 0, 1, 100),
            "differently standardized",
        )
        print(
            tstd_lls(1, np.array([0, 1]), 10),
            llt(1, 0, 1, 10),
            "differently standardized",
        )
        print(ts_lls(1, np.array([0, 1]), 10), llt(1, 0, 1, 10))
        print(tstd_lls(1, np.array([0, 1.0 * 10.0 / 8.0]), 10), llt(1.0, 0, 1.0, 10))
        print(ts_lls(1, np.array([0, 1]), 100), llt(1, 0, 1, 100))

        print(tstd_lls(1, np.array([0, 1]), 10), llt(1, 0, 1.0 * np.sqrt(8 / 10.0), 10))

    from numpy.testing import assert_almost_equal

    params = [(0, 1), (1.0, 1.0), (0.0, 2.0), (1.0, 2.0)]
    yt = np.linspace(-2.0, 2.0, 11)
    for loc, scale in params:
        dlldlo = misc.derivative(
            llnormloc, loc, dx=1e-10, n=1, args=(yt, scale), order=3
        )
        dlldsc = misc.derivative(
            llnormscale, scale, dx=1e-10, n=1, args=(yt, loc), order=3
        )
        gr = locscale_grad(yt, loc, scale, norm_dlldy)
        assert_almost_equal(dlldlo, gr[0], 5, err_msg="deriv loc")
        assert_almost_equal(dlldsc, gr[1], 5, err_msg="deriv scale")
    for df in [3, 10, 100]:
        for loc, scale in params:
            dlldlo = misc.derivative(
                lltloc, loc, dx=1e-10, n=1, args=(yt, scale, df), order=3
            )
            dlldsc = misc.derivative(
                lltscale, scale, dx=1e-10, n=1, args=(yt, loc, df), order=3
            )
            gr = locscale_grad(yt, loc, scale, ts_dlldy, df)
            assert_almost_equal(dlldlo, gr[0], 4, err_msg="deriv loc")
            assert_almost_equal(dlldsc, gr[1], 4, err_msg="deriv scale")
            assert_almost_equal(
                ts_lls(yt, np.array([loc, scale**2]), df),
                llt(yt, loc, scale, df),
                5,
                err_msg="loglike",
            )
            assert_almost_equal(
                tstd_lls(yt, np.array([loc, scale**2]), df),
                llt(yt, loc, scale * np.sqrt((df - 2.0) / df), df),
                5,
                err_msg="loglike",
            )
