#include "bp.h"

//---------------------------------BEGIN MAIN---------------------------------
int exampleUsage(int argc,char** argv)
{
	int i,sum,numCorrect,loaded,minIdx=0;
	float minErr;
	NeuralNet* nn;
	Data *learn,*test;
	FILE *errTrain,*errVal,*errTrainHid,*errTrainOut,*trainCorrect,*valCorrect;
	clock_t start;

	//Our main neural net structure
	printf("Initializing Neural Net ...\n");
	if(! (nn=(NeuralNet*)malloc(sizeof(NeuralNet))) ) {
		perror("NN struct"); exit(-1);
	}
	srand(1);
	
	//---Command Line Parameters
	if(argc>3) {
		printf("Usage: ./nn [agent.txt]\n");
		exit(-1);
	} else if (argc==2) {
		loadAgent(argv[1],nn);
		printf("  Loaded %s\n",argv[1]);
		loaded=1;
	} else {
		if(HALF_RESOLUTION==1) createNN(nn,1,INPUT_NODES/2,HIDDEN_NODES,OUTPUT_NODES,LEARNING_RATE);
		else createNN(nn,1,INPUT_NODES,HIDDEN_NODES,OUTPUT_NODES,LEARNING_RATE);
		assignRandomWeights(nn);
		loaded=0;
	}

	//Our error files
	if(WRITE_FILES==1) {
		errTrain=fopen("errTrain.txt","w");
		errTrainHid =	fopen("errTrainHid.txt","w");
		errTrainOut =	fopen("errTrainOut.txt","w");
		errVal =		fopen("errVal.txt","w");
		trainCorrect =	fopen("trainCorrect.txt","w");
		valCorrect =	fopen("valCorrect.txt","w");
	}

	//---Allocate and store all datasets (train,validate,test)
	printf("Loading Data ...\n");
	if(! (learn=(Data*)malloc(sizeof(Data))) ) {
		perror("learn data"); exit(-1);
	}
	if(! (test=(Data*)malloc(sizeof(Data))) ) {
		perror("test data"); exit(-1);
	}
	//Use macros defined at the top
	//initData(learn);
	//printf("  Loaded learn data\n");
	//initData(test);
	//printf("  Loaded test data\n");

	//---Learning
	printf("Learning ...\n");
	sum=0; minErr=RAND_MAX;
	printf("  Epoch: %d, MSE: %f\n",nn->nepochs,nn->a.mse);
	start=clock();
	while (1) {
		//Go through all input data once (one epoch) and learn each time
		onlineNN(nn,learn,1);
		//printf("%d: %0.6f\t",nn->nepochs,nn->a.mse);//Keep the user informed

		//Every few epochs, run on validation data
		if(nn->nepochs%NUM_EPOCHS==0) {
			printf("  Epoch: %d, MSE: %f",nn->nepochs,nn->a.mse);
			saveAgent(nn);

			/*
			//Original function for testing with optical character recognition
			//Write errors to file
			if(WRITE_FILES==1) {
				//MSE errors
				fprintf(errTrain,"%d\t%f\n",nn->nepochs,nn->a.mse);
				fprintf(errTrainHid,"%d\t",nn->nepochs);
				fprintf(errTrainOut,"%d\t",nn->nepochs);
				for(i=0;i<nn->nhid;i++) fprintf(errTrainHid,"%f\t",nn->a.sumErrHid[i]/nn->nepochs);
				for(i=0;i<nn->nout;i++) fprintf(errTrainOut,"%f\t",nn->a.sumErrOut[i]/nn->nepochs);
				fprintf(errTrainHid,"\n");
				fprintf(errTrainOut,"\n");
				//Number correct with percentage
				fprintf(trainCorrect,"%d\t",nn->nepochs);
				//for(i=0;i<nn->nout;i++) fprintf(trainCorrect,"%d\t",nn->a.numCorrect[i]);
				for(i=0;i<nn->nout;i++) fprintf(trainCorrect,"%0.2f\t",nn->a.numCorrect[i]/(float)learn->nitems*100.f);
				fprintf(trainCorrect,"\n");
			}

			//---Validate, don't learn
			printf("\nValidating ...");
			numCorrect=onlineNN(nn,validate,0);
			if(WRITE_FILES==1) {
				fprintf(errVal,"%d\t%f\n",nn->nepochs,nn->a.mse);
				fprintf(valCorrect,"%d\t",nn->nepochs);
				//for(i=0;i<nn->nout;i++) fprintf(valCorrect,"%d\t",nn->a.numCorrect[i]);
				for(i=0;i<nn->nout;i++) fprintf(valCorrect,"%f\t",nn->a.numCorrect[i]/(float)validate->nitems*100.f);
				fprintf(valCorrect,"\n");
			}
			//Find the best network
			if(nn->a.mse<minErr) {
				minErr=nn->a.mse;
				minIdx=nn->nepochs;
			}
			

			printf("  mse = %0.6f  (%0.2f%%)\n",nn->a.mse,
				((float)numCorrect / (float)(validate->nsup*validate->nitems)) * 100.f);
			if (nn->a.mse<(1e-5) || nn->nepochs>=MAX_EPOCHS) break;
			*/
		}
	}
	printf("  Training took %0.2f seconds, best network is %d\n",(clock()-start)/(float)CLOCKS_PER_SEC,minIdx);

	//---Testing
	printf("Testing ...\n");
	numCorrect=onlineNN(nn,test,0);
	if(WRITE_FILES==1) {
		fprintf(trainCorrect,"--------Test Result Accuracy--------\n");
		fprintf(trainCorrect,"%d\t",nn->nepochs);
		//for(i=0;i<nn->nout;i++) fprintf(trainCorrect,"%d\t",nn->a.numCorrect[i]);
		for(i=0;i<nn->nout;i++) fprintf(trainCorrect,"%f\t",nn->a.numCorrect[i]/(float)test->nitems*100.f);
		fprintf(trainCorrect,"\n");
	}

	//Output percentage of test data that the neural network got correct
	printf("Network is %0.2f%% correct\n", 
		  ((float)numCorrect / (float)(test->nsup*test->nitems)) * 100.f);

	//Close our error files so everything is written
	if(WRITE_FILES==1) {
		fclose(errTrain);
		fclose(errTrainHid);
		fclose(errTrainOut);
		fclose(errVal);
		fclose(trainCorrect);
		fclose(valCorrect);
	}
	return 0;
}
//----------------------------------END MAIN----------------------------------
/*
//Initialize data from which to learn (stores all in memory)
// Original function for testing with digital character recognition
void initData(Data* d, int start, int end,char* dir) {
	int i,j,k,fileNum,tmp;
	char fn[127],ext[7];
	FILE* fin;

	//Number of data items
	d->nfiles=end-start+1;//assumes inclusive data range; i.e. [1,100]
	d->nin=GRID_X*GRID_Y+1;//Extra bias input
	d->nsup=N_CHARS;

	//Allocate space safely
	if(! (d->item=(char*)malloc(d->nsup*sizeof(char))) ) {
		perror("init learn letter"); exit(-1);
	}
	if(! (d->in=(float***)malloc(d->nsup*sizeof(float**))) ) {
		perror("init learn in"); exit(-1);
	}
	if(! (d->sup=(float***)malloc(d->nsup*sizeof(float**))) ) {
		perror("init learn sup"); exit(-1);
	}
	//for each letter, allocate space for all files
	for(i=0;i<d->nsup;i++) {
		if(! (d->in[i]=(float**)malloc(d->nfiles*sizeof(float*))) ) {
			perror("init learn in[]"); exit(-1);
		}
		if(! (d->sup[i]=(float**)malloc(d->nfiles*sizeof(float*))) ) {
			perror("init learn sup[]"); exit(-1);
		}
		//for each file, allocate space for the data contained
		for(j=0;j<d->nfiles;j++) {
			if(! (d->in[i][j]=(float*)malloc(d->nin*sizeof(float))) ) {
				perror("init learn in[][]"); exit(-1);
			}
			if(! (d->sup[i][j]=(float*)malloc(d->nsup*sizeof(float))) ) {
				perror("init learn sup[][]"); exit(-1);
			}
		}
	}

	//Store our letters
	for(i=0;i<d->nsup;i++) {
		d->item[i]=DATA[i];
	}

	//Read in and store all data
	for(i=0;i<d->nsup;i++) {//for each letter
		fileNum=start;
		for(j=0;j<d->nfiles;j++) {//for each file

			//Setup filename and open
			strcpy(fn,dir);
			sprintf(ext,"%c%d",d->item[i],fileNum);
			strcat(fn,ext);
			fin=fopen(fn,"r");
			if(fin==NULL) {
				printf("Error: can't open dataset %s.\n", fn); exit(-1);
			}

			//Read this file's data into our struct
			d->in[i][j][0]=1.f;//Bias input
			for(k=1;k<d->nin;k++) {
				fscanf(fin,"%d ",&tmp);
				d->in[i][j][k]=(float)tmp;
			}
			fscanf(fin,"");//blank line
			for(k=0;k<d->nsup;k++) {
				fscanf(fin,"%d ",&tmp);
				d->sup[i][j][k]=(float)tmp;
			}
			fclose(fin);
			fileNum++;
		}
	}
}*/

