import pylab as pl
from scipy.misc import derivative
import inspect

def cramer_rao(model, p0, X, noise, show_plot=False):
    """Calulate inverse of the Fisher information matrix for model
    sampled on grid X with parameters p0. Assumes samples are not
    correlated and have equal variance noise^2.

    Parameters
    ----------
    model : callable
        The model function, f(x, ...).  It must take the independent
        variable as the first argument and the parameters as separate
        remaining arguments.
    X : array
        Grid where model is sampled.
    p0 : M-length sequence
        Point in parameter space where Fisher information matrix is
        evaluated.
    noise: scalar
        Squared variance of the noise in data.
    show_plot : boolean
        If True shows plot.

    Returns
    -------
    iI : 2d array
        Inverse of Fisher information matrix.
    """
    labels = inspect.getargspec(model)[0][1:]
    p0dict = dict(zip(inspect.getargspec(model)[0][1:], p0))

    D = pl.zeros((len(p0), X.size))
    for i, argname in enumerate(labels):
        D[i,:] = [
            derivative(
                lambda p: model(x, **dict(p0dict, **{argname: p})),
                p0dict[argname],
                dx=0.0001)
            for x in X ]

    if show_plot:
        pl.figure()
        pl.plot(X, model(X, *p0), '--k', lw=2, label='signal')
        for d, label in zip(D, labels):
            pl.plot(X, d, '.-', label=label)
        pl.legend(loc='best')
        pl.title('Parameter dependence on particular data point')

    I = 1/noise**2 * pl.einsum('mk,nk', D, D)
    iI = pl.inv(I)

    return iI

if __name__=='__main__':
    def sig(t, a, f, ph):
        return a * pl.sin(2 * pl.pi * f * t + ph)

    p0 = [2.12, 8, 0]
    noise = 0.09

    T = pl.arange(-1, 1, 0.01)
    C = cramer_rao(sig, p0, T, noise, show_plot=True)

    print('Inverse of the Fisher information matrix:')
    print(C)
    print('numerical: ', C.diagonal())

    var_a = noise**2 / T.size * 2
    var_w = 12 * noise**2 / ( p0[0]**2 * (T[1]-T[0])**2 * \
        (T.size**3 - T.size) ) * 2
    var_f = var_w / (2*pl.pi)**2
    var_ph = noise**2 / p0[0]**2 / T.size * 2

    print('teoretical:', pl.array([var_a, var_f, var_ph]))
    # ref: D. Rife, R. Boorstyn "Single-Tone Parameter Estimation from
    # Discrete-Time Observations", IEEE Transactions on Information Theory
    # 20, 5, 1974

    pl.show()