/***************************************************************************
local/competitive learning net, Jordan's algorithm
allows for arbitrary number of expertnets

If we use momentum term, must add another matrix to save all the
delta weights.  currently we do not use momentum - alpha is set to zero

    Nnets:  number of expert nets
    N:  input vector length
    R:  output vector length
    P:  number of nodes in 1st hidden layer
    x:  input, dimension N
    y:  desired output, dimension R
    x1: 1st hidden layer output of expert net, dimension p+1 
    y1: output of expert net, 
    x1m: stores x1 of all expert nets
    yt: stores y1 of all expert nets

    gw0: weight matrix of the gating net from the 0th layer (input
         layer), dimension Nnetsx(N+1)
    gw1: weight matrix of the gating net from the 1st layer (hidden
         layer), dimension Rx(Nnets+1)
    dgw0: delta weight matrix of gw0
    dgw1: delta weight matrix of gw1
    ew0: weight matrix of the expert net from the 0th layer(input
         layer) dimension Px(N+1)
    ew1: weight matrix of the expert net from the 1st layer(hidden
         layer) dimension Rx(P+1)
    dew0: delta weight matrix of ew0
    dew1: delta weight matrix of ew1

    g:  gating net output, dimension Nnets
    h:  posterior probability
    ewm: stores the weights of all the expert nets
    delta: backprop term of the output layer for the expert net
    c: the diagonal of the inverse covariance matrix for output,
        we assume they are diagonal only, so c is just a vector
    ntrain: number of input data sets
    ntest: number of input data sets
    wflag: if wflag =1, we want to read and write the weight matrices
           if wflag =2, we want to read but not write the weight matrice
           if wflag =0, we do not read but will write the weight matrice

***************************************************************************/
#include <math.h>
#include "nrutil.h"
#include <stdio.h>
#include <stdlib.h>

#define sig(x) x/(1.0 + fabs(x))
#define dsig(x) (1.0 - fabs(x))* (1.0 - fabs(x))

int Nnets,MP1,N,R,P,NP1,PP1,NR;
float lrate0, alpha,lrate1;	
void matvec();
void vecmat();
void gfeedf();
void gback();
void feedf1();
void backprop1();
void randomm();

