/*****************************************************************************
  Routine to perform MMSE for VPS.
  The noramlized vector is computed as the noisy vector - normalization term
  Normalization term = Sum{expectation_of_mode*(shift of mean of mode)}
******************************************************************************/

#include <stdio.h>
#include "VTS.h"

int mmse(float **x,
	 float **z,
	 float *prior,
	 float **mean,
	 float **var,
	 float *nmean,
	 float *nvar, float *q, int nframes, int nmodes, int ndim)
{
	float **f = NULL;
	float **wta = NULL;
	float **wtb = NULL;
	float *pk = NULL;
	float **zmean = NULL;
	float **zvar_inv = NULL;
	float *zcorprod = NULL;

	int i, j, k;

	if (!mmsealloc
	    (&f, &wta, &wtb, &zmean, &zvar_inv, &zcorprod, &pk, nmodes,
	     ndim)) {
		WARNING(("Unable to allocate for all variables in estimate\n"));
		mmsefree_all(f, wta, wtb, zmean, zvar_inv, zcorprod, pk);
		return (FAILURE);
	}

	compute_mwts(mean, var, nmodes, ndim, nmean, q, wta, wtb);
	compute_f(nmean, q, mean, f, wta, wtb, nmodes, ndim);
	compute_zstats(zmean, zvar_inv, zcorprod, wta, wtb, f, mean, var,
		       prior, nmean, q, nvar, nmodes, ndim);

	for (i = 0; i < nframes; ++i) {
		for (k = 0; k < ndim; ++k)
			x[i][k] = z[i][k];
		expectation(pk, z[i], zmean, zvar_inv, zcorprod, nmodes,
			    ndim);
		for (j = 0; j < nmodes; ++j)
			for (k = 0; k < ndim; ++k)
				x[i][k] -=
				    (zmean[j][k] - mean[j][k]) * pk[j];
	}

	mmsefree_all(f, wta, wtb, zmean, zvar_inv, zcorprod, pk);

	return (SUCCESS);
}
