/*****************************************************************************/
/*  Program mackey.c
       tree-structured nonlinear approximation for predicting the
       Mackey-Glass equation.

   this version uses Sutton's polynomial modifications
*/

#include <stdio.h>
#include <X11/Xlib.h>
#include <X11/X.h>
#include <X11/Xutil.h>
#include <strings.h>
#include <math.h>
#include <ctype.h>
#include "poly-tree.h"

#define LEARNING 400
#define ESTIMATION 200
#define N_EXAMPLES 1000000
#define pixel_range 255.0
#define TWO_PI 2.0*3.14159265456
#define N_INPUTS 6
#define RESOLUTION 1
#define TABLE_SIZE N_INPUTS*RESOLUTION
#define MULT_NODES 5  /* number added at a time */
#define BUF_SIZE 800
#define WIN_HEIGHT 200
#define ERROR_SAMPLES 100
#define MAX_PREDICT 100    /* MAX_PREDICT * DELTA = PREDICT_LENGTH */
                           /* N.B. MAX_PREDICT * DELTA + ERROR_SAMPLES
			      + N_INPUTS * DELTA must be < BUF_SIZE */

#define TAU 30
#define DELTA 6
#define P DELTA            /* calculate_error() assumes this */
#define TIME_STEP 1.0

#define INITIAL_RATE 0.005
#define DELTA_RATE 0.5
#define MAX_LEVELS 10
#define INIT_ALPHA 0.01

NODE *root_node;

float ring_buffer[BUF_SIZE];
float approx_buffer[BUF_SIZE];
float error_buffer[MAX_PREDICT];
float predict_buffer[MAX_PREDICT];

/* fourth-order Runge-Kutta integration */
float k1[BUF_SIZE], k2[BUF_SIZE], k3[BUF_SIZE], k4[BUF_SIZE];
float h = TIME_STEP;

int ring_pos = 0;
int skip_show = 1;
int show_results = 0;
int subtrees;
FILE *logfile;
float min_value = 0.2;
float max_value = 1.4;
float data_spread = 1.2;
float rate;
float alpha;
float basis_values[N_INPUTS];
float basis_potentials[N_INPUTS];
int max_levels = MAX_LEVELS;
FILE *fp;
float dat_variance, dat_mean;


double drand48();
char *malloc();
NODE *find_max_var();
float square();
long time();
FILE *fopen();

extern forward_propagate(), backward_propagate(), add_level(), zero_var();
extern NODE *nalloc();
extern graph_value(), make_graph_window();
extern double fabs();

/*****************************************************************************/

main()
{
  initialize();

  learn();

  done();
}


/*****************************************************************************/

initialize()
{
  make_graph_window(BUF_SIZE, WIN_HEIGHT);
  init_learning();

}

/*****************************************************************************/

done()
{
  fclose(fp);
  done_learning();
  printf( "press RETURN to quit...\n" );
  getchar();
}

/*****************************************************************************/

init_learning()
{
  int b,n,x,y,i;
  NODE *new;

  srand48(time(0));

  /* initial conditions */
  for (i=0 ; i<BUF_SIZE ; i++) ring_buffer[i] = 0.8;

  ring_pos = 0;
  rate = INITIAL_RATE;
  alpha = INIT_ALPHA;
  
  /* create root node */
  root_node = nalloc(1);
  root_node->below = NULL;
  root_node->next = NULL;
  root_node->output = 0;
  root_node->level = 0;
  root_node->number = 0;
  root_node->var = 0.0;
  root_node->input_var = 1.0;
  root_node->var_output = 0.0;
  root_node->var_weight = 0.0;
  root_node->weight = 0.0;
  root_node->flag = 1.0;  /* this will get filled */
  root_node->value = 1.0;  /* always "on" */

  new = root_node->below = nalloc(1);

  /* make all first-level nodes */
  for (i=0 ; i<TABLE_SIZE ; i++)
    {
      new->output = 0.0;  /* must be zero */
      new->weight = 0.0;  /* could be random */
      new->var = 0.0;
      new->input_var = 0.0;
      new->var_output = 0.0;
      new->var_weight = 0.0;
      new->value = 0.0;
      new->level = 1;
      new->number = i;
      new->below = NULL;  /* so we know it's a leaf */
      new->flag = 0;  /* allowed to add subtrees */
      if (i < TABLE_SIZE-1) new->next = nalloc(1);
      else new->next = NULL;

      new = new->next;
    }
}


