/****************************************************************************/
/*
/* File:		"rtrl_dsp.c"
/*				Real-Time Recurrent Learning for an RMLP (DSP version)
/*
/* Authors:		Scott A. Morrison, scott@cnel.ufl.edu
/*				Deniz Erdogmus, deniz@cnel.ufl.edu
/*				Justin Sanchez, justin@cnel.ufl.edu
/*
/*				Computational NeuroEngineering Laboratory
/*				Department of Electrical and Computer Engineering
/*				The University of Florida
/*				Gainesville, FL, 32611, USA
/*				
/* Created:		26 November 2002
/*
/*****************************************************************************/

/* BE SURE THAT READY.ASM IS PLACED BEFORE THIS MAIN() SECTION */
/* CHECK THE CMD FILE TO SEE THIS IS TRUE */       

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

/* assign parameters */
#define	TRAINING_START		0
#define	TRAINING_END		300
#define	TESTING_START		0
#define	TESTING_END			300
#define	TRAJ_LENGTH			30
#define	TRAJ_PER_UPDATE		10
#define	EPOCHS				50
#define	N0					104
#define	N1					5
#define	N2					3
#define	SLOPE				0.5f
#define	DERIVATIVE_OFFSET	0.001f
#define	ETA1_START			-0.0001f
#define	ETA2_START			-0.00001f
#define	ETAF_START			-0.0001f
#define	ADAPTIVE_STEP_SIZE	1
#define	ETA_MIN				0.00000000001f
#define	ETA_MAX				4.0f
#define	ETA_INCREASE		1.02f	/* ETAnew = ETAold * ETA_INCREASE; */
#define	ETA_DECREASE		0.1f	/* ETAnew = ETAold * ETA_DECREASE; */
#define	MOMENTUM			0.7f
#define	USE_INITS			1
 
#define	TRAINING_LENGTH		300
#define	UPDATES_PER_EPOCH	1

/* Give access to C33 Registers */
extern cregister volatile unsigned int IE;	/* CPU/MS interrupt enable register */
extern cregister volatile unsigned int IF;	/* CPU interrupt flag register */
extern cregister volatile unsigned int IOF;	/* I/O Flags */
extern cregister volatile unsigned int ST;	/* Status Register */

/*************************** Temporary Variables ***************************/
unsigned int	epoch	= 0;

float	d_power = 0.0f;	/* power of the desired signal, for normalized MSE */
float	e_power = 0.0f;	/* power of the error signal, for normalized MSE */

extern float W1_array[];	/* This is referencing W1_array in asm file */
float *W1 = W1_array;		/* Pointer to W1 */
extern float b1_array[];	/* This is referencing b1_array in asm file */
float *b1 = b1_array;		/* Pointer to b1 */
extern float W2_array[];	/* This is referencing W2_array in asm file */
float *W2 = W2_array;		/* Pointer to W2 */
extern float b2_array[];	/* This is referencing b2_array in asm file */
float *b2 = b2_array;		/* Pointer to b2 */
extern float Wf_array[];	/* This is referencing Wf_array in asm file */
float *Wf = Wf_array;		/* Pointer to Wf */     

extern float G_W1_asm[];	/* gradient updates for above */
float *G_W1 = G_W1_asm;
extern float G_b1_asm[];
float *G_b1 = G_b1_asm;
extern float G_W2_asm[];
float *G_W2 = G_W2_asm;
extern float G_b2_asm[];
float *G_b2 = G_b2_asm;
extern float G_Wf_asm[];
float *G_Wf = G_Wf_asm;

extern float U_W1_asm[];	/* weight updates for above */
float *U_W1 = U_W1_asm;
extern float U_b1_asm[];
float *U_b1 = U_b1_asm;
extern float U_W2_asm[];
float *U_W2 = U_W2_asm;
extern float U_b2_asm[];
float *U_b2 = U_b2_asm;
extern float U_Wf_asm[];
float *U_Wf = U_Wf_asm;

