/* em07.c
* xref: em05.tex, em06.c
* input:
* output:
* does:  Simulate mixture of 2 normals with equal known varainces.
        Z ~ Bern(theta)
        X0 ~ N(mu_0, 1)
        X1 ~ N(mu_1, 1)
        mu_0 < mu_1
        - MLE using EM algorithm
*/
/***************************************************/
#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <float.h>
#include <string.h>

#define N 100
/****************************************************/
double phi(double t) {return exp(-0.5*t*t);}
/***************************************************/
void swap (double* x, size_t i, size_t j)  /* swap x[i] and x[j] */
{
  double t=x[i]; x[i]=x[j]; x[j]=t;
}
/*************************************************/
void* emalloc (size_t n)
/* allocate, if alloc errors -> issue error msg and exit() */
{
  void *p = malloc(n);
  if (p == NULL) {
    printf("emalloc: Memory allocation error\n");
    exit(1);
  }
  return p;
}
/***************************************************/
double sum (double *x, size_t n, double init_val)
{
  size_t i;
  for (i = 0; i != n; ++i)   init_val += x[i];
  return init_val;
}
/***********************************************************/
double ssq (double *x, size_t n, double init_val, double subtract)
{
  size_t i;
  if (subtract == 0.0) {
    for (i = 0; i != n; ++i)  {
      init_val += x[i] * x[i];
    }
  }
  else {
    for (i = 0; i != n; ++i) {
      double tmp = x[i] - subtract;
      init_val += tmp * tmp;
    }
  }
  return init_val;
}

double mean(double *x, size_t n) {return sum(x, n, 0.0) / n;}
double var(double *x, size_t n)  {return ssq(x, n, 0.0, mean(x, n)) / (n-1);}
/***************************************************/
void  extract(double *mu_0, double *mu_1, double *theta, double *parms)
{
  *mu_0  = parms[0];
  *mu_1  = parms[1];
  *theta = parms[2];
}
/***************************************************/
void parmvec(double mu_0, double mu_1, double theta, double *parms)
{
  parms[0] = mu_0  ;
  parms[1] = mu_1  ;
  parms[2] = theta ;
}
/***************************************************/

void E_step (double* w, double* c0, double* c1, double* y,
               size_t n, double* parms)

  /*  w  <- E[Z  | Y] */
  /*  c0 <- E[X0 | Y] */
  /*  c1 <- E[X1 | Y] */
{
  double mu_0, mu_1, theta, d, s, psi, prior_odds, ww;
  size_t i;

  extract(&mu_0, &mu_1, &theta, parms);
  d = mu_1 - mu_0;
  s = 0.5*(mu_1 + mu_0);
  prior_odds = theta/(1.0-theta);

  for (i = 0; i != n; ++i)  {
    psi   = prior_odds * exp(d*(y[i]-s));
    ww    = w[i] = psi/(1.0+psi);
    c0[i] = ww*mu_0 + (1.0-ww)*y[i];
    c1[i] = ww*y[i] + (1.0-ww)*mu_1;
  }
}
/***************************************************/

void M_step (double* w, double* c0, double* c1,
               size_t n, double* parms)

/* obtain new parameter estimates */
{
  parmvec(mean(c0, n), mean(c1, n), mean(w, n), parms);
  if (parms[0] > parms[1]) {     /* force   mu_0 <= mu_1 */
    parmvec(parms[1], parms[0], 1-parms[2], parms);
  }
}
/***************************************************/

double loglkhd (double *y, size_t n, double *parms)

  /* the log-likelihood  based on y[0:n-1] */
{

  double mu_0, mu_1, theta, s=0;
  size_t i;

  extract(&mu_0, &mu_1, &theta, parms);

  for (i = 0; i != n; ++i)
    s += log(theta*phi(y[i]-mu_1) + (1.0-theta)*phi(y[i]-mu_0));

  return s - 0.5 * n * log(2.0*atan2(0.0, -1.0));
}
/***************************************************/
void  initvals (double *parms, double *y, size_t n)
  /* parms[0:2]<- initial values based on y[0:n-1] */
{
  double  ybar = mean(y, n);
  double  sd   = sqrt(var(y, n));
  parmvec(ybar-sd, ybar+sd, 0.5, parms);
}
/***************************************************/

double monitor (double *y, size_t n, double *parms, size_t iter)

  /* print log-lkhd and estimates */
{
  double ll = loglkhd (y, n, parms);
  printf ("Iteration %4d: l=%12.10f theta=%8.6f mu_0=%8.6f, mu_1=%8.6f\n",
             iter, loglkhd (y, n, parms),parms[2], parms[0], parms[1] );
  return ll;
}
/***************************************************/