//Allocate memory for our neural network
void initNN(NeuralNet* nn) {
	int i;

	//Initialize input-to-hidden weights
	if(! (nn->wih=(float**)malloc((nn->nin+1)*sizeof(float*))) ) {
		perror("init input-to-hidden weights"); exit(-1);
	}
	for(i=0;i<nn->nin+1;i++) {
		if(! (nn->wih[i]=(float*)malloc(nn->nhid*sizeof(float))) ) {
			perror("init input-to-hidden weights"); exit(-1);
		}
	}

	//Initialize hidden-to-output weights
	if(! (nn->who=(float**)malloc((nn->nhid+1)*sizeof(float*))) ) {
		perror("init hidden-to-output weights"); exit(-1);
	}
	for(i=0;i<nn->nhid+1;i++) {
		if(! (nn->who[i]=(float*)malloc(nn->nout*sizeof(float))) ) {
			perror("init hidden-to-output weights"); exit(-1);
		}
	}

	//Setup space for activations
	if(! (nn->a.in=(float*)malloc(nn->nin*sizeof(float))) ) {
		perror("init input activation"); exit(-1);
	}
	if(! (nn->a.hid=(float*)malloc(nn->nhid*sizeof(float))) ) {
		perror("init hidden activation"); exit(-1);
	}
	if(! (nn->a.sup=(float*)malloc(nn->nout*sizeof(float))) ) {
		perror("init output activation"); exit(-1);
	}
	if(! (nn->a.actual=(float*)malloc(nn->nout*sizeof(float))) ) {
		perror("init actual output activation"); exit(-1);
	}

	//Setup space for storing errors for backprop change
	if(! (nn->a.errOut=(float*)malloc(nn->nout*sizeof(float))) ) {
		perror("init errOut nodes"); exit(-1);
	}
	if(! (nn->a.errHid=(float*)malloc(nn->nhid*sizeof(float))) ) {
		perror("init errHid nodes"); exit(-1);
	}
	if(! (nn->a.sumErrOut=(float*)malloc(nn->nout*sizeof(float))) ) {
		perror("init avg errOut nodes"); exit(-1);
	}
	if(! (nn->a.sumErrHid=(float*)malloc(nn->nhid*sizeof(float))) ) {
		perror("init avg errHid nodes"); exit(-1);
	}
	if(! (nn->a.numCorrect=(int*)malloc(nn->nout*sizeof(int))) ) {
		perror("init class sup"); exit(-1);
	}
	nn->a.mse=RAND_MAX;
	for(i=0;i<nn->nhid;i++) nn->a.sumErrHid[i]=0.f;
	for(i=0;i<nn->nout;i++) nn->a.sumErrOut[i]=0.f;
	for(i=0;i<nn->nout;i++) nn->a.numCorrect[i]=0;
	nn->a.guess=NULL;
}