extern float c_asm[];		/* intermediate update parameters */
float *c = c_asm;
extern float g_asm[];
float *g = g_asm;
extern float h_asm[];
float *h = h_asm;

extern float y1p_asm[];		/* previous Y1 */
float *y1p = y1p_asm; 
extern float y1p_save_asm[];	/* previous Y1 */
float *y1p_save = y1p_save_asm;
extern float z1_asm[];		/* input to hidden layer */
float *z1 = z1_asm;
extern float Y1_asm[];		/* output of hidden layer */
float *Y1 = Y1_asm;
extern float y2_asm[];		/* network output */
float *y2 = y2_asm;
extern float e_asm[];		/* network error */
float *e = e_asm;
extern float a_asm[];		/* save sech(z1)^2 */
float *a = a_asm;
extern float temp_asm[];
float *temp = temp_asm;

extern float chunkMSE_End[];
float *endMSE = chunkMSE_End;

far extern float input_array[];		/* pointers to data */ 
far extern float desired_array[];
far float *input = input_array;
far float *desired = desired_array;
         
float	*G_W1,*G_b1,*G_W2,*G_b2,*G_Wf;	/* pointers to gradients */
float	*U_W1,*U_b1,*U_W2,*U_b2,*U_Wf;	/* pointers to weight updates */
float	*c,*g,*h;						/* pointers to update parameters */
float	*y1p,*z1,*Y1,*y2,*e,*a,*temp;	/* pointers to network feedforward */

float	*x_ptr = NULL;
float	*d_ptr = NULL;

float	MSE = 0.0f;
float	MSEp = 0.0f;

float	ETA1 = ETA1_START;
float	ETA2 = ETA2_START;
float	ETAF = ETAF_START;

/*************************** Function Prototypes ***************************/
void updateGradient();
void doTraining();

/***************************** Main program **********************************/
void main(void)
{
    unsigned int	i = 0;
	unsigned int	j = 0;
	unsigned int	lines = 0;

	x_ptr = input;		/* initialize the pointer to the input data */
	d_ptr = desired;	/* initialize the pointer to the desired data */
  
	doTraining();
  
	/* we're done, so quit */

} /* end of main() */


