__author__ = 'Tom Schaul, tom@idsia.ch'

from pybrain.utilities import blockCombine
from scipy import mat, dot, outer
from scipy.linalg import inv, cholesky



def calcFisherInformation(sigma, invSigma=None, factorSigma=None):
    """ Compute the exact Fisher Information Matrix of a Gaussian distribution,
    given its covariance matrix.
    Returns a list of the diagonal blocks. """
    if invSigma == None:
        invSigma = inv(sigma)
    if factorSigma == None:
        factorSigma = cholesky(sigma)
    dim = sigma.shape[0]
    fim = [invSigma]
    for k in range(dim):
        D = invSigma[k:, k:].copy()
        D[0, 0] += factorSigma[k, k] ** -2
        fim.append(D)
    return fim



def calcInvFisher(sigma, invSigma=None, factorSigma=None):
    """ Efficiently compute the exact inverse of the FIM of a Gaussian.
    Returns a list of the diagonal blocks. """
    if invSigma == None:
        invSigma = inv(sigma)
    if factorSigma == None:
        factorSigma = cholesky(sigma)
    dim = sigma.shape[0]

    invF = [mat(1 / (invSigma[-1, -1] + factorSigma[-1, -1] ** -2))]
    invD = 1 / invSigma[-1, -1]
    for k in reversed(list(range(dim - 1))):
        v = invSigma[k + 1:, k]
        w = invSigma[k, k]
        wr = w + factorSigma[k, k] ** -2
        u = dot(invD, v)
        s = dot(v, u)
        q = 1 / (w - s)
        qr = 1 / (wr - s)
        t = -(1 + q * s) / w
        tr = -(1 + qr * s) / wr
        invF.append(blockCombine([[qr, tr * u], [mat(tr * u).T, invD + qr * outer(u, u)]]))
        invD = blockCombine([[q , t * u], [mat(t * u).T, invD + q * outer(u, u)]])

    invF.append(sigma)
    invF.reverse()
    return invF

