/* poly-tree.c
**
**    routines for performing nonlinear tree-oriented lms approximation
**  11/2/90
**   modified 11/7/90 for variable structure trees
**   modified 2/20/91 to include Rich Sutton's ideas about growing one
**       subleaf at a time, using weight change variance...
**   modified 5/23/91 for compress.c application, and fixed n->input_var
*/

#include <math.h>
#include <stdio.h>
#include "poly-tree.h"

NODE *find_max_var();
NODE *nalloc();

/* this routine sets n->output by recursively descending the tree below n */
/* note that this version, unlike others, adds in the root node weight */

forward_propagate(n)
     NODE *n;
{
  NODE *next;
  int lev;
  
  lev = n->level;  /* figure out which parts participate here */
  next = n->below;  /* start of subnode list */
  n->output = n->weight;   /* initialize result; adds own weight */
  n->var_output = n->var_weight;

  while(next)  /* returns if this is a leaf */
    {
      forward_propagate(next);
      n->output += next->output * basis_values[next->number] ;
      n->var_output += next->var_output * basis_values[next->number]
	* basis_values[next->number] ;

      next = next->next;  /* follow chain until NULL */
    }
}

/* note that, unlike other versions, this adapts the root node weight */
/* assumes n->error has been set */

backward_propagate(n)
     NODE *n;
{
  NODE *next;
  
  next = n->below;
  n->weight += rate * n->error;

  n->var_weight += rate * n->var_error ;  /* regress variances */

  n->var *= 1.0 - alpha;
  n->var += alpha * n->error * n->error;
  
  if (n->level > 0)  /* root node has input_var = 1.0 always */
    {
      n->input_var *= 1.0 - alpha;
      n->input_var += alpha * n->value * n->value ; 
    }
  
  while(next) 	/* return if this is a leaf */
    {
      next->error = n->error * basis_values[next->number];
      next->var_error = n->var_error * basis_values[next->number]
	* basis_values[next->number] ;
      next->value = n->value * basis_values[next->number];  /* prod basis */
      backward_propagate(next);

      next = next->next;
    }
}

/* add a new basis below the selected node with high error variance */

add_level(n, basis_variances, table_size)
     NODE *n;  /* usually the root node */
     float *basis_variances;   /* input variance for each basis number */
     int table_size;
{
  float max_var, max_basis_var;
  int i, j, max_basis_num, lev, *subnode_list, full_flag;
  NODE *worst, *next, *new;

  subnode_list = (int *) malloc(table_size * sizeof(int));
  
  max_basis_var = -1001.0;
  while (max_basis_var < -1000.0)  /* repeats until a possible node is found */
    {
      /* find the worst node below n to add a sub-basis to */
      max_var = 0.0;
      worst = find_max_var(n, &max_var);  /* can find intermediate nodes */
      lev = worst->level;
      if (max_var < -1000.0) printf("no room in tree %.1f\n",max_var);
      
      /* eliminate disallowed new nodes from basis_vars[] list */
      for (i=0 ; i<table_size ; i++) subnode_list[i] = 0;
      for (next = worst->below; next; next = next->next) /* if already below */
	subnode_list[next->number] = 1;
      
      /* find input var which is not yet a subnode and which has the 
	 largest variance */
      max_basis_var = -1001.0;
      full_flag = 1;
      for (i=0 ; i<table_size ; i++)
	if ((subnode_list[i] == 0) && (basis_variances[i] > max_basis_var))
	  {
	    max_basis_var = basis_variances[i];
	    max_basis_num = i;
	    full_flag = 0;
	  }
      
      if ((lev == 0) && (max_basis_var == 0.0))
	{
	  max_basis_num = 0;
	  printf("no variance in tree \n");
	}

      /* if nothing was possible at this node, try to find another */
      if (full_flag == 1)
	{
	  printf("full intermediate node - continuing\n");
	  worst->flag = 1;  /* so find_max_var won't */
	}
    }

  printf("adding %d th node number %d below level %d node %d\n", 
	 subtrees, max_basis_num, worst->level, worst->number);
      
  /* allocate a new subnode below the worst node */
  if (worst->below)
    {
      /* get to end of chain */
      for (next = worst->below; next->next; next = next->next);
      new = next->next = nalloc(1);
    }
  else new = worst->below = nalloc(1);

  if (new == NULL)
    {
      printf("nalloc: out of memory\n");
      exit(1);
    }

  /* initialize the new node */
  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 = worst->level + 1;
  new->number = max_basis_num;
  new->below = NULL;  /* so we know it's a leaf */
  new->next = NULL;
  new->flag = 0;  /* allowed to add subtrees */

  free(subnode_list);
  
}

/* finds current worst node based on max. error variance */

NODE *find_max_var(n, p_max_var)
     NODE *n;
     float *p_max_var;
{
  NODE *worst, *new, *next;
  float max_var, var;
  int i, lev;

  /* don't build trees too deep */
  lev = n->level;
  if (lev >= max_levels)
    {
      *p_max_var = -1002.0;
      return (n);
    }

  /* search subtree recursively to find worst */
  worst = n;		/* returns self if nothing below is worse */

  /* sutton's criterion: maximize the "potential" */
  max_var = fabs(n->var_weight);  /* *n->input_var  */
    /* n->var / n->input_var;  /* normalize by "input" variance */

  if (n->flag) max_var = -1001.0;  /* don't use this node */
  
  for (next = n->below ; next ; next = next->next)
    {
      new = find_max_var(next, &var);  
      if (var > max_var)  /* strictly > so tends not to go deep ... */
	{
	  worst = new;
	  max_var = var;
	}
    }
  
  *p_max_var = max_var;
  return(worst);
}

/* zero all variances so we can start incrementally computing a new estimate */

zero_var(n)
     NODE *n;
{
  NODE *next;

  n->var = 0.0;
  n->input_var = 0.0; 
  /*  n->var_weight = 0.0;  */
  
  for (next = n->below; next ; next = next->next)
    zero_var(next);
}

NODE *nalloc(i)
     int i;
{
  return (NODE *) malloc(i * sizeof(NODE));
}