/*  
 *  src/org/metaqtl/algo/EMAlgorithm.java
 * 
 *  $Author: veyrieras $, $Date: 2006-10-17 15:41:55 $, $Version$
 *
 *  MetaQTL : a Java package for QTL Meta-Analysis
 * 
 *  Copyright (C) 2003-2006  Jean-Baptiste Veyrieras, INRA, France.
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA * 
 */
package org.metaqtl.algo;

import java.util.Random;

import org.metaqtl.EMResult;
import org.metaqtl.numrec.NumericalUtilities;

/**
 * This class defines methods to perform the EM-algorithm
 * for univariate mixture of gaussians with known variances. 
 */
public final class EMAlgorithm {
	/**
	 * The number of replicates for a run of the EM-algorithm.
	 */
	public static int EM_START 	  = 50;
	/**
	 * The EM maximum number of iterations.
	 */
	public static int EM_ITER_MAX = 1000;
	/**
	 * The EM convergence tolerance.
	 */
	public static double EM_ERR   = 1.e-8;
	/**
	 * The minimal mahalanobis distance between
	 * components.  
	 */
	public static double EM_MIN_DISTANCE = 1.e-5;
	/**
	 * If true the best starting point is used to
	 * do Supplemental EM (SEM) in order to
	 * compute the DM Matrix.. 
	 */
	public static boolean DO_SEM = true;
	/**
	 * The minimal value for Z proba in EM-clustering
	 */
	public static final double TINY_Z_PROBA 	= Double.MIN_VALUE;
	/**
	 * The EM convergence status : success 
	 */
	public static final int EM_OK = 0;
	/**
	 * The EM convergence status : continue 
	 */
	public static final int EM_CONTINUE = 1;
	/**
	 * The EM convergence status : failure 
	 */
	public static final int EM_FAILURE  = 2;
	/**
	 * Apply the EM-Algorithm on the given data set where <code>x</code>
	 * is an array of observed value and <code>sd</code> an array of same
	 * size than <code>x</code> which stores the standard deviations of 
	 * the observed values. The underlying mixture model will have  
	 * <code>k</code> distinct components and its parameter estimates
	 * will be returned as an {@link EMResult}. If <code>soint</code>
	 * is not null then the EM will start at the given point specified by
	 * the mixture parameters of <code>soint</code>. To control some 
	 * behaviours of the EM-algorithm use the global variables of the class.
	 * 
	 * @param x   the observed positions.
	 * @param sd  the standard deviation of the observed positions.
	 * @param k   the number of mixture components. 
	 * @param spoint the starting point for the EM-algortihm
	 * @return a <code>EMResult</code> from which the parameter estimates can be obtained.
	 * 
	 * @see EMResult
	 */
	public static EMResult doEM(double[] x, double[] sd, int k, EMResult spoint) {
		int i,j,status,iter,start,nstart;
		double max,mu[];
		EMResult   stheta=null,theta=null;
		EMResult[] buff1,buff2;
		
		if (k < 1) return null;
		if (x==null || sd==null) return null;
		
		if (EM_START > 0 && spoint == null) {
			if (k > 1 && k < x.length)
				nstart=EM_START;
			else
				nstart=1;
		} else {
			nstart=1;
		}
		
		buff1 = new EMResult[nstart];
		buff2 = new EMResult[nstart];
		
		Random rng  = new Random();
		
		for(start=0;start<nstart;start++) {
			
			theta   	 = new EMResult(x.length, k);
			buff2[start] = new EMResult(x.length, k);
			mu      	 = new double[k];
			
			if (spoint == null) {
				initEM(x,sd,theta,rng);
				EMResult.copy(buff2[start], theta);
			}
			else {
				EMResult.copy(theta, spoint);
			}
			
			iter = 0;
			do {
				iter++;
				for(i=0;i<theta.k;i++) mu[i] = theta.mu[i];
				status=iterate(x,sd,theta);
				if (iter > EM_ITER_MAX) {
					status = EM_FAILURE;
					break;
				}
				//if (DM_COMPUTE)
					//computeDM(x,sd,mu,theta);
			} while(status == EM_CONTINUE);
			
			buff1[start] = theta;
 			
		}
		
		max = Double.NEGATIVE_INFINITY;
		
		theta=stheta=null;
		
		for(i=0;i<start;i++) {
			if (buff1[i].olog > max) {
				boolean ok=true;
				if (spoint==null && k > 1 && k < x.length) {
					// test if this model is valid
					double[] d=buff1[i].getEuclidean();
					for(j=0;j<d.length;j++)
						if (d[j]<EM_MIN_DISTANCE) break;
					ok=(j==d.length);
				}
				if (ok) {
					max    = buff1[i].olog;
					theta  = buff1[i];
					stheta = buff2[i];
				}
			}
		}
		/*
		 * Do Supplemental EM (SEM) to compute the
		 * DM matrix which will be used to compute
		 * the observed variance covariance matrix.
		 */
		if (theta!=null && DO_SEM) {
			iter = 0;
			mu   = new double[k];
			do {
				iter++;
				for(i=0;i<stheta.k;i++) 
					mu[i] = stheta.mu[i];
				status=iterate(x,sd,stheta);
				if (iter > EM_ITER_MAX) {
					status = EM_FAILURE;
					break;
				}
				if (DO_SEM)
					computeDM(x,sd,mu,stheta);
			} while(status == EM_CONTINUE);
			
			theta.dm = stheta.dm;
			
		}
		
		if (theta!=null) computeCOV(sd, theta);
		
		//System.out.println("log-lik"+theta.olog);
		
		//for(i=0;i<theta.k;i++) {
	    	//System.out.println("mu["+i+"] " + theta.mu[i] +" " + theta.ccov[i] + " " + theta.ocov[i][i]);
	    	//System.out.println(" pi=" + theta.pi[i]);
	    //}
		
		return theta;
	}
	/**
	 * Initializes the EM algorithm by randomly assigning observations to
	 * the clusters and do one M-Step to compute the first values of the 
	 * parameters.
	 * @param x the observed values.
	 * @param sd the standard deviations of the observed values. 
	 * @param theta the current parameter estimates. 
	 */
	public static void initEM(double[] x, double[] sd, EMResult theta, Random rng) {
		 int i,j;
		
		 if (theta.k == 1) {
		 	for(i=0;i<theta.n;i++)
	  	    	theta.z[0][i]=1.0;
	  		theta.pi[0] = 1.0;
	  		mStep(x,sd,theta,EM_ERR);
		 } else if (theta.k == theta.n) {
		 	for(i=0;i<theta.k;i++) {
		   	 	theta.mu[i]   = x[i];
		   	   	theta.pi[i]   = 1.0/(double)theta.k;
		   	   	for(j=0;j<theta.k;j++) {
		   	   		if (j!=i) theta.z[i][j]=0.0;
		   	   		else theta.z[i][j]=1.0;
		   	   	}
		   	 }
		  	 theta.sortCluster();
		 } else {
		   int l,n;
	   	   int[] cidx;
	   	   
	   	   /* A litte boring code to insure
	   	    * that the random clustering is valid, i.e
	   	    * there is as many partition as the number
	   	    * of clusters k
	   	    */
	   	   n    = theta.k;
	   	   cidx = new int[n];
	   	   for(i=0;i<theta.k;i++) cidx[i]=i;
		   for(i=0;i<theta.k;i++) {
		   	  j=rng.nextInt(n);
		   	  for(l=0;l<theta.k;l++) {
		   	  	if (l==cidx[j]) {
		   	  		theta.z[l][i]=1.0;
		   	  	}
		   	  	else theta.z[l][i]=0.0;
		   	  }
		   	  for(l=j;l<n-1;l++)
		   	  	cidx[l]=cidx[l+1];
		   	  n--;
		   }
		   for(i=theta.k;i<theta.n;i++) {
			   /* Draw origin */
			   j=rng.nextInt(theta.k);
		   	   for(l=0;l<theta.k;l++) {
		   	  	if (l==j) {
		   	  		theta.z[l][i]=1.0;
		   	  	}
		   	  	else theta.z[l][i]=0.0;
		   	  }
		   }
		   mStep(x,sd,theta,EM_ERR);
		 }
		 updateLoglikelihood(x, sd, theta);
	}
	/**
	 * This method performs one iteration of the EM-algorithm. An
	 * iteration consists in two steps :
	 * <ul>
	 *    <li> Expectation (E-step) : where the expectation of the 
	 * cluster membership missing variable is computed.</li>
	 *    <li> Maximization (M-step) : where the parameters are updated
	 * to the values which maximize the expectation.</li>
	 * </ul>
	 * If the iteration leads to an update of the parameters which is under
	 * the suited range of convergence error then the method will return
	 * <code>EM_OK<code>, otherwise it will return <code>EM_CONTINUE</code>.
	 * In some cases it can happen - mainly due to numerical roundoff errors -
	 * that the likelihood value after the iteration is lower than the likelihood
	 * before the iteration which in theory cannot happen. Then in this case the
	 * method will return <code>EM_FAILURE</code>.
	 * 
	 * @param x the observed values.
	 * @param sd the standard deviations of the observed values. 
	 * @param theta the current parameter estimates. 
	 * @return the status of the iteration.
	 */
	public static int iterate(double[] x, double[] sd, EMResult theta) {
		  if (theta.k > 1 && theta.k < theta.n) {
			  /* E-Step */
			  eStep(x, sd, theta);
			  /* M-Step */
		  	  return mStep(x, sd, theta, EM_ERR);
		  	  
		  } else if (theta.k == 1) {
		  	
		  	return EM_OK;
	        
		  } else {
		  	 
		   	 return EM_OK;
		  }
	}
	/**
	 * The Expectation step : it just consists in updating the 
	 * Z matrix, i.e the cluster membership probabilities.
	 * 
	 * @param x the observed values.
	 * @param sd the standard deviations of the observed values. 
	 * @param theta the current parameter estimates.
	 * 
	 * @see EMAlgorithm#updateZMatrix(double[], double[], EMResult)
	 */
	public static void eStep(double[] x, double[] sd, EMResult theta)
	{
	    updateZMatrix(x, sd, theta);
	}
	/**
	 * The Maximization step : here this step is straighforward and 
	 * simple analytical formula are applied to obtain the new parameter
	 * estimates. The method returns the status of the maximization.
	 * If the maximization leads to an new value of the parameters which is under
	 * the suited range of convergence error then the method will return
	 * <code>EM_OK<code>, otherwise it will return <code>EM_CONTINUE</code>.
	 * In some cases it can happen - mainly due to numerical roundoff errors -
	 * that the likelihood value after the maximization is lower than the likelihood
	 * before the iteration which in theory cannot happen. Then in this case the
	 * method will return <code>EM_FAILURE</code>
	 * 
	 * @param x the observed values.
	 * @param sd the standard deviations of the observed values. 
	 * @param theta the current parameter estimates.
	 * 
	 * @see EMAlgorithm#updateMuVector(double[], double[], EMResult)
	 * @see EMAlgorithm#updatePiVector(double[], double[], EMResult)
	 * @see EMAlgorithm#updateLoglikelihood(double[], double[], EMResult)
	 */
	public static int mStep(double[] x, double[] sd, EMResult theta, double err)
	{
		double olog;
		double[] mu,pi,tmp1,tmp2;
		
	   	/* Update parameters */
	    mu   = updateMuVector(x, sd, theta);
	    pi   = updatePiVector(theta);
	    /* Swap */
	    tmp1     = theta.mu;
	    theta.mu = mu;
	    mu       = tmp1;
	    tmp2     = theta.pi;
	    theta.pi = pi;
	    pi       = tmp2;
	    olog     = -theta.olog;
	    theta.sortCluster();
	    updateLoglikelihood(x,sd,theta);
	    olog     += theta.olog;
	    
	    if (olog < 0.0) {
	    	/* Swap */
	    	tmp1      = theta.mu;
		    theta.mu = mu;
		    mu       = tmp1;
		    tmp2     = theta.pi;
		    theta.pi = pi;
		    pi       = tmp2;
		    theta.sortCluster();
		    updateLoglikelihood(x,sd,theta);
		    
	    	return EM_FAILURE;
	    	
	    } else if (olog < err) {
	    	
	    	emRate(theta,mu,pi);
	    	
	    	return EM_OK;
	    	
	    } else {
	    	
	    	emRate(theta,mu,pi);
	    	
	    	return EM_CONTINUE;
	    	
	    }
	    
	}
	/**
	 * Compute the euclidean distance between the new parameters and the
	 * old ones to obtain a simple approximation of the EM convergence
	 * rate.
	 * 
	 * @param theta the current parameter estimates.
	 * @param mu the last estimates of the means.
	 * @param pi the last estimates of the mixing.
	 */
	public static void emRate(EMResult theta, double[] mu, double[] pi) {
		int i;
		double r=0.0;
		
		if (theta.k > 1) {
		
			for(i=0;i<theta.k;i++) {
				r+=(mu[i]-theta.mu[i])*(mu[i]-theta.mu[i]);
				r+=(pi[i]-theta.pi[i])*(pi[i]-theta.pi[i]);
			}
			if (theta.edist > Double.MIN_VALUE) {
				theta.rate  = Math.sqrt(r)/theta.edist;
			} else {
				theta.rate  = 0.0;
			}
			theta.edist = Math.sqrt(r);
		} else {
			theta.edist=0.0;
			theta.rate =0.0;
		}
		
		
	}
	
