/*-----------------------------------------------------------------------------
  Estimation routine
  Note: zvar_inv = 1/zvar
------------------------------------------------------------------------------*/
#include <math.h>
#include <stdlib.h>
#include <stdio.h>

#include "VTS.h"



float improvement(float a, float b)
{
	if (a != 0) {
		if (a > 0)
			return ((a - b) / a);
		else
			return ((b - a) / a);
	} else
		return (100);
}


int estimate(float **z,
	     float *prior,
	     float **xmean,
	     float **xvar,
	     float *nvar, float *q, float *n, int nz, int nmodes, int ndim)
{
	float **f = NULL, **wta = NULL, **wtb = NULL, **wtaa =
	    NULL, **wtab = NULL, **wtbb = NULL;
	float *pk = NULL, **zmean = NULL, **zvar_inv = NULL, *zcorprod =
	    NULL;
	float *A = NULL, *B = NULL, *C = NULL, *D = NULL, *E = NULL, *F =
	    NULL, *G = NULL;
	float likelihood, prevlikelihood, deltalikelihoodratio, T;
	float oldq[256], oldn[256], oldnv[256];
	float oldoldq[256], oldoldn[256], oldoldnv[256];

	int i, j, k, iterations = 0;

	if (!allocate
	    (&f, &wta, &wtb, &wtaa, &wtab, &wtbb, &zmean, &zvar_inv,
	     &zcorprod, &pk, &A, &B, &C, &D, &E, &F, &G, nmodes, ndim)) {
		WARNING(("Unable to allocate for all variables in estimate\n"));
		free_all(f, wta, wtb, wtaa, wtab, wtbb, zmean, zvar_inv,
			 zcorprod, pk, A, B, C, D, E, F, G);
		return (FAILURE);
	}

	prevlikelihood = -1e+30;
	deltalikelihoodratio = 10;
	while ((deltalikelihoodratio > SIGNIFICANT) && (iterations < 5)) {
		compute_wts(xmean, xvar, nmodes, ndim, n, q, wta, wtb,
			    wtaa, wtab, wtbb);
		compute_f(n, q, xmean, f, wta, wtb, nmodes, ndim);
		compute_zstats(zmean, zvar_inv, zcorprod, wta, wtb, f,
			       xmean, xvar, prior, n, q, nvar, nmodes,
			       ndim);

		for (i = 0; i < ndim; ++i)
			A[i] = B[i] = C[i] = D[i] = E[i] = F[i] = G[i] = 0;
		likelihood = 0;
		for (i = 0; i < nz; ++i) {
			likelihood +=
			    expectation(pk, z[i], zmean, zvar_inv,
					zcorprod, nmodes, ndim);
			for (k = 0; k < ndim; ++k) {
				for (j = 0; j < nmodes; ++j) {
					A[k] +=
					    pk[j] * (z[i][k] - f[j][k] -
						     xmean[j][k] *
						     wta[j][k]) *
					    wta[j][k] * zvar_inv[j][k];
					B[k] +=
					    pk[j] * wtaa[j][k] *
					    zvar_inv[j][k];
					C[k] +=
					    pk[j] * wtab[j][k] *
					    zvar_inv[j][k];
					D[k] +=
					    pk[j] * (z[i][k] - f[j][k] -
						     xmean[j][k] *
						     wta[j][k]) *
					    wtb[j][k] * zvar_inv[j][k];
					E[k] +=
					    pk[j] * wtbb[j][k] *
					    zvar_inv[j][k];
					T = z[i][k] -
					    wta[j][k] * (xmean[j][k] +
							 q[k]) -
					    wtb[j][k] * n[k] - f[j][k];
					F[k] +=
					    pk[j] * zvar_inv[j][k] *
					    zvar_inv[j][k] * (T * T *
							      wtbb[j][k] -
							      xvar[j][k] *
							      wtaa[j][k] *
							      wtbb[j][k]);
					G[k] +=
					    pk[j] * zvar_inv[j][k] *
					    zvar_inv[j][k] * wtbb[j][k] *
					    wtbb[j][k];
				}
			}
		}
		likelihood /= nz;
		deltalikelihoodratio =
		    improvement(likelihood, prevlikelihood);
		/*
		 * Update ONLY if likelihood increases
		 */
		for (i = 0; i < ndim; i++) {
			oldoldnv[i] = oldnv[i];
			oldoldq[i] = oldq[i];
			oldoldn[i] = oldn[i];
			oldnv[i] = nvar[i];
			oldq[i] = q[i];
			oldn[i] = n[i];
		}
		if (deltalikelihoodratio > 0)
			compute_params(n, q, nvar, A, B, C, D, E, F, G,
				       wta, wtb, xvar, nmodes, ndim);
		else {
			for (i = 0; i < ndim; i++) {
				q[i] = oldoldq[i];
				n[i] = oldoldn[i];
				nvar[i] = oldoldnv[i];
			}
		}

		prevlikelihood = likelihood;
		++iterations;
		printf("Likelihood after %d iterations = %f\n", iterations,
		       likelihood);
		fflush(stdout);
	}

	free_all(f, wta, wtb, wtaa, wtab, wtbb, zmean, zvar_inv, zcorprod,
		 pk, A, B, C, D, E, F, G);
	if (deltalikelihoodratio > 0)
		return (SUCCESS);
	else {
		//printf("Likelihoods did not converge! VTS failed!\n");
		printf("Likelihoods did not converge! Use result from iteration [%d]\n",iterations-1);
		fflush(stdout);
		//return (FAILURE);
		return (SUCCESS);
	}
}
