/*
 *  File : backprop.C
 *
 *  History
 *     Date       Author               Description
 *     -------    -----------------    --------------------------------------
 *     08/23/91   Andrew H. Fagg       Original Release
 *     09/19/93   John C. Weng	       Release With NSL 2.5
 *
 */

#include "nsl_include.h"
#include "backprop.h"
#include "alib.h"

BPForward::BPForward(char* str, NslModule* parent) : 
	NslModule(str,parent), 
 	fInput("fInput", this),
    mf("mf", this),			
   	mp("mp", this),		  
	h("h", this),      		  
	dh("dh", this),    	
	w("w", this),   	
	dw("dw", this)		
{
}

void BPForward::memAlloc(int inSize,int hidSize)
{
	fInput.memAlloc(inSize);			
	mf.memAlloc(hidSize);		
	mp.memAlloc(hidSize);		 
	h.memAlloc(hidSize);      	  
	dh.memAlloc(hidSize);    			
	w.memAlloc(inSize, hidSize);                         
	dw.memAlloc(inSize, hidSize); 	
}

void BPForward::initSys()
{
  	randomize(w);
  	w = w/2;
  	randomize(h);
  	h = h/2;

 	dw = 0.0;
  	dh = 0.0;
}


void BPForward::forwardPass()
{ 	 
    mp = mat_mult_col_vec(w, fInput);    	
  	mf = sigmoid(mp + h);
}

void BPForward::simRun() 
{  
	forwardPass();
}

void BPForward::simTrain() 
{
 	w = w + dw;
 	h = h + dh;
	forwardPass();
}

void BPForward::endTrain()
{
// 	w = w + dw;
//	h = h + dh;
}

// Backward Hidden Layer

BPBackward::BPBackward(char* str, NslModule* parent) : 
	NslModule(str,parent), 
 	fInput("fInput", this),
 	bInput("bInput", this),
    mf("mf", this),			
   	delta("delta", this),		  
	dh("dh", this),    	
	dw("dw", this), 	
	lrate("lrate", this),   	
	momentum("momentum", this)
{
}

void BPBackward::memAlloc(int inSize,int hidSize)
{
	fInput.memAlloc(inSize);		
	bInput.memAlloc(hidSize);		
	mf.memAlloc(hidSize);			
	delta.memAlloc(hidSize);      	        
	dh.memAlloc(hidSize);    			          
	dw.memAlloc(inSize, hidSize); 	
}

void BPBackward::simTrain()
{
	backwardPass();
}
void BPBackward::backwardPass()
{     	   	
  	delta = (mf ^ (1.0 - mf)) ^ bInput; 
  	dw = lrate * vec_mult_vec_trans(delta, fInput); // + momentum * dw; // dw + 
  	dh = lrate * delta; // + momentum * dh; // dh + 
}

// Backward Prop 

BPBackwardProp::BPBackwardProp(char* str, NslModule* parent) : 
	BPBackward(str,parent),
   	bOutput("bOutput", this),
	w("w", this)
{
}

void BPBackwardProp::memAlloc(int hidSize, int outSize)
{
	BPBackward::memAlloc(hidSize,outSize);		
	bOutput.memAlloc(hidSize);	
	w.memAlloc(hidSize,outSize);                      
}

void BPBackwardProp::simTrain()
{
    backwardPass();

  	bOutput = mat_trans_mult_col_vec(w, delta);
}

// Backward Error 

BPBackwardError::BPBackwardError(char* str, NslModule* parent) : 
	NslModule(str,parent),
   	mf("mf", this),
	desired("desired",this),
	eOutput("eOutput",this),
	stopError("stopError", this),
	tss("tss", this),
	pss("pss", this)
{
}

void BPBackwardError::initSys()
{
  	tss = 0.0;
}

void BPBackwardError::memAlloc(int outSize)
{
	mf.memAlloc(outSize);	
	desired.memAlloc(outSize);		
	eOutput.memAlloc(outSize);       
}

void BPBackwardError::initTrain()
{
  	pss = 0.0;
//	mf = 0;
}

void BPBackwardError::simTrain()
{
  	eOutput = desired - mf;
  	pss = pss + (eOutput ^ eOutput).sum();
}

void BPBackwardError::endTrain()
{
	tss = pss;
	int tepoch = (int) NSL_SYSTEM->getEpochStep();	 
	cmd_out("Epoch: ", tepoch);
	cmd_out("tss = ", tss.get_value());
	if (tss < stopError) {
		cmd_out("Convergence");
		nslBreakEpoch();
		return;
	}
}

// AllLayers