//Set the parameters for our neural network and allocate space
void createNN(NeuralNet* nn, int id, int nin, int nhid, int nout,float lr) {
	nn->id=id;
	nn->nepochs=0;
	nn->nin=nin;
	nn->nhid=nhid;
	nn->nout=nout;
	nn->lr=lr;
	initNN(nn);
}

//Allows saving of agent's weights to a file
// Header: NN ID, number of epochs, number of learning iterations
//         number input nodes, number hidden, number out, learning rate
//         input-to-hidden weights, hidden-to-output weights
void saveAgent(NeuralNet* nn) {
	int inp,hid,out;
	FILE* fp;
	char filename[128];

	sprintf(filename,"agent%d.txt",nn->nepochs);
	fp=fopen(filename,"w");
	fprintf(fp,"%d %d %f\n",nn->id,nn->nepochs,nn->a.mse);
	fprintf(fp,"%d %d %d %f\n",nn->nin,nn->nhid,nn->nout,nn->lr);
	for(inp=0;inp<nn->nin;inp++) {
		for(hid=0;hid<nn->nhid;hid++) {
			fprintf(fp,"%e\t",nn->wih[inp][hid]);
		}
		fprintf(fp,"\n");
	}
	for(hid=0;hid<nn->nhid;hid++) {
		for(out=0;out<nn->nout;out++) {
			fprintf(fp,"%e\t",nn->who[hid][out]);
		}
		fprintf(fp,"\n");
	}
	fclose(fp);
}

