#include <math.h>
#include <stdlib.h>
#include <stdio.h>
#include <iostream>
#include "matrix.h"
 
/* MATRIX template library necessary staff  */
#ifndef _NO_NAMESPACE
using namespace std;
using namespace math;
#define STD std
#else
#define STD
#endif

#ifndef _NO_EXCEPTION
#  define TRYBEGIN()	try {
#  define CATCHERROR()	} catch (const STD::exception& e) { \
                     cerr << "Error: " << e.what() << endl; }
#else
#  define TRYBEGIN()
#  define CATCHERROR()
#endif

typedef matrix<double> Matrix;


#define NN 1000					/* max # of the input samples */
#define MAX_X 100
#define DIM 4
#define PI 3.14159265358979
#define DEBUG
#define PREC 0.05				/* threshold */

typedef Matrix TPoint;			/* Tpoint in DIM space when we go farther */

double Co;						/* Constant for gaus */

class Gaussian {				// very simple gaussians with equal variances
 private:
  TPoint _mean;
  double _stdv;
public:
  
  void SetParam (TPoint& mean, double stdv=1) {
	_mean = mean; _stdv = stdv;
  }
  
  Gaussian (TPoint& mean, double stdv) {
	_mean.SetSize(DIM,1);
	SetParam(mean,stdv);
  };
  
  Gaussian() {
	TPoint NUL(DIM,1); NUL.Null();
	SetParam(NUL,10);
  }
  
  TPoint GetMean () {return _mean;}
  double Getstdv () {return _stdv;}
  double Value(TPoint& p) {
	return ((Co*1/pow(_stdv,DIM))*exp(-1*pow((p-_mean).Norm(),2)/(2*pow(_stdv,2))));
  }
};


TPoint x [NN];					/* samples */
int N = 0;						/* # of data */
int Nc = 0;						/* # of clusters */

double * yp;					/* prob. to be in class p after E step */

void generate_data (TPoint x [],	/* data */
					int Nd		/* number of data */
					//	int mean,	/* mean of the distribution */
					//int stdv
					) {	/* varience */
  TPoint mean(DIM);
  int stdv = 5 + rand() % 20;
  cout << "Generating "<< Nd << " points drawn from the gaussian with mean [";
  for (int i=0; i<DIM; i++)	{
	mean(i)=rand()%MAX_X;
	cout << mean(i) << " ";
  }
  cout << "] and std. dev " << stdv << endl;
  
  for (int i=0; i<Nd; i++) {
	x[N+i].SetSize(DIM,1);
	for (int k=0; k<DIM; k++) {
	  double x1=1.0*rand()/RAND_MAX,x2=1.0*rand()/RAND_MAX,
		z=(sqrt(-2*log(x1))*cos(2*PI*x2));
	  /* z - random normal with mean 0 and var =1 */
	  x[N+i](k) = (int) ceil(mean(k) + z*stdv);		/* move and scale */
	}
  }
  N += Nd;
  Nc++;
}
	   
int main () {
  int i,k,was;
  TPoint Unit(DIM,1); for (k=0;k<DIM;k++) Unit(k)=1;
  Co	= pow(1/sqrt(2*PI),DIM);
  cout << "Generating data..." << endl;
  generate_data(x, 4*NN/17);
  generate_data(x, 3*NN/17);
  generate_data(x, 2*NN/17);
  generate_data(x, 4*NN/17);
  generate_data(x, 4*NN/17);

  Gaussian * clusters = new Gaussian [Nc];

  yp    = (double *) calloc ( N*Nc, sizeof(double));
  // initialize distributions
  cout << "Initializing distributions..." << endl;
  for (i=0; i< Nc; i++) {
	TPoint mean(DIM);
	mean = (MAX_X*(i+1)*1.0/(Nc+2))*Unit;
	clusters[i].SetParam( mean,10);
  }
  TRYBEGIN()
	{
	do {
	  was = 0;
	  // E-step
	  for (i=0;i<N;i++) {
		double sum=0;
		for (k=0; k<Nc; k++)  sum += clusters[k].Value(x[i]); 
		for (k=0; k<Nc; k++)  {
		  if (sum!=0) 
			yp [i*Nc+k] = clusters[k].Value(x[i])/sum;	/* P(k|x(i)) */
		  else
			yp [i*Nc+k] = 1/Nc;	/* point is very far from all of
								   them... so let's assume that it is
								   equaly probable from all classes */
		}
	  }
	  
	  // M-step
	  for (k=0; k<Nc; k++) {
		TPoint sum(DIM,1), Pmean(DIM,1);
		double sump=0, sumd=0;
		for (i=0; i<N; i++) {
		  sum += yp[i*Nc+k]*x[i];
		  sump+= yp[i*Nc+k];
		}
		if(fabs((sum/sump - clusters[k].GetMean()).Norm())>PREC) was=1;
		Pmean = sum/sump;
		
		for (i=0; i<N; i++)
		  sumd += yp[i*Nc+k]*pow((x[i]-clusters[k].GetMean()).Norm(),2);
		if ((1-fabs(sqrt(sumd/sump)/clusters[k].Getstdv()))>PREC) was=1;
		
		clusters[k].SetParam(Pmean,sqrt(sumd/(sump*DIM)));
		if (clusters[k].Getstdv()<1) {
		  printf("---!!!!!Tends to have varience -> 0. Artificially broaden!!!---\n");
		  clusters[k].SetParam(Pmean,100);
		}
		
		printf("New distribution #%d  stdv=%3.2f MEAN=\n", k, clusters[k].Getstdv());
		cout << clusters[k].GetMean();
		
	  }
	  printf("\n");
	}  while (was);
	} CATCHERROR();
  delete [] clusters;
  free(yp);
  return 0;
}