/*****************************************************************************/
void doTraining()
{
	unsigned int update	= 0;
	unsigned int traj	= 0;
	unsigned int i		= 0;
	unsigned int T		= 0;
 	
 	asm("	XOR		01000000b,IOF   ; TOGGLE XF1");                     
	/* one epoch is a single pass through 300 samples */
	for(epoch=0;epoch<EPOCHS;epoch++)
	{
		asm("	XOR		00000100b,IOF   ; TOGGLE XF0"); 
		/* reset the hidden state before starting epoch */
        /* resetting the state is the job of the PC */

		/* restore the previous state before continuing */
		for(i=0;i<N1;i++)
		{
			y1p[i] = y1p_save[i];
		}
 
		/* reset x_ptr */
		x_ptr = input;
		d_ptr = desired;

		/* reset for the error calculation */
		e_power = 0.0f;
   		
		/* reset gradients and momentum terms */
		for (i=0;i<N0*N1;i++)
		{
			U_W1[i] = 0.0f;
			G_W1[i] = 0.0f; 
		}

		for (i=0;i<N1;i++)
		{
			U_b1[i] = 0.0f;
			G_b1[i] = 0.0f; 
		}

		for (i=0;i<N1*N1;i++)
		{
			U_Wf[i] = 0.0f;
			G_Wf[i] = 0.0f; 
		}

		for (i=0;i<N1*N2;i++)
		{
			U_W2[i] = 0.0f;
			G_W2[i] = 0.0f; 
		}

		for (i=0;i<N2;i++)
		{
			U_b2[i] = 0.0f;
			G_b2[i] = 0.0f; 
		}

		for(traj=0;traj<TRAJ_PER_UPDATE;traj++)
		{
			/* reset c, g, and h */
			for(i=0;i<N1*N1;i++)
			{
				c[i] = 0.0f;
			}
			for(i=0;i<N1*N1*N0;i++)
			{
				g[i] = 0.0f;
			}				
			for(i=0;i<N1*N1*N1;i++)
			{
				h[i] = 0.0f;
			}
			for(T=0;T<TRAJ_LENGTH;T++)
			{	
				updateGradient();

				/* save the state */
				for(i=0;i<N1;i++)
				{
					y1p[i] = Y1[i];
				}

				/* move to next input and desired */
				x_ptr = x_ptr + N0;
				d_ptr = d_ptr + N2;
			}

		} /* end of loop over all trajectories within this block */
 
         /* save previous MSE */  
		MSEp = *(endMSE);
		        
		/* write MSE to memory for PC reading */
        *(endMSE) = e_power / (N2*TRAINING_LENGTH);

#if (ADAPTIVE_STEP_SIZE == 1)		 
		/* adjust the step size */
		if (epoch != 0)
		{
 			if (*(endMSE) < MSEp)
			{
				/* increase the step size */
				ETA1 = ETA1 * ETA_INCREASE;
				ETA2 = ETA2 * ETA_INCREASE;
				ETAF = ETAF * ETA_INCREASE;

				if (ETA1 > ETA_MAX)
				{
					ETA1 = ETA_MAX;
				}
				if (ETA2 > ETA_MAX)
				{
					ETA2 = ETA_MAX;
				}
				if (ETAF > ETA_MAX)
				{
					ETAF = ETA_MAX;
				}
			}
			else
			{
				/* decrease the step size */
				ETA1 = ETA1 * ETA_DECREASE;
				ETA2 = ETA2 * ETA_DECREASE;
				ETAF = ETAF * ETA_DECREASE;

				/* make sure the step size isn't too small */
				if (ETA1 < ETA_MIN)
				{
					ETA1 = ETA_MIN;
				}
				if (ETA2 < ETA_MIN)
				{
					ETA2 = ETA_MIN;
				}
				if (ETAF < ETA_MIN)
				{
					ETAF = ETA_MIN;
				}
			}
	
		} /* end of adaptive step size block */
#endif /* end of if ADAPTIVE_STEP_SIZE */
 
		/* update weights with the new step size */
		for (i=0;i<N0*N1;i++)
		{
			U_W1[i] = ETA1 * G_W1[i] + MOMENTUM * U_W1[i];
			W1[i] = W1[i] + U_W1[i];
		}

		for (i=0;i<N1;i++)
		{
			U_b1[i] = ETA1 * G_b1[i] + MOMENTUM * U_b1[i];
			b1[i] = b1[i] + U_b1[i]; 
		}

		for (i=0;i<N1*N1;i++)
		{
			U_Wf[i] = ETAF * G_Wf[i] + MOMENTUM * U_Wf[i];
			Wf[i] = Wf[i] + U_Wf[i]; 
		}

		for (i=0;i<N1*N2;i++)
		{
			U_W2[i] = ETA2 * G_W2[i] + MOMENTUM * U_W2[i];
			W2[i] = W2[i] + U_W2[i]; 
		}

		for (i=0;i<N2;i++)
		{
			U_b2[i] = ETA2 * G_b2[i] + MOMENTUM * U_b2[i];
			b2[i] = b2[i] + U_b2[i]; 
		}
         
	} /* end of loop over epochs */  
	
	/* save the state for next chunk of data, invoked by PC */
	for(i=0;i<N1;i++)
	{
		y1p_save[i] = y1p[i];
	}

} /* end of doTraining() */