int  estimate (double *parms, double *y, size_t n)

/* input: data y[1,...,n] */
/* output: parms[2]<- theta, parms[0]<- mu_0, parms[1]<- mu_1  */
/* returns 1 if converged, 0 otherwise */
{
  const double tol = 1e-9;
  const size_t    maxiter = 200;
  double *w, *c0, *c1, ll, pll;
  size_t i;
  int  conv;

  if (n < 3) return 0;
  w = (double*) emalloc((n+1)*sizeof(*w));
  c0= (double*) emalloc((n+1)*sizeof(*c0));
  c1= (double*) emalloc((n+1)*sizeof(*c1));

  initvals (parms, y, n);       /* initial values */
  pll = monitor (y, n, parms, 0);       /* previous ll */

  for (i=1; i <= maxiter; ++i)  {
    E_step  (w, c0, c1, y, n, parms);
    M_step  (w, c0, c1,    n, parms);
    ll = monitor (y, n, parms, i);
    conv =  (fabs(ll-pll) < tol);       /* converged ? */
    if (conv) break;
    pll = ll;
  }

  free(c1);
  free(c0);
  free(w);
  return conv;
}
/***************************************************/

void summary (double *tru_parms, double *est_parms, double *y, size_t n)

{
  printf ("\nSample size n=%d\n", n);

  printf ("True parameter values: theta=%8.6f, mu_0=%8.6f, mu_1=%8.6f\n",
               tru_parms[2], tru_parms[0], tru_parms[1]);

  printf ("Parameter Estimates:   theta=%8.6f, mu_0=%8.6f, mu_1=%8.6f\n",
               est_parms[2], est_parms[0], est_parms[1]);

  printf ("Log-likelihood = %12.10f\n", loglkhd(y, n, est_parms));

}
/***************************************************/

void  simulate (double *parms, double *y, size_t n)

/* input: parms[2]== theta, parms[0]== mu_0, parms[1]== mu_1  */
/* output: simulated data y[0:n-1] */
{
  const double a[] = { // simulated with parms: 0, 2, 0.3
    0.0007668, -0.939475, 0.8456077, -0.493304, 0.3805769, 2.6011146,
    2.08912, 0.1207658, 2.1957545, -0.959282, 0.1548358, 0.3394777,
    0.4788278, 1.1161578, 0.7300012, -2.120032, -1.085473, 0.2274149,
    0.8690178, -0.932938, 0.4580891, 1.2779732, -1.400311, 3.0479398,
    -1.200237, 2.4490362, 1.1866626, -0.693796, -1.651684, -2.222105,
    0.1401175, 1.9019002, 0.6396009, -1.382346, -0.718541, -0.279965,
    0.4651524, 4.3747447, 1.3001852, -1.359992, 0.2648149, 1.524874,
    1.6180265, 1.7134072, -0.870176, -2.099657, 0.3258149, 1.9970998,
    3.4406849, 0.5807176, -0.93498, 1.2195432, 1.0252208, 2.7270676,
    -1.145333, 0.9615463, -0.595377, 0.320387, 3.0322296, 0.0399386,
    0.6169643, 1.0346161, -1.311096, 0.7927389, -0.37912, 1.2072907,
    0.3468948, 0.5203926, -0.932413, -0.597823, -0.854034, 2.5486596,
    -0.502315, 3.2556189, -0.598408, 0.3365086, 0.4282481, 2.5592742,
    1.9806215, -0.912323, 0.4092242, -0.365883, 0.9473954, 0.4847945,
    0.0288647, 0.4406987, -1.126491, 4.0417427, -1.564552, 2.8244872,
    1.8998997, 3.807505, -0.419282, 2.7611861, 2.0530032, 1.2842716,
    0.8478354, -1.082433, 2.3424769, 0.2944996
  };
  if (sizeof a != n * sizeof y[0])  exit(1); /* error */
  memmove (y, a, sizeof a);
}
/***************************************************/
int main ()
{
  double y[N];
  size_t n = N;

  double tru_parms[3] = {0, 2, 0.3};
  double est_parms[3];

  /* mu_0 := parms[0], mu_1 := parms[1], theta := parms[2] */

  simulate (tru_parms, y, n);        /* simulate data */
  estimate (est_parms, y, n);         /* estimate parameters  */
  summary  (tru_parms, est_parms, y, n);
  return 0;
}