/***************************************************************************
local/competitive learning net, tlong's modification implemented
allows for arbitrary number of expertnets  5/1/93

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

    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;
  float **data,**testdata,**gw0,**dgw0,**gw1,**dgw1;
  float **ew0,**ew1,**ewm;
  float **dew0,**dew1;
  float *error, **yt,*y_out;

  fp1=fopen("train.in","r");
  fp2=fopen("test.in","r");
  fp3=fopen("comp.in","r");
  fp4=fopen("tcomp.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);
  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, we are only doing it for the training data,
 assuming that it is a bigger set than testing data */

  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 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 <= R; i++)
	    error[i]  *=  c[i];
	  for(i = 1;i <= Nnets;i++)
	    for(j = 1;j <= R;j++)
	      yt[i][j] = y[j]-yt[i][j];
	  gback(g,yt,error,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] =  -g[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);

/************************************************************************
*p        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(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,yt,error,x,gx1,gw0,gw1,dgw0,dgw1)
float *g,**yt,*error,*x,*gx1,**gw0,**gw1,**dgw0,**dgw1;
/* backprop of the gating net
*/
{
  extern int NP1,R,Nnets,MP1;
  extern float lrate0,lrate1, alpha;
  int i,j;
  float *delta1,*delta2,**dgds,**w1temp,*ye;

  delta1  =  vector(1,Nnets);
  delta2  =  vector(1,Nnets);
  ye  =  vector(1,Nnets);
  dgds  =  matrix(1,Nnets,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++)
      {
        for(j = 1;j <= Nnets;j++)
          dgds[i][j] = -g[i]*g[j];
        dgds[i][i] = g[i]+dgds[i][i];
      }
  matvec(yt,error,ye,Nnets,R);
  matvec(dgds,ye,delta2,Nnets,Nnets);

    for(i = 1;i <= Nnets;i++)
      for (j = 1;j <= MP1;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(dgds,1,Nnets,1,Nnets);
  free_matrix(w1temp,1,Nnets,1,Nnets);
  free_vector(ye,1,Nnets);
  free_vector(delta2,1,Nnets);
  free_vector(delta1,1,Nnets);
}