main()
{

  FILE *fp1,*fp2,*fp3,*fp4,*fp5,*fp6;
  int nepoch,ntrain,ntest,i,ii,jj,j,k,kk,wflag;
  int PxNP1,n_w;
  long idum;
  float inner,tss,tss_last,ysum,ss,*delta;
  float *x,*y,*gx1,**x1m,*x1,*y1;
  float *g,*c,*h;
  float **data,**testdata,**gw0,**dgw0,**gw1,**dgw1;
  float **ew0,**ew1,**ewm;
  float **dew0,**dew1;
  float *error, **yt,*y_out,denom;

  fp1=fopen("train.in","r");
  fp2=fopen("test.in","r");
  fp3=fopen("comp.in","r");
  fp4=fopen("jcomp.out","w");
  fp5 = fopen("weight.i","r");
  fp6 = fopen("weight.s","w");

  fscanf(fp3,"%d %d %d %f %f %f",&nepoch,&ntrain,&ntest,&lrate0,&lrate1,&alpha);
  fscanf(fp3,"%d %d %d %d ",&Nnets,&N,&R,&P);
  fscanf(fp3,"%d %ld",&wflag,&idum);
  fclose(fp3);

  NP1= N+1;
  PP1= P+1;
  PxNP1 = NP1*P;
  n_w = PxNP1 + PP1*R;
  MP1= Nnets+1;
  NR= N+R;
  tss_last = 2.0;

  data = matrix(1,ntrain,1,NR);       
  testdata = matrix(1,ntest,1,NR);       
  x = vector(1,NP1);
  gx1 = vector(1,MP1);
  y = vector(1,R);
  y1 = vector(1,R);
  x1 = vector(1,PP1);
  x1m = matrix(1,Nnets,1,P);
  error = vector(1,R);
  g = vector(1,Nnets);
  h = vector(1,Nnets);
  delta = vector(1,R);
  c = vector(1,R);
  gw0 = matrix(1,Nnets,1,NP1);
  dgw0 = matrix(1,Nnets,1,NP1);
  gw1 = matrix(1,Nnets,1,MP1);
  dgw1 = matrix(1,Nnets,1,MP1);
  ew0 = matrix(1,P,1,NP1);
  dew0 = matrix(1,P,1,NP1);
  ew1 = matrix(1,R,1,PP1);
  dew1 = matrix(1,R,1,PP1);
  ewm = matrix(1,Nnets,1,n_w);
  yt  =  matrix(1,Nnets,1,R);
  y_out  =  vector(1,R);

/*  if wflag !=0, we want to read the weight matrices from file  */
  if(wflag != 0)
    {
      for(i = 1; i<= Nnets; i++)
	for(j = 1; j <= NP1; j++)
	    fscanf(fp5,"%f",&gw0[i][j]);
      for(i = 1; i<= Nnets; i++)
	for(j = 1; j <= MP1; j++)
	    fscanf(fp5,"%f",&gw1[i][j]);
      for(i = 1; i<= Nnets; i++)
	for(j = 1; j <= n_w; j++)
	    fscanf(fp5,"%f",&ewm[i][j]);
    }
  else
    {
/*   initialize weight matrices */
      randomm(gw0,Nnets,NP1,&idum);
      randomm(gw1,Nnets,MP1,&idum);
      randomm(ewm,Nnets,n_w,&idum);
    }

/* zero delta weight matrices  */
  for(i = 1;i <= Nnets;i++)
    for(j = 1;j <= NP1;j++)
      dgw0[i][j] = 0.0;
  for(i = 1;i <= Nnets;i++)
    for(j = 1;j <= MP1;j++)
      dgw1[i][j] = 0.0;

  for(i = 1;i <= P;i++)
    for(j = 1;j <= NP1;j++)
      dew0[i][j] = 0.0;
  for(i = 1;i <= R;i++)
    for(j = 1;j <= PP1;j++)
      dew1[i][j] = 0.0;

/* enter the constant x value for layers 0 and 1.  They will
be used to generate bias weights    */
  x[NP1] = 1.0;
  gx1[MP1] = 1.0;
  x1[PP1] = 1.0;

/* right now all the train and test data are in train.in, if test data
is in a separate test.in, change fscanf for test data below  */

  for (k = 1;k <= ntrain; k++)	
    for (i = 1;i <= NR;i++)
      fscanf(fp1,"%f",&data[k][i]);
  for (k = 1;k <= ntest; k++)	
    for (i = 1;i <= NR;i++)
      fscanf(fp1,"%f",&testdata[k][i]);
/*  compute the normalizing factor ysum for tss, and the covariance c
 for each element of y  */

  for(j = NP1;j <= NR;j++)
    {
      ysum = 0.0;
      for(k = 1;k <= ntrain;k++)
	ysum += data[k][j]*data[k][j];
      c[j-N] = ysum/ntrain;
    }
/* now get ysum, then invert c, if the data is already normalized
   we can be lazy and set c to 1  */
  ysum  =  0.0;
  for(j = 1;j <= R;j++)
    {
      ysum += c[j];
      c[j] = 1.0;
    }
  fclose(fp1);
  fclose(fp2);
  fclose(fp5);

  for(kk = 1; kk <= nepoch; kk++)
    {
/*  zero tss  */

      tss = 0.0;

      for (k = 1; k <= ntrain; k++)	
	{
	  for (i = 1;i <= N;i++)
	    x[i] = data[k][i];

	  for (i = 1;i <= R;i++)
	    y[i] = data[k][i+N];
	
/*  feedforward gating net  */

	  gfeedf(gw0,gw1,x,gx1,g);

/*  feedforward expert net  */
	  for(ii=1; ii<=Nnets; ii++)
	    {
/* copy weights from the big multi-net matrix ewm  */	      
	      for(i=1; i<= P; i++)
		for(j=1; j<=NP1; j++)
		  {
		    jj = (i-1)*NP1 + j;
		    ew0[i][j] = ewm[ii][jj];
		  }
	      for(i=1; i<= R; i++)
		for(j=1; j<= PP1; j++)
		  {
		    jj = PxNP1 +(i-1)*PP1 + j;
		    ew1[i][j] = ewm[ii][jj];
		  }

	      feedf1(ew0,ew1,x,y1,x1);

/* save the  output x1,y1-- yt stores the transpose of expert net outputs  */
	      for(j = 1;j <= P;j++)
		x1m[ii][j]  =  x1[j];

	      for(j = 1;j <= R; j++)
		yt[ii][j]  =  y1[j];
            }

/*  compute error(error) of final network output( y_out)  */
	  vecmat(g,yt,y_out,Nnets,R);
	  inner = 0.0;
	  for (i = 1;i <= R;i++)
	    {
	      error[i]  =  y[i] - y_out[i];
	      inner  +=  error[i]*error[i];
	    }
/*  compute the summed square error of the final output  */
 
	  tss += inner;

/*  backprop the gating net, note that yt is now to store the error
 of the expert net output, each expert net takes up 1 row.
this time, need to redo the error term to include the covariance  */

	  for(i = 1;i <= Nnets;i++)
	    for(j = 1;j <= R;j++)
	      yt[i][j] = y[j]-yt[i][j];


          denom = 0.0;
 
          for (ii = 1;ii <= Nnets;ii++)
            {
	      inner = 0.0;
	      for(i = 1; i<= R; i++)
		inner += yt[ii][i];
              h[ii] = g[ii]*exp(-inner/2.0);
              denom  += h[ii];
            }
 
          for (ii = 1;ii <= Nnets;ii++)
            h[i] /= denom;
 
/*  backprop the gating net  */
          gback(g,h,x,gx1,gw0,gw1,dgw0,dgw1);
 
 
/*  backprop the expert net, copy the weight matrix from ewm again,
    also, copy x1 from x1m   */

	  for (ii = 1;ii <= Nnets;ii++)
	    {
	      for (i = 1;i <= R;i++)
		delta[i] =  -h[ii]*c[i]*yt[ii][i];
	      
	      for(i=1; i<= P; i++)
		for(j=1; j<=NP1; j++)
		  {
		    jj = (i-1)*NP1 + j;
		    ew0[i][j] = ewm[ii][jj];
		  }
	      for(i=1; i<= R; i++)
		for(j=1; j<= PP1; j++)
		  {
		    jj = PxNP1 +(i-1)*PP1 + j;
		    ew1[i][j] = ewm[ii][jj];
		  }
	      
	      for(i=1; i<= P; i++)
		x1[i] = x1m[ii][i];

	      backprop1(x,x1,delta,dew0,dew1,ew0,ew1);

/* store the updated weight matrices in ewm  */
	      for(i=1; i<= P; i++)
		for(j=1; j<=NP1; j++)
		  {
		    jj = (i-1)*NP1 + j;
		    ewm[ii][jj] =  ew0[i][j];
		  }
	      for(i=1; i<= R; i++)
		for(j=1; j<= PP1; j++)
		  {
		    jj = PxNP1 +(i-1)*PP1 + j;
		    ewm[ii][jj] = ew1[i][j];
		  }
	    }
	}
      tss  /= ntrain;
      tss  =  sqrt(tss/ysum);

      printf("%d  %f ",kk,tss);  
      fprintf(fp4,"%d  %f  ",kk,tss);

/************************************************************************
        CROSS VALIDATION
*************************************************************************/
/*  zero tss  */
      tss = 0.0;

      for (k = 1;k <= ntest;k++)	
	{
	  for (i = 1;i <= N;i++)
	    x[i] = testdata[k][i];
	  for (i = 1;i <= R;i++)
	    y[i] = testdata[k][i+N];
/*  feedforward gating net  */
	  gfeedf(gw0,gw1,x,gx1,g);
/*  feedforward expert net, copy the weight matrices from ewm  */
	  for(ii=1; ii<=Nnets; ii++)
	    {
	      
	      for(i=1; i<= P; i++)
		for(j=1; j<=NP1; j++)
		  {
		    jj = (i-1)*NP1 + j;
		    ew0[i][j] = ewm[ii][jj];
		  }
	      for(i=1; i<= R; i++)
		for(j=1; j<= PP1; j++)
		  {
		    jj = PxNP1 +(i-1)*PP1 + j;
		    ew1[i][j] = ewm[ii][jj];
		  }
	      feedf1(ew0,ew1,x,y1,x1);

/*  compute error of each expert net output  */

	      for(j = 1;j <= R;j++)
		yt[ii][j]  =  y1[j];
            }
/*  compute error of final network output -- y_out  */

          vecmat(g,yt,y_out,Nnets,R);
          inner = 0.0;
          for (i = 1;i <= R;i++)
            {
              error[i]  =  y[i] - y_out[i];
              inner  +=  error[i]*error[i];
            }

	  tss +=  inner;
/*	  
	  fprintf(fp4,"%f ",x[1]);
	  for(i=1;i<=Nnets;i++)
	    fprintf(fp4,"%f ",g[i]);
	  fprintf(fp4,"\n ");
*/	  
	}
      tss  =  sqrt(tss/ntest/ysum);
      
      fprintf(fp4," %f \n",tss);
      printf(" %f \n",tss);
      if(tss > tss_last)
	{
	  lrate0 *= 0.9;
	  lrate1 *= 0.9;
	}
      tss_last = tss;
    }
  printf("lrates = %f %f\n",lrate0,lrate1);
/*  if wflag != 2, we want to save the weight matrices into file  */
  if(wflag != 2)
    {
      for(i = 1; i<= Nnets; i++)
	{
	  for(j = 1; j <= NP1; j++)
	    fprintf(fp6," %f ",gw0[i][j]);
	  fprintf(fp6,"\n");
	}
      for(i = 1; i<= Nnets; i++)
	{
	  for(j = 1; j <= MP1; j++)
	    fprintf(fp6," %f ",gw1[i][j]);
	  fprintf(fp6,"\n");
	}
      for(i = 1; i<= Nnets; i++)
	{
	  for(j = 1; j <= n_w; j++)
	    fprintf(fp6," %f ",ewm[i][j]);
	  fprintf(fp6,"\n");
	}
    }

  free_vector(x,1,NP1);
  free_vector(gx1,1,MP1);
  free_vector(y,1,R);
  free_vector(y1,1,R);
  free_vector(x1,1,PP1);
  free_vector(error,1,R);
  free_vector(g,1,Nnets);
  free_vector(h,1,Nnets);
  free_vector(delta,1,R);
  free_matrix(gw0,1,Nnets,1,NP1);
  free_matrix(dgw0,1,Nnets,1,NP1);
  free_matrix(gw1,1,Nnets,1,MP1);
  free_matrix(dgw1,1,Nnets,1,MP1);
  free_matrix(ew0,1,P,1,NP1);
  free_matrix(dew0,1,P,1,NP1);
  free_matrix(ew1,1,R,1,PP1);
  free_matrix(dew1,1,R,1,PP1);
  free_matrix(ewm,1,Nnets,1,n_w);
  free_matrix(data,1,ntrain,1,NR);      
  free_matrix(testdata,1,ntest,1,NR);      
  free_matrix(yt,1,Nnets,1,R);
  free_matrix(x1m,1,Nnets,1,P);
  free_vector(y_out,1,R);

  fclose(fp4);
  fclose(fp6);
}