//Allows the loading of agent's Q-table from a file
void loadAgent(char* filename, NeuralNet* nn) {
	int inp,hid,out;
	float mse;
	FILE* fp;
	fp=fopen(filename,"r");
	if(fp==NULL) {
		printf("  No such agent file: %s\n",filename);
		exit(-1);
	}

	fscanf(fp,"%d %d %f\n",&(nn->id),&(nn->nepochs),&mse);
	fscanf(fp,"%d %d %d %f\n",&(nn->nin),&(nn->nhid),&(nn->nout),&(nn->lr));
	initNN(nn);
	nn->a.mse=mse;
	
	for(inp=0;inp<nn->nin;inp++) {
		for(hid=0;hid<nn->nhid;hid++) {
			fscanf(fp,"%e\t",&(nn->wih[inp][hid]));
		}
	}
	for(hid=0;hid<nn->nhid;hid++) {
		for(out=0;out<nn->nout;out++) {
			fscanf(fp,"%e\t",&(nn->who[hid][out]));
		}
	}
	fclose(fp);
}

//Assign random weights to our neural network
void assignRandomWeights(NeuralNet* nn) {
	//#define GET_RANDOM_WEIGHT ( 0.02f * ( ((float)rand()/(float)RAND_MAX) - 0.5f ) )//[-.01,.01]
	#define GET_RANDOM_WEIGHT	( ((float)rand() / (float)RAND_MAX) - 0.5f)//[-0.5,0.5]
	int hid,inp,out;

	for(inp=0; inp<nn->nin+1; inp++) {
		for(hid=0; hid<nn->nhid; hid++) {
		  nn->wih[inp][hid] = GET_RANDOM_WEIGHT;
		}
	}

	for(hid=0; hid<nn->nhid+1; hid++) {
		for(out=0; out<nn->nout; out++) {
		  nn->who[hid][out] = GET_RANDOM_WEIGHT;
		}
	}
}

//Sigmoid activation function
float sigmoid(float x) {
	return (1.f / (1.f + (float)exp(-x)));
}
//Derivative of activation function
float sigmoidDerivative(float x)
{
	return (x * (1.f - x));
}

//Calculate activations throughout the neural network
void feedForward(NeuralNet* nn)
{
	int inp,hid,out;
	float sum;

	//Input to hidden nodes
	for(hid=0; hid<nn->nhid; hid++) {
		sum = 0.f;
		for (inp=0; inp<nn->nin; inp++) {
			sum += nn->a.in[inp] * nn->wih[inp][hid];
		}
		//Bias
		sum += nn->wih[nn->nin][hid];

		//Hidden node's activation
		nn->a.hid[hid] = sigmoid(sum);
	}

	//Input to output nodes
	for(out=0; out<nn->nout; out++) {
		sum = 0.f;
		for(hid=0; hid<nn->nhid; hid++) {
			sum += nn->a.hid[hid] * nn->who[hid][out];
		}
		//Bias
		sum += nn->who[nn->nhid][out];

		//Actual output node's activation
		nn->a.actual[out] = sigmoid( sum );
	}
}