/*****************************************************************************/
learn()
{
  int example, i, j;
  float network_output;
  float rms_error;
  NODE *next;
  
  fp = fopen("error.dat","w");
  
  subtrees = 0;
  rms_error = dat_variance = dat_mean = 0.0;

  for(example=0 ; example<N_EXAMPLES ; example++)
    {
      refresh_graph_window();

      /* add a new subtree (except the first time) */
      if ((example>1) && ((LEARNING * ((int) example/LEARNING)) == example))
	{
	  rate = INITIAL_RATE / (1.0 + subtrees * DELTA_RATE); 
	  rms_error = sqrt(rms_error/ESTIMATION) / 
	    sqrt(dat_variance/ESTIMATION - square(dat_mean/ESTIMATION));
	  printf("6-step nrms error = %.4f\n", rms_error);
	  fprintf(fp,"%.4f\n", rms_error);
	  
	  /* basis potentials are the first layer in the net */
	  for (next = root_node->below ; next ; next=next->next)
	    basis_potentials[next->number] = fabs(next->var_weight);

	  /* now add the new trees */
	  for (i=0 ; i<MULT_NODES ; i++)
	    add_level(root_node, basis_potentials, (int) TABLE_SIZE);
	  subtrees++;
	}

      /* only count variance after net has converged a little */
      /*  (also initializes variances for new subtrees) */
      if ((ESTIMATION * ((int) example/ESTIMATION)) == example) 
	{
	  zero_var(root_node);
	  rms_error = 0.0;
	}

      /* interpolate and calculate mserror */
      if (show_results == 1 && ring_pos == BUF_SIZE - 1) calculate_error();

      /* integrate differential equation for current time */
      ring_pos = buf_mod(1 + ring_pos);
      mackey_glass();
      if (ring_pos == 0) clear_graph_window();
      graph_line(ring_pos, ring_buffer[buf_mod(ring_pos - 1)], 
		 ring_buffer[ring_pos], min_value, max_value);
      dat_mean = (dat_mean * example + ring_buffer[ring_pos])/(example+1);
      dat_variance = (dat_variance * example 
		       + square(ring_buffer[ring_pos] - dat_mean))
	/ (example+1);

      /* forward propagation */
      get_basis_values();
      forward_propagate(root_node);

      /* calculate output error */
      network_output = root_node->output;
      approx_buffer[ring_pos] = network_output;
      graph_line(ring_pos, approx_buffer[buf_mod(ring_pos - 1)], 
		 approx_buffer[ring_pos], min_value, max_value);
      root_node->error = ring_buffer[ring_pos] - network_output;
      root_node->var_error = root_node->error * root_node->error
	- root_node->var_output;
      rms_error += square(root_node->error);

      /* modify weights */
      backward_propagate(root_node);
    
    }
}

int buf_mod(j)
     int j;
{
  int i;

  if (j >= BUF_SIZE) i = j - BUF_SIZE;
  else if (j < 0) i = j + BUF_SIZE;
  else i = j;
  return i;
}

euler_mackey_glass()
{
  float x, dx, xtau, a, b;

  a = 0.2;
  b = 0.1;
  x = ring_buffer[buf_mod(ring_pos - 1)];
  xtau = ring_buffer[buf_mod(ring_pos - 1 - TAU)];

  dx = (a * xtau / (1.0 + pow(xtau, 10.0))) - (b * x);
  x += TIME_STEP * dx;

  ring_buffer[ring_pos] = x;
}

float mg(x, xtau)
     float x, xtau;
{
  float a,b;

  a = 0.2;
  b = 0.1;
  return (a * xtau / (1.0 + pow(xtau, 10.0))) - (b * x);
}