/*~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~*/

void randomm(w,mk,nk,idum)
float **w;
long *idum;
int mk,nk;
{
  float ran2();
  int i,j;
  for(i = 1;i <= mk;i++)
    for(j = 1;j <= nk;j++)
      w[i][j]  =  (ran2(idum)-0.5)*4.0;
}

/*~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~*/
void gfeedf(w0,w1,x,x1,y)
float **w0,**w1,*x,*x1,*y;
/* feedforward for gating net
w: weight matrix
x: input
y: output
*/
{
  int i,j;
  extern int Nnets,NP1,MP1;
  float denom,*s;
  s = vector(1,Nnets);

/* input layer  */
  matvec(w0,x,s,Nnets,NP1);
/*  hidden layer.
 has the same number of nodes as the output layer plus a bias node*/
  for(i = 1;i <= Nnets;i++)
    x1[i] = sig(s[i]);
  matvec(w1,x1,s,Nnets,MP1);
/* output layer -- softmax  */
  denom = 0.0;
  for(i = 1;i <= Nnets;i++)
    {
      y[i] = exp(s[i]);
      denom +=  y[i];
    }
  for(i = 1;i <= Nnets;i++)
    y[i] /= denom;
  free_vector(s,1,Nnets);
}
/*  *********************************** */
void gback(g,h,x,gx1,gw0,gw1,dgw0,dgw1)
float *g,*h,*x,*gx1,**gw0,**gw1,**dgw0,**dgw1;
/* backprop of the gating net, jordan's algorithm
*/
 
