/*
 *  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"

BackPropHiddenLayer::BackPropHiddenLayer(char* str, NslModule* parent) : 
	NslModule(str,parent), 
 	fInput("fInput", this),
  	bInput("bInput", this),
    mf("mf", this),			
   	mp("mp", this),		
   	delta("delta", this),          
	h("h", this),      		  
	dh("dh", this),    	

	w("w", this),   	
	dw("dw", this), 		
          
	lrate("lrate", this),
	momentum("momentum", this)
{
}

void BackPropHiddenLayer::memAlloc(int inSize,int hidSize)
{
	fInput.memAlloc(inSize);		
	bInput.memAlloc(hidSize);		
	mf.memAlloc(hidSize);		
	mp.memAlloc(hidSize);		
	delta.memAlloc(hidSize);      	      
	h.memAlloc(hidSize);      	  
	dh.memAlloc(hidSize);    			

	w.memAlloc(inSize, hidSize);                         
	dw.memAlloc(inSize, hidSize); 	
}

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

 	dw = 0.0;
  	dh = 0.0;
}
void BackPropHiddenLayer::initTrain()
{		
//	dw = 0.0;
//  dh = 0.0;
}

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

void BackPropHiddenLayer::calculateDelta()
{
  	delta = (mf ^ (1.0 - mf)) ^ bInput; 
}

void BackPropHiddenLayer::updateWeights() 
{
  	dw = lrate * vec_mult_vec_trans(delta, fInput); // + momentum * dw; // dw + 
  	w = w + dw;
}

void BackPropHiddenLayer::updateThresholds()
{
  	dh = lrate * delta; // + momentum * dh; // dh + 
  	h = h + dh;
}

void BackPropHiddenLayer::backwardPass()
{     	   	
      	calculateDelta();
        updateWeights();
        updateThresholds();
}

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

// OutputLayer

BackPropOutputLayer::BackPropOutputLayer(char* str, NslModule* parent) : 
	BackPropHiddenLayer(str,parent),
   	bOutput("bOutput", this)
{
}

void BackPropOutputLayer::memAlloc(int hidSize, int outSize)
{
	BackPropHiddenLayer::memAlloc(hidSize,outSize);		
	bOutput.memAlloc(hidSize);	
}

void BackPropOutputLayer::propagateError()
{
  	bOutput = mat_trans_mult_col_vec(w, delta);
}

void BackPropOutputLayer::backwardPass()
{	      	   	
      	calculateDelta();
      	propagateError();
        updateWeights();
	updateThresholds();
}


// AllLayers

BackPropLayers::BackPropLayers(char* str, NslModule* parent) : 
	AslSchema(str,parent), 
	hl("hiddenLayer",this),
	ol("outputLayer",this),
	input("input",this),
	desired("desired",this),
	output("output",this),
	eOutput("eOutput",this),
	stopError("stopError", this),
	tss("tss", this),
	pss("pss", this)
{
        initModule();
}

void BackPropLayers::initModule()
{
	relabel(input,hl.fInput);
	relabel(ol.mf,output);
	relabel(eOutput,ol.bInput);
		
	connect(hl.mf,ol.fInput);
	connect(ol.bOutput,hl.bInput);
}

void BackPropLayers::memAlloc(int inSize,int hidSize, int outSize)
{
	input.memAlloc(inSize);			
	desired.memAlloc(outSize);		
 	output.memAlloc(outSize);
	eOutput.memAlloc(outSize);
	
	hl.memAlloc(inSize,hidSize);
	ol.memAlloc(hidSize,outSize);  
}

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


void BackPropLayers::initTrain()
{
  	pss = 0.0;
  	output = 0;
}

void BackPropLayers::calculateError()
{ 	      	     	
  	eOutput = desired - output;
  	pss = pss + (eOutput ^ eOutput).sum();
}

void BackPropLayers::simTrain()
{   	
      	hl.forwardPass();
      	ol.forwardPass();
      	
  	calculateError();
  	
  	ol.backwardPass();
  	hl.backwardPass();
}

void BackPropLayers::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;
	}
}

void BackPropLayers::simRun()
{ 	   
      	input.print(cout);
      		   	
      	hl.forwardPass();
      	ol.forwardPass();
      	
      	output.print(cout);
}
    
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("train",this), bp("layers",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");