BackPropLayers::BackPropLayers(char* str, NslModule* parent) : 
	NslModule(str,parent), 
	fh("fh",this),
	fo("fo",this),
	be("be",this),
	bo("bo",this),
	bh("bh",this),
	input("input",this),
	output("output",this),
	desired("desired",this)
{
    makeConn();
}

void BackPropLayers::makeConn()
{
	relabel(input,fh.fInput);
	relabel(input,bh.fInput);
	relabel(fo.mf,output);
	relabel(desired,be.desired);
		
	connect(fh.mf,fo.fInput);
	connect(fh.mf,bo.fInput);
	connect(fh.mf,bh.mf);
	connect(fo.mf,bo.mf);
	connect(fo.mf,be.mf);

	connect(bh.dw,fh.dw);
	connect(bh.dh,fh.dh);

	connect(bo.dw,fo.dw);
	connect(bo.dh,fo.dh);
	connect(fo.w,bo.w);

	connect(be.eOutput,bo.bInput);
	connect(bo.bOutput,bh.bInput);
}

void BackPropLayers::memAlloc(int inSize,int hidSize, int outSize)
{
	input.memAlloc(inSize);			
	desired.memAlloc(outSize);		
 	output.memAlloc(outSize);
	
	fh.memAlloc(inSize,hidSize);
	fo.memAlloc(hidSize,outSize);  
	bh.memAlloc(inSize,hidSize);
	bo.memAlloc(hidSize,outSize);  
	be.memAlloc(outSize);  
}

// train manager
    
TrainManager::TrainManager(char* str, NslModule* parent) : 
	NslModule(str,parent), 
	dInput("dInput", this),		
	dOutput("dOutput", this),
	pInput("pInput", this),
	pOutput("pOutput", this),
	numPats("numPats",this),
	fname("fname",this)
{
}
void TrainManager::memAlloc(int nPats,int inSize,int outSize)
{
	dInput.memAlloc(inSize);	
	dOutput.memAlloc(outSize);
	pInput.memAlloc(nPats, inSize);
	pOutput.memAlloc(nPats, outSize);
}

void TrainManager::readFile(int inSize,int outSize)
{
  	int i,nPats,tmp;
  	int pat;
	nsl_file file("temp"); // dummy name !?

	cmd_out("fname: ",fname.elem());
	file.init_file(fname.elem(),""); // join with next call
  	if (file.open_file(NSL_INPUT) == -1)
  		cmd_out("Bad file name: ",fname.elem());
  	else {
  		file.read(nPats);
//  	file.read(iSize);
// 		file.read(oSize);
		memAlloc(nPats,inSize,outSize);
		numPats = nPats;
 	
// 		if (iSize == inSize && oSize == outSize && nPats == numPats) {
 		
  	   	    for(pat = 0; pat < nPats; ++pat) {
      	      		for(i = 0; i < inSize; ++i) {
	  			file.read(tmp);
	  			pInput[pat][i] = tmp;
	      	    	}
      	      	    	for(i = 0; i < outSize; ++i) {
	  			file.read(tmp);
	  			pOutput[pat][i] = tmp;
	      	    	}
    	   	    }
/*    		}
    		else {
    			cmd_out("Bad data from file");
    			cmd_out("Pattern number: ",nPats);
    			cmd_out("Input Layer Size: ",iSize);
    			cmd_out("Output Layer Size: ",oSize);    		
    		}
*/   		
  		file.close_file();
  	}
}

void TrainManager::simTrain()
{
	int timeStep = NSL_SYSTEM->getTrainCurrentTimeStep() - 1;
	int endStep = (int) NSL_SYSTEM->getTrainEndTime() + 0.5;
	int current_pat = timeStep%endStep;
      	extract_row_vector(pInput, current_pat, dInput);
      	extract_row_vector(pOutput, current_pat, dOutput); 
}

BackPropModel::BackPropModel() : NslModule("BackPropModel"),
	tf("tf",this), bp("bp",this), 
	inSize("inSize",this),
	hidSize("hidSize",this),
	outSize("outSize",this)
{
        makeConn();
}

void BackPropModel::makeConn()
{
	connect(tf.dInput,bp.input);
	connect(tf.dOutput,bp.desired);
}
void BackPropModel::initSys()
{
//	tf.memAlloc(numPats.elem(),inSize.elem(),outSize.elem());		
	bp.memAlloc(inSize.elem(),hidSize.elem(),outSize.elem());

	tf.readFile(inSize.elem(),outSize.elem());
}

AslSchemaModel<BackPropModel> _BackPropModel("BackPropModel");