/*
 * Example from section 1 of "Sequence Estimation EM Examples",
 * R. Perry, June 1998
 */


SNR = p = c = N = v = b = y = b_save = c_save = 0;


void get_y( double _SNR, double _p, double _c, double _N)
{
  SNR = (narg > 0) ? _SNR : 6.0;	/* signal-to-noise ratio */
  p   = (narg > 1) ? _p   : 0.4;	/* P(b==1) */
  c   = (narg > 2) ? _c   : 2.0;	/* scale factor */
  N   = (narg > 3) ? _N   : 50;		/* number of samples */

  b = (rand(1,N) < p);			/* 0's and 1's */

  avg_b = sum(b) / N;

  var_b = (sum(b.*b) - N * avg_b * avg_b) / (N-1);

  v = (c*c*var_b)/(10^(SNR/10));	/* noise variance */

  y = c*b + sqrt(v)*rand(1,N,"gauss");	/* observed data */

  b_save = b;
  c_save = c;

  b = ones( 1, N);	/* initial estimate */
}


void iter( int _n)
{
/*
 * perform n iterations of EM for y = c*b + z problem
 *
 * Given values for y, b:
 * E-step: estimate value for c
 * M-step: compute ML estimate of b
 */

  int i, n = 1;
  vector j;

  n = (narg > 0) ? _n : 1;

  for( i = 0; i < n; ++i) {

    /* E-step */

	j = find( b == 1);
	c = sum(y(j)) / cols(j);

    /* M-step */

        b = (c > 0) ? (y > c/2) : (y < c/2);

    printf( "c = %g, sum(b) = %g\n", c, sum(b));
  }
}


void check_ml( void)
{
/*
 * Verify that ML solution is at least a local minimum of
 * -log(f(Y|B,c)) by examining nearby points:
 *   with individual bits flipped one at a time
 *   or using 0.999*c or 1.001*c for c.
 */

  int i;
  double s, s0;
  vector t, ct;

  t = (y - c*b);
  s0 = sum(t.*t);

  printf( "check_ml: min = %g\n", s0);	/* supposed minimum */

  ct = [0.999*c, 1.001*c];

  for( i = 1; i <= 2; ++i) {

    t = (y - ct(i)*b);
    s = sum(t.*t);

    if( s < s0)
      printf("check_ml: found better solution %g with c = %g\n", s, ct(i));
  }

  for( i = 1; i <= N; ++i) {

    b(i) = !b(i);
    t = (y - c*b);
    s = sum(t.*t);
    b(i) = !b(i);

    if( s < s0)
      printf("check_ml: found better solution %g with bit %g flipped\n", s, i);
  }

}
