import fastlmm.inference.lmm_cov as lmm_cov
import numpy as np

def est_h2(Y, K, covariates=None, nGridH2=10000, plot=True, verbose=True):
    """
    This function implements the Bayesian heritability estimate from Furlotte et al., 2014

    Furlotte, Nicholas A., David Heckerman, and Christoph Lippert.
    "Quantifying the uncertainty in heritability." Journal of human genetics 59.5 (2014): 269-275.

    Args:
        Y:              [N x 1] np.ndarray of phenotype values
        K:              [N x N] np.ndarray of kinship values
        covariates:     [N x D] np.ndarray of covariate values [default: None]
        plot:           Boolean, create a plot? [default: True]
        verbose:        print results? [default: True]

    returns:
        REML estimate of h^2 (as in Yang et al. 2010)
        posterior mean of h^2
        posterior variance of h^2
        h2 values on a grid
        posterior for h2 values on a grid
    """
    lmm = lmm_cov.LMM(forcefullrank=False, X=covariates, linreg=None, Y=Y, G=None, K=K, regressX=True, inplace=False)
    h2 = lmm.findH2()
    h2_posterior = lmm.posterior_h2(nGridH2=nGridH2)
    logp = -h2_posterior[2]
    grid = h2_posterior[1]
    post_h2 = np.exp(logp-logp.max())/np.exp(logp-logp.max()).sum()*logp.shape[0]
    h2_mean = (np.exp(logp-logp.max())*grid[:,np.newaxis]).sum()/np.exp(logp-logp.max()).sum()
    h2_var = (np.exp(logp-logp.max())*(grid[:,np.newaxis] - h2_mean)**2.0).sum()/np.exp(logp-logp.max()).sum()
    if plot:
        import pylab as plt
        plt.figure()
        plt.plot([h2['h2'],h2['h2']],[0,1],"r")
        plt.plot([h2_mean,h2_mean],[0,1],"g")
        plt.legend(["REML-estimate = %.3f"%h2['h2'],"posterior mean = %.3f"%h2_mean])
        plt.plot(grid.flatten(),post_h2.flatten())
        plt.xlabel("$h^2$")
        plt.ylabel("$p( h^2 | Data)$")
    if verbose:
        print "max[h^2] = %.5f  E[h^2] = %.5f +- %.5f" % (h2['h2'], h2_mean,np.sqrt(h2_var))
    return h2, h2_mean, h2_var, grid, post_h2 