mackey_glass()
{
  float x, dx, xtau, a, b;
  int p, ptau;
  
  a = 0.2;
  b = 0.1;
  
  p = buf_mod((ring_pos - 1));
  ptau = buf_mod((ring_pos - 1 - TAU));

  /* fourth order runge-kutta */
  x = ring_buffer[p];
  xtau = ring_buffer[ptau];
  k1[p] = mg(x, xtau);   /* k values assoc. with previous x value */
  k2[p] = mg((x + 0.5 * h * k1[p]), (xtau + 0.5 * h * k1[ptau]));
  k3[p] = mg((x + 0.5 * h * k2[p]), (xtau + 0.5 * h * k2[ptau]));
  k4[p] = mg((x + h * k3[p]), (xtau + h * k3[ptau]));
  ring_buffer[ring_pos] = x + (1.0/6.0) * h * 
    (k1[p] + 2.0*k2[p] + 2.0*k3[p] + k4[p]);
  
}

float square(x)
     float x;
{
  return x*x;
}

/* do iterated predictions and calculated mserror ; assumes DELTA == P */
calculate_error()
{
  int estimate, i, b, n, flag;
  float data_mean, data_variance[MAX_PREDICT], network_output;

  /* init. estimate variance buffer and calculate data variance */
  data_mean = 0.0;
  for (i=0 ; i<BUF_SIZE ; i++) data_mean += ring_buffer[i];
  data_mean /= (float) BUF_SIZE;

  for (i=0 ; i<MAX_PREDICT ; i++) 
    {
      error_buffer[i] = 0.0;
      data_variance[i] = 0.0;
    }

  for (estimate=0 ; estimate < ERROR_SAMPLES ; estimate++)
    {
      /* initialize the predict buffer from real data */
      for (i=0 ; i<N_INPUTS ; i++)
	{
	  predict_buffer[i] = ring_buffer[(estimate + i*DELTA)];
	}

      /* calculate iterated predictions for this estimate */
      for (i=N_INPUTS ; i<MAX_PREDICT ; i++)
	{
	  for (n=0 ; n<N_INPUTS ; n++)
	    basis_values[n] = predict_buffer[(i-n-1)];
	  
	  /* get network output */
	  forward_propagate(root_node);
	  network_output = root_node->output;
	  predict_buffer[i] = network_output;
	  if (estimate == 0)
	    graph_long_line((estimate + (i-1) * DELTA), predict_buffer[(i-1)],
			    (estimate + i * DELTA), predict_buffer[i],
			    min_value, max_value);
	  error_buffer[i] += 
	    square((network_output - ring_buffer[(estimate+i*DELTA)]));
	  data_variance[i] += 
	    square((ring_buffer[(estimate+i*DELTA)] - data_mean));
	}
    }
  flag = 0;
  for (i=N_INPUTS ; i<MAX_PREDICT ; i++)
    {
      error_buffer[i] = sqrt(error_buffer[i]) / sqrt(data_variance[i]);
      if (error_buffer[i] > 1.0 && flag == 0) 
	{
	  printf("error > 1.0 at time %d\n", (DELTA * (i-N_INPUTS+1)));
	  flag = 1;
	}
    }
  printf("single-step error = %f\n", error_buffer[N_INPUTS]);
  show_error();
  show_results = 0;
}

/* show the estimated error graph */
show_error()
{
  int i;

  printf("type c at graph window to see error graph\n");
  show_results = 0;
  while (show_results == 0) refresh_graph_window();

  clear_graph_window();

  for (i=1 ; i<=MAX_PREDICT - N_INPUTS ; i++)
    graph_long_line(((i-1) * DELTA), error_buffer[(i+N_INPUTS-2)],
			  (i * DELTA), error_buffer[(i+N_INPUTS-1)],
			  0.0, 1.0);

  printf("type c at graph window to continue\n");
  show_results = 0;
  while (show_results == 0) refresh_graph_window();
  clear_graph_window();
}

/*****************************************************************************/

get_basis_values()
{
  int i;
  
  for (i=0 ; i<N_INPUTS ; i++)
    basis_values[i] = ring_buffer[buf_mod(ring_pos - P - i*DELTA)] - dat_mean;
}

/*****************************************************************************/

done_learning()
{
}

/*****************************************************************************/