/*
  Copyright (C) 1999 E. H. Haley

  (cwm.c was 1D CWM.  crwth.c was 2D with separable gaussians.
  crwton.c uses the full covariance matrix.)

  trwth is crwton for the triagely image blending.
  
*/

#include <stdlib.h>
#include <math.h>
#include <stdio.h>

#define RTP 2.5066283 //root two pi
//#define NDATA 21
#define NX 21
#define NY 21
#define NCLU 4
#define NITS 100

int main(void)
{
  int i,j, //data plane indices
    m, //cluster #
    d,d2, //dimension 
    it; //iteration
  float num,denom, small=0.01, ploti,plotj, rdcim, term;
  float x[NX][NY][2], y[NX][NY];
  float mux[NCLU][2], sigy[NCLU], f[NX][NY][NCLU];
  float pc[NCLU], norm[NCLU];
  float cov[NCLU][2][2], invcov[NCLU][2][2];
  float px_c[NX][NY][NCLU], py_xc[NX][NY][NCLU], pc_yx[NX][NY][NCLU];
  FILE *fd;
  char fname[32];
  
  //  srandom(2);
  
  for(i=0; i<NX; i++)
    for(j=0; j<NY; j++)
      {
	x[i][j][0]=i-10.0;
	x[i][j][1]=j-10.0;
	y[i][j]=(i-10.0>0 ? 1.0 : 0.0)+(j-10.0>0 ? 1.0 : 0.0);
	//      printf("x[%d]=%.3f, y[%d]=%.3f\n", i,x[i][j], i,y[i][j]);
      }
  
  for(m=0; m<NCLU; m++)
    {
      for(d=0; d<2; d++)
	{
	  mux[m][d]=20.0*((float)random()/((float)RAND_MAX+1.0))-10.0;
	  sigy[m]=1.0; //?
	}
      for(d=0; d<2; d++)
	for(d2=0; d2<2; d2++)
	  {
	    cov[m][d][d2]=100.0*(1-(d+d2)%2); //e.g. {{100, 0}, {0, 100}} //100 bcs ~ sigx^2
	    invcov[m][d][d2]=0.001*(1-(d+d2)%2); //e.g. {{1/100, 0}, {0, 1/100}}
	  }
      pc[m]=1.0/(float)NCLU; //?
      
      // For the image proj, we actually never update f:
      // just initialize each f[i][j][m] to its image[i][j].
      // So for 1st ex, use a constant=m.
      for(i=0; i<NX; i++)
	for(j=0; j<NY; j++)
	  f[i][j][m]=m%2+m/2;
    }
  
  
  for(it=0; it<NITS; it++)
    {
      //update px_c
      for(m=0; m<NCLU; m++)
	{
	  rdcim = sqrt(invcov[m][0][0]*invcov[m][1][1] - invcov[m][1][0]*invcov[m][0][1]);
	  for(i=0; i<NX; i++)
	    for(j=0; j<NY; j++)
	      {
		px_c[i][j][m]= rdcim/RTP/RTP;
		for(d=0; d<2; d++)
		  for(d2=0; d2<2; d2++)
		    px_c[i][j][m]*=
		      exp(-(x[i][j][d]-mux[m][d])*invcov[m][d][d2]*(x[i][j][d2]-mux[m][d2])/2.0);
	      }
	}

      //update py_xc
      for(m=0; m<NCLU; m++)
	for(i=0; i<NX; i++)
	  for(j=0; j<NY; j++)
	    py_xc[i][j][m]=
	      exp(-(y[i][j]-f[i][j][m])*(y[i][j]-f[i][j][m])/(2.0*sigy[m]*sigy[m]))
	      /(RTP*sigy[m]);
      
      //update pc_yx
      for(i=0; i<NX; i++)
	for(j=0; j<NY; j++)
	  {
	    denom=0.0;
	    for(m=0; m<NCLU; m++)
	      {
		pc_yx[i][j][m]=py_xc[i][j][m]*px_c[i][j][m]*pc[m];
		denom+=pc_yx[i][j][m];
	      }
	    for(m=0; m<NCLU; m++)
	      pc_yx[i][j][m]/=denom;
	  }
      
      //update pc
      for(m=0; m<NCLU; m++)
	{
	  pc[m]=0.0;
	  for(i=0; i<NX; i++)
	    for(j=0; j<NY; j++)
	      pc[m]+=pc_yx[i][j][m];
	  pc[m]/=NX*NY;
	}

      //get norm[m] = sum_i pc_yx[i][j]  ..maybe just use f(pc[m])?
      for(m=0; m<NCLU; m++)
	{
	  norm[m]=0.0;
	  for(i=0; i<NX; i++)
	    for(j=0; j<NY; j++)
	      norm[m]+=pc_yx[i][j][m];
	}
	  
      //update mux
      for(d=0; d<2; d++)
	for(m=0; m<NCLU; m++)
	  {
	    mux[m][d]=0.0;
	    for(i=0; i<NX; i++)
	      for(j=0; j<NY; j++)
		mux[m][d]+=x[i][j][d]*pc_yx[i][j][m];
	    mux[m][d]/=norm[m];
	  }
      
      //update cov
      for(m=0; m<NCLU; m++)
	{
	  for(d=0; d<2; d++)
	    for(d2=0; d2<2; d2++)
	      {
		cov[m][d][d2]=0.0;
		for(i=0; i<NX; i++)
		  for(j=0; j<NY; j++)
		    cov[m][d][d2]+=(x[i][j][d]-mux[m][d])*(x[i][j][d2]-mux[m][d2])*pc_yx[i][j][m];
		cov[m][d][d2]/=norm[m];
	      }
	  term=1.0/(cov[m][0][0]*cov[m][1][1] - cov[m][1][0]*cov[m][0][1]);
	  invcov[m][0][0]= term*cov[m][1][1];
	  invcov[m][1][1]= term*cov[m][0][0];
	  invcov[m][1][0]= -term*cov[m][0][1];
	  invcov[m][0][1]= -term*cov[m][1][0];		
	}
      
      //update sigy -- to cluster-expec of <(y-beta_m)^2>_m
      for(m=0; m<NCLU; m++)
	{
	  sigy[m]=0.0;
	  for(i=0; i<NX; i++)
	    for(j=0; j<NY; j++)
	      sigy[m]+=(y[i][j]-f[i][j][m])*(y[i][j]-f[i][j][m])*pc_yx[i][j][m];
	  sigy[m]/=norm[m];
	  sigy[m]=pow(sigy[m],0.5);
	  sigy[m]+=small;
	}
    }
  
  
  for (m=0; m<NCLU; m++)
    printf("f[%d]=%.1f, mux[%d][0]=%.3f, mux[%d][1]=%.3f\ncovariance matrix = \n%f\t%f\n%f\t%f\n",
	   m, f[0][0][m], m, mux[m][0], m, mux[m][1], cov[m][0][0], cov[m][0][1], cov[m][1][0], cov[m][1][1]);
  //cov row<-->col?
  
  sprintf(fname, "crwton%d.dat", NCLU);
  fd = fopen(fname, "w");
  for (ploti=-11.0; ploti<11.0; ploti+=0.5)
    for (plotj=-11.0; plotj<11.0; plotj+=0.5)
      {
	num=denom=0.0;
	for (m=0; m<NCLU; m++)
	  {
	    rdcim = sqrt(invcov[m][0][0]*invcov[m][1][1] - invcov[m][1][0]*invcov[m][0][1]);
	    //maybe already had that, could keep?
	    
	    term =  rdcim * pc[m] /RTP/RTP;
	    
	    for(d=0; d<2; d++)
	      for(d2=0; d2<2; d2++)
		term *= 
		  exp(-(((d==0)?ploti:plotj) -mux[m][d])
		      *invcov[m][d][d2]
		      *(((d2==0)?ploti:plotj) -mux[m][d2])
		      /2.0);
	    num+=term * f[0][0][m];
	    denom+=term;
	  }
	fprintf(fd,"%f\t%f\t%f\n", ploti, plotj, num/denom);
      }
  fclose(fd);	    
  
  exit(0);
}