	/**
	 * Updates the vector of mixture components and returns the new
	 * values.
	 *  
	 * @param x the observed values.
	 * @param sd the standard deviations of the observed values. 
	 * @param theta the current parameter estimates.
	 * @return the new vector of means as an array of double. 
	 */
	public static double[] updateMuVector(double[] x, double[] sd, EMResult theta)
	{
		int i,j;
		double w;
		double[] mu;
		
		mu = new double[theta.k];
		
		for(i=0;i<theta.k;i++) {
		   mu[i]=w=0.0;
		   for(j=0;j<x.length;j++) {
		   	  mu[i] += x[j]*theta.z[i][j]/(sd[j]*sd[j]);
			  w     += theta.z[i][j]/(sd[j]*sd[j]);
		   }    	  
           mu[i] /= w;
		}
		
		return mu;
	}
	/**
	 * Updates the vector of mixture mixings and returns the new
	 * values.
	 * 
	 * @param theta the current parameter estimates.
	 * @return the new vector of mixings as an array of double. 
	 */
	public static double[] updatePiVector(EMResult theta)
	{
	   int i,j;
	   double[] pi = new double[theta.k];
	   
	   for(i=0;i<theta.k;i++) {
	   	  pi[i] = 0.0;
	   	  for(j=0;j<theta.n;j++) {
	   	  	pi[i] += theta.z[i][j];
	   	  }
	   	  pi[i] /= (double) theta.n;
	   }
	   	
	   return pi;
	}
	/**
	 * Updates the Z matrix.
	 * 
	 * @param x the observed values.
	 * @param sd the standard deviations of the observed values. 
	 * @param theta the current parameter estimates.
	 */
	public static void updateZMatrix(double[] x, double[] sd, EMResult theta)
	{
	    int i,j;
	    double[] s;
	    
	    s = new double[theta.n];
	    for(j=0;j<theta.n;j++) s[j]=0.0;
	    
		for(i=0;i<theta.k;i++) {
	    	for(j=0;j<theta.n;j++) {
	      		theta.z[i][j] = 
	      			(theta.pi[i]/sd[j])*NumericalUtilities.gauss((x[j]-theta.mu[i])/sd[j]);
	      		s[j]+=theta.z[i][j];
	   		}
		}
		for(i=0;i<theta.k;i++) {
			for(j=0;j<theta.n;j++) {
	    		if (s[j] > TINY_Z_PROBA) {
		    		theta.z[i][j] /= s[j];
		    		theta.z[i][j] = (theta.z[i][j] < TINY_Z_PROBA) ? 0.0 : theta.z[i][j];
	    		} else {
	    			/* This means that this observations is removed ! */
	    			theta.z[i][j] = 0.0;
	    		}
	      		//System.out.print(" "+theta.z[i][j]+" "+s[j]);
	    	}
			//System.out.println();
		}
	}
	
