/*
 * Example from section 5 of "Dual Binomial EM Examples",
 * R. Perry, June 1998
 */

/*** need to load fmin before using p_ml() from here ***/

vector G(vector _n, double _m, double _v)
{
/* Gaussian probability density function, mean m, variance v */

  vector t = _n-_m;
  double pi = 4 * atan(1);

  return exp( (t.*t)/(-2*_v)) / sqrt(2 * pi * _v);
}


vector f( vector _y, double _p, double _v)
{
/* probability density function for Bernoulli + Gaussian */

  return _p * G(_y,1,_v) + (1 - _p) * G(_y,0,_v);
}


void get_y( double _v, double _p, double _N)
{
  v = (narg > 0) ? _v : 0.05;	/* noise variance */
  p = (narg > 1) ? _p : 0.4;	/* P(1) */
  N = (narg > 2) ? _N : 50;	/* number of samples */

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

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

  avg_y = sum(y) / N;

  var_y = (sum(y.*y) - N * avg_y * avg_y) / (N-1);

  p_save = p;
  x_save = x;

  p = 0.5;	/* initial estimate */
}


void iter( int _n)
{
/*
 * perform n iterations of EM for Bernoulli + noise problem
 *
 * Given values for y, v, p:
 * E-step: estimate values for x
 * M-step: compute ML estimate of p
 */

  int i, n = 1;

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

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

    /* E-step */

      ex = p * G(y,1,v) / f(y,p,v);

      n1 = sum(ex);

    /* M-step */

      p = n1 / N;

      print(p);
  }
}


double fn( double _p)
{
/*
 * returns  -log( f(y(1)) * f(y(2)) * ... * f(y(N)))
 */

  return -sum(log(f(y,_p,v)));
}


double fmin();	/* need to load fmin before using p_ml() from here */


double p_ml( void)
{
/*
 * directly compute ML estimate of p
 */

  int i;
  vector j;		/* indices of minimum points */
  int j1, j2;		/* indices bounding the interval of the minimum  points */
  int NG = 1001;	/* number of points on grid */

/* first check for minimum fn() on a grid */

  P = [0:(NG-1)] / (NG-1);
  FN = zeros(1,NG);

  for( i = 1; i <= NG; ++i)
    FN(i) = fn(P(i));

  j = find( FN == min(FN));

  if( cols(j) == 1) {	/* single minimum point found on grid */

    if( j == 1) {		/* minimum is near p=0 */
      j1 = 1;
      j2 = 2;
    } else if( j == NG) {	/* minimum is near p=1 */
      j1 = NG - 1;
      j2 = NG;
    } else {			/* minumum is somewhere in the middle */
      j1 = j - 1;
      j2 = j + 1;
    }
    
  } else {

    print( "p_ml: Warning: multiple minimum points found on grid");

    j1 = j(1);			/* interval end points */
    j2 = j(cols(j));

    if( j1 > 1)			/* not at p=0 */
      --j1;

    if( j2 < NG)		/* not at p=1 */
      ++j2;
  }

/* now use interval halving to refine the estimate of the minimum */

  return fmin( "fn", P(j1), P(j2), 1e-12);
}