{
  int i,j;
  float *delta1,*delta2,**w1temp;

  delta1 = vector(1,Nnets);
  delta2 = vector(1,Nnets);
  w1temp = matrix(1,Nnets,1,Nnets);
 
  for (i = 1;i <= Nnets;i++)
    for (j = 1;j <= Nnets;j++)
      w1temp[i][j] = gw1[i][j];
 
  for (i = 1;i <= Nnets;i++)
    {
      delta2[i] = h[i] -g[i];
      for (j = 1;j <= NP1;j++)
        {
          dgw1[i][j] = alpha*dgw1[i][j]+lrate1*delta2[i]*gx1[j];
          gw1[i][j] += dgw1[i][j];
        }
    }
  vecmat(delta2,w1temp,delta1,Nnets,Nnets);
  for (i = 1;i <= Nnets;i++)
     delta1[i] *=  dsig(gx1[i]);
 
  for (i = 1;i <= Nnets;i++)
    for (j = 1;j <= NP1;j++)
      {
        dgw0[i][j] = alpha*dgw0[i][j]-lrate1*delta1[i]*x[j];
        gw0[i][j] += dgw0[i][j];
      }
 
 
  free_matrix(w1temp,1,Nnets,1,Nnets);
  free_vector(delta2,1,Nnets);
  free_vector(delta1,1,Nnets);
}