/*****************************************************************************/
inline void updateGradient()
{
	/* temporary variables used in this function */
	int		i = 0;
	int		j = 0;
	int		m = 0;
	int		k = 0;
	int		gamma = 0;
	int		alpha = 0;
	int		beta = 0;
	
	register float regFloat;

	/* forward pass through the network */
	for(j=0;j<N1;j++)
	{
        regFloat = 0;       
         
		for (i=0;i<N0;i++)
		{
			regFloat = regFloat + x_ptr[i]*W1[j*N0+i];       
		}

		for (m=0;m<N1;m++)
		{
			regFloat = regFloat + y1p[m] * Wf[j*N1+m];       			
		}

		regFloat = regFloat + b1[j];       
		Y1[j] = tanh(SLOPE*regFloat);	 	       
		z1[j] = regFloat;	       	
		a[j] = SLOPE*(1-(Y1[j]*Y1[j])+DERIVATIVE_OFFSET);  
	}

	for(k=0;k<N2;k++)
	{
		/* y2[k]= 0;   */
        regFloat = 0;	       

		for (j=0;j<N1;j++)
		{
			regFloat = regFloat + Y1[j]*W2[k*N1+j];     
		}

		regFloat = regFloat + b2[k];	     
		
		e[k] = d_ptr[k] - regFloat;       
		y2[k] = regFloat;      
		e_power = e_power + (e[k] * e[k]); /* cumulate the error power */
	}


	/* The forward pass is done. Now backpropagate the updates. */

	/* update c */
	for(beta=0;beta<N1;beta++)
	{
		for(j=0;j<N1;j++)
		{
			regFloat = 0.0f;      
			for(i=0;i<N1;i++)
			{
				regFloat = regFloat + Wf[j*N1+i] * c[beta*N1+i];  
			}

			if (j == beta)
			{
				c[beta*N1+j] = regFloat + 1;  
			}
			else
			{
				c[beta*N1+j] = regFloat;  
			}
			
			c[beta*N1+j] = c[beta*N1+j] * a[j];
			
		}
	}

	/* update g */
	for(gamma=0;gamma<N0;gamma++)
	{
		for(beta=0;beta<N1;beta++)
		{
			for(alpha=0;alpha<N1;alpha++)
			{
				regFloat = 0.0f;    
				for (i=0;i<N1;i++)
				{
                    regFloat = regFloat + Wf[alpha*N1+i] * g[gamma*N1*N1+beta*N1+i];    
				}
				if (alpha == beta)
				{					
					g[gamma*N1*N1+beta*N1+alpha] = regFloat + x_ptr[gamma]; 
				}
				else
				{
					g[gamma*N1*N1+beta*N1+alpha] = regFloat; 
				}
			
				g[gamma*N1*N1+beta*N1+alpha] = g[gamma*N1*N1+beta*N1+alpha] * a[alpha];

			}
		}
	}

	/* update h */
	for(gamma=0;gamma<N1;gamma++)
	{
		for(beta=0;beta<N1;beta++)
		{
			for(alpha=0;alpha<N1;alpha++)
			{
				regFloat = 0.0f;   
				for (i=0;i<N1;i++)
				{
					regFloat = regFloat + Wf[alpha*N1+i] * h[gamma*N1+beta*N1+i];  					
				}
				if (alpha == beta)
				{
					h[gamma*N1+beta*N1+alpha] = regFloat + y1p[gamma];  
				}
				else
				{
					h[gamma*N1+beta*N1+alpha] = regFloat;  
				}
				
				h[gamma*N1+beta*N1+alpha] = h[gamma*N1+beta*N1+alpha] * a[alpha];

			}
		}
	}

	/* update the gradients */
	for (k=0;k<N2;k++)
	{
		G_b2[k] = G_b2[k] - e[k]; 

		for(j=0;j<N1;j++)
		{
			G_W2[k*N1+j] = G_W2[k*N1+j] - e[k]*Y1[j];
		}	

	}

	for (j=0;j<N1;j++)
	{
		regFloat = 0.0f; 
		for (i=0;i<N2;i++)
		{
			regFloat = regFloat + e[i]*W2[i*N1+j]; 
		}

		G_b1[j] = G_b1[j] - c[j*N1+j] * regFloat; 

		for (i=0;i<N0;i++)
		{
			G_W1[j*N0+i] = G_W1[j*N0+i] - g[i*N1*N1+j*N1+j] * regFloat; 
		}
		
		for (i=0;i<N1;i++)
		{
			G_Wf[j*N1+i] = G_Wf[j*N1+i] - h[i*N1+j*N1+j] * regFloat; 
		}

	}

} /* end of updateGradient() */