	/**
	 * Updates the loglikelihood value.
	 * 
	 * @param x the observed values.
	 * @param sd the standard deviations of the observed values. 
	 * @param theta the current parameter estimates.
	 */
	public static void updateLoglikelihood(double[] x, double[] sd, EMResult theta) {
		int i,j;
		double olog,clog,o,u;
		
		olog=clog=0.0;
		
		for(i=0;i<theta.n;i++) {
			for(o=0.0,j=0;j<theta.k;j++) {
				u=NumericalUtilities.gauss((x[i]-theta.mu[j])/sd[i]);
				if (u > Double.MIN_VALUE) {
					o+=(theta.pi[j]/sd[i])*u;
					clog+=theta.z[j][i]*Math.log(u/sd[i]);
				}
			}
			if (o>0.0) {
				olog+=Math.log(o);
			}
		}
		
		theta.olog = olog;
		theta.clog = clog;
		
	}
	/**
	 * Computes the matrix of derivatives of the update
	 * function. This is done by applying one iteration
	 * of the EM algorithm to a sligh modified vector of 
	 * the current parameter estimates.
	 * 
	 * @param x   the data points.
	 * @param sd  the standard deviation of the data points.
	 * @param mu  the previous estimate of the means
	 * @param theta the current parameter estimates.
	 */
	public static void computeDM(double[] x, double[] sd, double[] mu, EMResult theta)
	{
		int i,j;
		double d,z;
		
		for(i=0;i<theta.k;i++) {
			
			EMResult theta_i = new EMResult(theta.n, theta.k);
			
			d = theta.mu[i]-mu[i];
			
			if (Math.abs(d) > Double.MIN_VALUE) {
				
				for (j=0;j<theta.k;j++) {
					if (j==i) theta_i.mu[i]=theta.mu[i];
					else theta_i.mu[j]=mu[j];
					theta_i.pi[j] = theta.pi[j];
				}
				
				/* Do EM step */
				eStep(x,sd,theta_i);
				mStep(x,sd,theta_i,EM_ERR);
				
				for (j=0;j<theta.k;j++) {
					
					z = theta_i.mu[j]-theta.mu[j];
					
					if (Math.abs(z) > Double.MIN_VALUE)
						theta.dm[i][j] = z/d;
					else theta.dm[i][j] = 0.0;
				}
				
			} else {
				for (j=0;j<theta.k;j++)
					theta.dm[i][j]=0.0;
			}
		}
	}
	/**
	 * Computes the variance-covariance matrix of the 
	 * mixture component estimates. If the global
	 * variable {@link EMAlgorithm#DO_SEM} has been
	 * set to false then the observed information won't be
	 * computed and the variance-covariance matrix will be
	 * obtained by using only the complete information.
	 * 
	 * @param sd the standard deviations.
	 * @param theta the maximum likelihood estimate.
	 * 
	 */
	public static void computeCOV(double[] sd, EMResult theta) {
		int i,j,l;
		boolean singular;
		
		if (theta.k < theta.n) {
			for(i=0;i<theta.k;i++) {
		    	theta.ccov[i]=0.0;
		    	for(j=0;j<theta.n;j++) {
		    		theta.ccov[i]+=theta.z[i][j]/(sd[j]*sd[j]);		    		
		    	}
		    	if (theta.ccov[i]!=0.0)
		    		theta.ccov[i]=1.0/theta.ccov[i];
		    	else
		    		theta.ccov[i]=0.0;
			}
		} else {
			for(i=0;i<theta.k;i++)
				theta.ccov[i]=sd[i]*sd[i];
		}
		for(i=0;i<theta.k;i++) {
			   theta.ocov[i][i]=theta.ccov[i];
			   for(j=i+1;j<theta.k;j++)
			      theta.ocov[j][i]=theta.ocov[i][j]=0.0;
			}
		if (theta.k>1 && theta.k<theta.n && DO_SEM) {
			 double wmax;
			 double[] w;
			 double[][] t,u,v;
			 
			 u = new double[theta.k][theta.k];
			 v = new double[theta.k][theta.k];
			 t = new double[theta.k][theta.k];
			 w = new double[theta.k];
			 
			 /*
		      * Thus, we now have DM in mle->dm.
		      * Compute (I-DM) and store it into info->um
		      */
		      for(i=0;i<theta.k;i++) {
	      	    u[i][i]= 1.0-theta.dm[i][i];
	      	    for(j=i+1;j<theta.k;j++) {
	      	 		u[i][j] = -theta.dm[i][j];
		      		u[j][i] = -theta.dm[j][i];
		      	}
		      }
		      /* SVD of (I-DM) */
		      SVDAlgorithm.SVDecomposition(u, v, w, theta.k, theta.k);
		      singular=false;
		      wmax = 0.0;
			  for(i=0;i<theta.k;i++)
				if (w[i] > wmax) wmax=w[i];
			  for(i=0;i<theta.k;i++) {
				if (w[i] < wmax*1.e-8) {
					w[i]=0.0;
					singular=true;
				}
			  }
			  //System.out.println("SINGULAR  " + singular);
			  if (!singular) {
				 /*
				  * Now, we have to compute (I-DM)^{-1}
				  */
				  for(i=0;i<theta.k;i++) {
				  	 for(j=0;j<theta.k;j++) {
				 	 	t[i][j]=0.0;
					 	for(l=0;l<theta.k;l++) {
				 			if (w[l] != 0.0)
				 				t[i][j]+=v[i][l]*u[j][l]/w[l];
				 		}
				 	 }
				  }	
				  /* Compute DV = Ioc^{-1} * DM * inv(I-DM)^{-1} */
				  for(i=0;i<theta.k;i++) {
				  	 for(j=0;j<theta.k;j++) {
				  	 	u[i][j]=0.0;
				  	 	for(l=0;l<theta.k;l++) {
				  	 		u[i][j] += theta.dm[i][l]*t[l][j];
				  	 	}
			  	  	 }
				  }
				  /* Finally V = Ioc^{-1} + DV */
				  for(i=0;i<theta.k;i++) {
				  	 theta.ocov[i][i] = Math.max(theta.ccov[i],theta.ccov[i] + u[i][i]);
				  	 for(j=i+1;j<theta.k;j++)
				  	 	theta.ocov[j][i]=theta.ocov[i][j]=0.5*(u[i][j]+u[j][i]);
				  }
			  } else {
			  	for(i=0;i<theta.k;i++) {
				   theta.ocov[i][i]=theta.ccov[i];
				   for(j=i+1;j<theta.k;j++)
				      theta.ocov[j][i]=theta.ocov[i][j]=0.0;
				}
			  }
			  /*for(i=0;i<theta.k;i++) {
			  	 for(j=0;j<theta.k;j++) {
			  	 	System.out.print(" "+theta.ocov[i][j]);
			  	 }
			  	 System.out.println();
			  }*/
		}
	}


}
