/* Stefan Posch, University of Halle, 2003/04

   EM for mixture of gaussians (using GSL library)

   - (light) pseudo code version (by B. Moeller) -
*/

/* ========================================================= */
/* Update step                                               */
/* ========================================================= */

/* Estimation step:
   Assign to each data element the corresponding component or
   calculate the weight, respectively. 
*/

void em_mg2d_e_step( Tem_mg2d *em,          // data structure for EM-component
                     Tmixgaussian2d *paras, // current parameters of gaussians (paras->gaussians[i] -> component i)
                     Tc2dsamples *cdata,    // data points: cdata[i][0] -> x, cdata[i][1] -> y
                     double **gamma,        // weights to be calculated: gamma[i][k] ->
                                            //         weight of k-th component for data point i
                     int *u_t)              // index of component with maximal weight for each data point
{
    double       *eta;
    gsl_vector   *x;
    int          i, g, k, max_k;
    double       sum_eta, cum_prob, max_value, r;

    // paras->K ==> number of components
    eta = ALLOC1( double, paras->K, "eta");
    x = gsl_vector_alloc(2);

    em->likelihood = 0;

    /* ============================  */
    // iterate over all data points
    for ( i=0 ; i<cdata->N ; i++) {
        sum_eta = 0.0;

	/* ==============================================================  */
	// for each datapoint, calculate probabilities for all components...
        for ( g=0; g<paras->G ; g++ ) {
            GSL_vector_set( x, 0, cdata->datapoints[i][0]);
            GSL_vector_set( x, 1, cdata->datapoints[i][1]);

            // weight for component g, given data point i
            eta[g] = paras->p[g] *
                        gaussian2d_pdf ( x, paras->gaussians[g]);
            sum_eta += eta[g];
        }

	// ... and scale them to one
        // (cope with very small eta's summing to zero
        //  for reasons of numerical stability)
        if ( sum_eta == 0 ) {
            for ( k=0; k<paras->K ; k++ ) {
                eta[k] = 1.0/paras->K;
                gamma[i][k] = 1.0/paras->K;
            }
        } else {
            for ( k=0; k<paras->K ; k++ ) {
                gamma[i][k] = eta[k] / sum_eta;
            }
        }

	/* =================  */
	// update likelihood
        em->likelihood += log( sum_eta);
        
	/* ====================== */
        // standard EM or K-means
        switch ( em->em_type ) {
        case EM_STANDARD:
            // all is ready, weights have been calculated ....
            break;

        case EM_K_MEANS:
	    // search maximum to find component 
            max_value = -DBL_MAX; max_k = 0;
            for ( k=0; k<paras->K ; k++ ) {
                if ( eta[k] > max_value ) {
                    max_value = eta[k];
                    max_k = k;
                }
            }

            u_t[i] = max_k;
            break;

        }
    }
}


/* ========================================================= */

/* Maximization step:
   Calculate new parameters.
*/
void em_mg2d_m_step( Tmixgaussian2d *paras,    // current parameters of gaussians
                     Tc2dsamples *cdata,       // data points
                     double **gamma)           // component-wise weights
{

    int        i, g, k;
    double     s11, s22, s12, mean_x, mean_y, x_norm, y_norm;
    gsl_matrix *SSD = gsl_matrix_alloc(2,2);
    double     *sum_gamma_k;


    /* ====================================  */
    // update probabilities for each component 
    for ( k=0; k<paras->K ; k++ ) {
        sum_gamma_k[k] = 0;
        for ( i=0 ; i<cdata->N ; i++) 
            sum_gamma_k[k] += gamma[i][k];

        paras->p[k] = sum_gamma_k[k] / cdata->N;
    }

    /* ====================================  */
    // update mean of gaussians
    for ( g=0 ; g<paras->G ; g++ ) {
        mean_x = mean_y = 0;
        for ( i=0 ; i<cdata->N ; i++) {
            mean_x += gamma[i][g] * cdata->datapoints[i][0];
            mean_y += gamma[i][g] * cdata->datapoints[i][1];
        }
        GSL_vector_set( paras->gaussians[g]->mu, 0, mean_x/sum_gamma_k[g]);
        GSL_vector_set( paras->gaussians[g]->mu, 1, mean_y/sum_gamma_k[g]);
    }

    /* ====================================  */
    // update covariance matrix of gaussians 
    for ( g=0 ; g<paras->G ; g++ ) {
        s11 = s22 = s12= 0;
        for ( i=0 ; i<cdata->N ; i++) {
            x_norm = cdata->datapoints[i][0] - 
                     GSL_vector_get( paras->gaussians[g]->mu, 0);
            y_norm = cdata->datapoints[i][1] - 
                     GSL_vector_get( paras->gaussians[g]->mu, 1);

            s11 += gamma[i][g] * x_norm * x_norm;
            s22 += gamma[i][g] * y_norm * y_norm;
            s12 += gamma[i][g] * x_norm * y_norm;
        }

        GSL_matrix_set( SSD, 0, 0, s11 / sum_gamma_k[g]);
        GSL_matrix_set( SSD, 0, 1, s12 / sum_gamma_k[g]);
        GSL_matrix_set( SSD, 1, 0, s12 / sum_gamma_k[g]);
        GSL_matrix_set( SSD, 1, 1, s22 / sum_gamma_k[g]);

        gaussian2d_setcov( paras->gaussians[g], paras->gaussians[g]->Cov);
    }
}

/* ========================================================= */

/* density of x for given gaussian                         */
double gaussian2d_pdf( gsl_vector *x, Tgaussian2d *g) {
    return exp( gaussian2d_ln_pdf( x, g));

}

/* ======================================================= */

/* ln density of x for given gaussian                      */
double gaussian2d_ln_pdf( gsl_vector *x, Tgaussian2d *g) {
    double res, exponent;
    static gsl_vector *x_norm = NULL;
    static gsl_vector *x_help = NULL;

    if ( x_norm == NULL ) {
        x_norm = gsl_vector_alloc(2);
        x_help = gsl_vector_alloc(2);
    }

    // normalize x (i.e. subtrat mu)
    gsl_vector_memcpy( x_norm, x);
    gsl_vector_sub( x_norm, g->mu);

    gsl_blas_dgemv( CblasNoTrans, 1, g->CovInv, x_norm, 0, x_help);
    exponent = 0;
    gsl_blas_ddot( x_norm, x_help, &exponent);
    exponent *= -0.5;

    res = log( 1 / (2 * M_PI * sqrt(g->CovDet))) +  exponent;

    return res;
}