//Correct weights so our neural network gets better
void backPropagate(NeuralNet* nn)
{
	int i,inp,hid,out;

	//Output node errors
	for(out=0; out<nn->nout; out++) {
		nn->a.errOut[out] = (nn->a.sup[out]-nn->a.actual[out]) * sigmoidDerivative(nn->a.actual[out]);
	}

	//Hidden node errors
	for(hid=0; hid<nn->nhid; hid++) {
		nn->a.errHid[hid]=0.f;
		for(out=0; out<nn->nout; out++) {
			nn->a.errHid[hid] += nn->a.errOut[out] * nn->who[hid][out];
		}
		nn->a.errHid[hid] *= sigmoidDerivative(nn->a.hid[hid]);
	}

	//Average weights for output
	if(WRITE_FILES==1) {
		for(i=0;i<nn->nhid;i++) {
			if(nn->a.errHid[i]>0) nn->a.sumErrHid[i]+=nn->a.errHid[i];
			else nn->a.sumErrHid[i]+=(-1.f*nn->a.errHid[i]);
		}
		for(i=0;i<nn->nout;i++) {
			if(nn->a.errOut[i]>0) nn->a.sumErrOut[i]+=nn->a.errOut[i];
			else nn->a.sumErrOut[i]+=(-1.f*nn->a.errOut[i]);
		}
	}

	//Update hidden-to-output weights
	for(out=0; out<nn->nout; out++) {
		for(hid=0; hid<nn->nhid; hid++) {
			nn->who[hid][out] += (nn->lr * nn->a.errOut[out] * nn->a.hid[hid]);
		}
		//Bias
		nn->who[nn->nhid][out] += (nn->lr * nn->a.errOut[out]);
	}

	//Update input-to-hidden weights
	for(hid=0; hid<nn->nhid; hid++) {
		for(inp=0; inp<nn->nin; inp++) {
			nn->wih[inp][hid] += (nn->lr * nn->a.errHid[hid] * nn->a.in[inp]);
		}
		//Bias
		nn->wih[nn->nin][hid] += (nn->lr * nn->a.errHid[hid]);
	}
}

//Select the most highly activated letter
int hardmax(int nout, float *vector)
{
	int index,sel;
	float max;
	sel = 0;
	max = vector[sel];

	for (index = 1 ; index < nout ; index++) {
		if (vector[index] > max) {
			max = vector[index]; sel = index;
		}
	}
	return(sel);
}

//Select the most highly activated letter probabilistically (derivative exists)
void softmax(int nout, float* vector) {
	int i;
	float sum;

	sum=0.f;
	for(i=0;i<nout;i++) {
		sum+=(float)exp(vector[i]);
	}
	for(i=0;i<nout;i++) {
		vector[i]=(float)(vector[i]/sum);
	}
}

//Main learn/validate/test loop for 1 epoch on all data
int onlineNN(NeuralNet* nn, Data* learn, int train) {
	int i,item,sum,guess,correct;
	float err;

	//Initialize data
	sum=0;
	nn->a.mse=0.f;
	if(train==1) nn->nepochs++;
	for(i=0;i<nn->nout;i++) nn->a.numCorrect[i]=0;

	if(train==0 && nn->a.guess!=NULL) free(nn->a.guess);
	nn->a.nguess=learn->nitems;
	nn->a.guess=(int*)malloc(nn->a.nguess*sizeof(int));

	//Main epoch loop
	for(item=0;item<learn->nitems;item++) {
		//Load input data for this letter
		for(i=0;i<learn->nin;i++) {
			nn->a.in[i] = learn->in[item][i];
			if(HALF_RESOLUTION==1) i++;
		}
		for(i=0;i<learn->nsup;i++) {
			nn->a.sup[i] = learn->sup[item][i];
		}

		//Hey neural network, classify this sucker
		feedForward(nn);

		//Tell exactly how many we get right
		guess=hardmax(nn->nout,nn->a.actual);
		correct=hardmax(nn->nout,nn->a.sup);
		if (guess==correct) {
			nn->a.numCorrect[correct]+=1;
			sum++;
		}
		nn->a.guess[item]=guess;

		//Calculate the error
		err = 0.f;
		for(i=0; i<nn->nout; i++) {
			err += ((nn->a.sup[i] - nn->a.actual[i])*(nn->a.sup[i] - nn->a.actual[i]));
		}
		err = 0.5f*err;
		nn->a.mse+=err;

		//If neural network is training, then adjust weights
		if(train==1) backPropagate(nn);
	}
	nn->a.mse=nn->a.mse/(float)(learn->nitems*learn->nsup);
	return sum;
}

