
import numpy as np
from scipy import special

def t_fit(X, dof=4, iter=200, tol=1e-7):
    '''t
    Estimates the mean and covariance of the dataset
    X (rows are datapoints) assuming they come from a
    student t likelihood with no priors and dof degrees
    of freedom using the EM algorithm.
    :param X: np.array[n,d]
    :param dof: float > 2
    :param iter: int, max iterations of algorithm
    :param tol: float, tolerance
    :return: estimated covariance, estimated mean, list of
             objectives at each iteration.
    :rtype: np.array[d,d], np.array[d], list[float]
    '''
    # initialize parameters
    D = X.shape[1]
    N = X.shape[0]
    cov = np.cov(X,rowvar=False)
    mean = X.mean(axis=0)
    mu = X - mean[None,:]
    delta = np.einsum('ij,ij->i', mu, np.linalg.solve(cov,mu.T).T)
    z = (dof + D) / (dof + delta)
    obj = [
        -N*np.linalg.slogdet(cov)[1]/2 - (z*delta).sum()/2 \
        -N*special.gammaln(dof/2) + N*dof*np.log(dof/2)/2 + dof*(np.log(z)-z).sum()/2
    ]

    # iterate
    for i in range(iter):
        # M step
        mean = (X * z[:,None]).sum(axis=0).reshape(-1,1) / z.sum()
        mu = X - mean.squeeze()[None,:]
        cov = np.einsum('ij,ik->jk', mu, mu * z[:,None])/N

        # E step
        delta = (mu * np.linalg.solve(cov,mu.T).T).sum(axis=1)
        delta = np.einsum('ij,ij->i', mu, np.linalg.solve(cov,mu.T).T)
        z = (dof + D) / (dof + delta)

        # store objective
        obj.append(
            -N*np.linalg.slogdet(cov)[1]/2 - (z*delta).sum()/2 \
            -N*special.gammaln(dof/2) + N*dof*np.log(dof/2)/2 + dof*(np.log(z)-z).sum()/2
        )   
        
        if np.abs(obj[-1] - obj[-2]) < tol:
            break

    return mean.squeeze(), cov

def t_generate(mu, S, dof=4, n=1):
    '''generate random variables of multivariate t distribution
    Parameters
    ----------
    mu : array_like
        mean of random variable, length determines dimension of random variable
    S : array_like
        square array of covariance  matrix
    d0f : int or float
        degrees of freedom
    n : int
        number of observations, return random array will be (n, len(mu))
    Returns
    -------
    rvs : ndarray, (n, len(mu))
        each row is an independent draw of a multivariate t distributed
        random variable
    '''
    mu = np.asarray(mu)
    d = len(mu)
    if dof == np.inf:
        x = 1.
    else:
        x = np.random.chisquare(dof, n)/dof
    z = np.random.multivariate_normal(np.zeros(d),S,(n,))
    return mu + z/np.sqrt(x)[:,None]