#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <time.h>
#include <math.h>
#include <float.h>
#include "defs.h"

#include "cover_tree.h"

#include "par_est.h"
#include "per_eval.h"
#include "fast_lsvm.h"
#include "svm.h"
#include "knn.h"

sFastLSVM::sFastLSVM(struct lkm_parameter *par, struct problem *prob)
{
	
	this->par = par;
	coverTree = NULL;

	if( par->unbalanced == 1 )
		{
		prob_pl = Malloc(struct problem,1);

		prob_pl->l = 0;
		for( int  i = 0; i < prob->l; i++ )
			if( prob->y[i] == 1 )
				prob_pl->l++;

		prob_pl->x = Malloc(struct node*,prob_pl->l);
		prob_pl->y = Malloc(double,prob_pl->l);


		int ppi = 0;
		for( int  i = 0; i < prob->l; i++ )
			if( prob->y[i] == 1 )
				{
				prob_pl->x[ppi] = prob->x[i];
				prob_pl->y[ppi] = prob->y[i];
				ppi++;
				}

		}

	this->prob = prob;

	

	parEst = new ParEst(par->opt);
}

sFastLSVM::sFastLSVM(int pred_strategy)
{
	coverTree = NULL;
	parEst = new ParEst(pred_strategy);
	this->pred_strategy = pred_strategy;
}

sFastLSVM::~sFastLSVM()
{
	delete parEst;
	if( coverTree != NULL ) delete coverTree;
}

void sFastLSVM::cross_validation_int(const problem *prob, int nr_fold, double *target)
{
	int i;
	int *fold_start = Malloc(int,nr_fold+1); if( fold_start == NULL ) MallocError("fold_start");
	int l = prob->l;
	int *perm = Malloc(int,l); if( perm == NULL ) MallocError("perm");
	int nr_class;


	int *start = NULL;
	int *label = NULL;
	int *count = NULL;
	svm_group_classes(prob,&nr_class,&label,&start,&count,perm);

	// random shuffle and then data grouped by fold using the array perm
	int *fold_count = Malloc(int,nr_fold); if( fold_count == NULL ) MallocError("fold_count");
	int c;
	int *index = Malloc(int,l); if( index == NULL ) MallocError("index");
	for(i=0;i<l;i++)
		index[i]=perm[i];
	for (c=0; c<nr_class; c++) 
		for(i=0;i<count[c];i++)
		{
			int j = i+rand()%(count[c]-i);
			swap(index[start[c]+j],index[start[c]+i]);
		}
	for(i=0;i<nr_fold;i++)
	{
		fold_count[i] = 0;
		for (c=0; c<nr_class;c++)
			fold_count[i]+=(i+1)*count[c]/nr_fold-i*count[c]/nr_fold;
	}
	fold_start[0]=0;
	for (i=1;i<=nr_fold;i++)
		fold_start[i] = fold_start[i-1]+fold_count[i-1];
	for (c=0; c<nr_class;c++)
		for(i=0;i<nr_fold;i++)
		{
			int begin = start[c]+i*count[c]/nr_fold;
			int end = start[c]+(i+1)*count[c]/nr_fold;
			for(int j=begin;j<end;j++)
			{
				perm[fold_start[i]] = index[j];
				fold_start[i]++;
			}
		}
	fold_start[0]=0;
	for (i=1;i<=nr_fold;i++)
		fold_start[i] = fold_start[i-1]+fold_count[i-1];
	free(start);	
	free(label);
	free(count);	
	free(index);
	free(fold_count);


	for(i=0;i<nr_fold;i++)
	{
		if( !par->silent ) printf("\n\n====== Fold %d ======================================================================\n",i);
		int begin = fold_start[i];
		int end = fold_start[i+1];
		int j,k;
		struct problem subprob;

		subprob.l = l-(end-begin);
		subprob.x = Malloc(struct node*,subprob.l); if( subprob.x == NULL ) MallocError("subprob.x");
		subprob.y = Malloc(double,subprob.l); if( subprob.y == NULL ) MallocError("subprob.y");
			
		k=0;
		for(j=0;j<begin;j++)
		{
			subprob.x[k] = prob->x[perm[j]];
			subprob.y[k] = prob->y[perm[j]];
			struct node *n = subprob.x[k];
			while( n->index > 0 ) n++;
			n->index = -k;
			++k;
		}
		for(j=end;j<l;j++)
		{
			subprob.x[k] = prob->x[perm[j]];
			subprob.y[k] = prob->y[perm[j]];
			struct node *n = subprob.x[k];
			while( n->index > 0 ) n++;
			n->index = -k;
			++k;
		}

		sFastLSVM *flmf = new sFastLSVM(par,&subprob);
		flmf->train();

		double *tmp_res = Malloc(double,end-begin);
		double *tmp_c = Malloc(double,end-begin);

		for(j=begin;j<end;j++) 
			{
			if( par->svm_par->probability == 0 ) target[perm[j]] = flmf->predict(prob->x[perm[j]]);
			else target[perm[j]] = flmf->predict_p(prob->x[perm[j]]);
			tmp_res[j-begin] = target[perm[j]];
			tmp_c[j-begin] = prob->y[perm[j]];
			}

		PerEval *perEval = new PerEval(tmp_c,end-begin);
		if( par->per_meas < 4 ) perEval->calcAll(tmp_res,par->per_meas);
		if( par->per_meas == 4 ) perEval->calcLift20P(tmp_res,true);
		
		free(tmp_res);
		free(tmp_c);
		delete flmf;


		free(subprob.x);
		free(subprob.y);
		if( !par->silent ) printf("====== END Fold %d ==================================================================\n",i);
	}		
	free(fold_start);
	free(perm);	
}

double* sFastLSVM::cross_validation()
{
	double * target = Malloc(double,prob->l);
	cross_validation_int(prob, par->nr_fold, target);
	return target;
}



void sFastLSVM::train()
{
	coverTree = new CoverTree(par->svm_par,par->base,false,0,par->silent);

	int n_kps = par->n_k_lms;
	int *KPs = Malloc(int,n_kps);
	for( int i = 0; i < n_kps; i++ )
		KPs[i] = (int)(((double)par->k_lms[i])*par->kp_ratio);
		
	ct_prob = coverTree->addPoints(prob);

	if( par->local_model_selection != NO_LMS ) 
		{
		struct problem **nlp = NULL;
		
		int kkk = par->k_lms[par->n_k_lms-1];
		int kkp = (int)(((double)par->k_lms[0])*par->kp_ratio);
		nlp = coverTree->createNLocalProbs(par->n_loc_models,kkk,kkp,par->loc_prob_bil_th);

		if( par->svm_par->kernel_type == LINEAR )
			{
			double G_fake[1] = {1.0};
			parEst->estCGK(nlp,par->n_loc_models,par,par->c_lms,par->n_c_lms,G_fake,1,par->k_lms,par->n_k_lms,KPs,n_kps);
			}
		if( par->svm_par->kernel_type == RBF )
			parEst->estCGK(nlp,par->n_loc_models,par,par->c_lms,par->n_c_lms,par->g_lms,par->n_g_lms,par->k_lms,par->n_k_lms,KPs,n_kps);
		
		for( int i = 0; i < par->n_loc_models; i++ )
			free(nlp[i]);
		free(nlp);
		}

	coverTree->setKs(par->k,par->kp);
	coverTree->createLocalProbs();

	int n = 0;

	bool est_gamma = false;
	double quant = -par->svm_par->gamma;
	if(  par->svm_par->kernel_type == RBF && par->svm_par->gamma < 0.0 ) est_gamma = true;

	int h_warning = 0;
	int shrinked = 0;
	double h_warning_th = 0.5;

	struct problem **probs = coverTree->get_trained_probs();

	for( int i = 0; i < coverTree->get_n_local_probs(); i++ )
		{
		struct problem *lp = probs[i];

		if( par->unbalanced == 1 )
			{
			lp = Malloc(struct problem,1);
			lp->l = probs[i]->l;
			lp->x = Malloc(struct node*,lp->l);
			lp->y = Malloc(double,lp->l);

			for( int h = 0; h < lp->l; h++)
				{
				lp->x[h] = probs[i]->x[h];
				lp->y[h] = probs[i]->y[h];
				}

			for( int h = 0; h < prob_pl->l; h++ )
				{
				struct node *nod = prob_pl->x[h];

				bool found = false;
				for( int s = 0; s < lp->l; s++ )
					if( lp->y[s] == 1 && lp->x[s] == nod )
						found = true;

				if( !found )
					{
					lp->l++;
					lp->x = Realloc(lp->x,struct node*,lp->l);
					lp->y = Realloc(lp->y,double,lp->l);
					lp->x[lp->l-1] = nod;
					lp->y[lp->l-1] = prob_pl->y[h];
					}
				}
			}

		struct svm_parameter *parameter = copy_svm_parameter(par->svm_par);

		if( est_gamma )	parameter->gamma = parEst->estG_quant(lp,quant);
		if( par->local_model_selection == LMS_ALL )
			{
			parEst->estCG(lp,parameter,par->c_lms,par->n_c_lms,par->g_lms,par->n_g_lms,par->kp/2,par->k);
			parameter->gamma = parEst->estG_quant(lp,-parameter->gamma);
			}

		int www = 0;
		if( parameter->weight != NULL && parameter->weight[0] < 0.0 && parameter->weight_label[0] == 1.0 )
			{
			www = 1;
			int np = 0;
			for( int u = 0; u < lp->l; u++ )
				if( lp->y[u] == 1.0 )
					np++;
			double cp = 1.0;
			if( np > 1 ) cp = ((double)lp->l)/((double)np)*parameter->weight[0];
			if( cp > 256.0 ) cp = 256.0;
			parameter->weight[0] = cp;
			}

		if( i > 15 && (double)h_warning / (double)shrinked > h_warning_th )
			{
			parameter->shrinking = 0;
			if( !par->silent )
				{
				if( www ) printf("train the %d local SVM model [c %f, cw1 %f, g %f --> %f, tot %d, size %d - shrinking deactivated]: ",(n++)+1,par->svm_par->C,parameter->weight[0],par->svm_par->gamma,parameter->gamma,coverTree->get_n_local_probs(),lp->l);
				else printf("train the %d local SVM model [c %f, g %f --> %f, tot %d, size %d - shrinking deactivated]: ",(n++)+1,par->svm_par->C,par->svm_par->gamma,parameter->gamma,coverTree->get_n_local_probs(),lp->l);
				}
			}
		else 
			{
			if( !par->silent )
				{
				if( www ) printf("train the %d local SVM model [c %f, cw1 %f, g %f --> %f, tot %d, size %d]: ",(n++)+1,par->svm_par->C,parameter->weight[0],par->svm_par->gamma,parameter->gamma,coverTree->get_n_local_probs(),lp->l);	
				else printf("train the %d local SVM model [c %f, g %f --> %f, tot %d, size %d]: ",(n++)+1,par->svm_par->C,par->svm_par->gamma,parameter->gamma,coverTree->get_n_local_probs(),lp->l);
				}
			}

		
		struct svm_model *m = svm_train(lp,parameter);

		if( par->unbalanced == 1 )
			{
			free(lp->x);
			free(lp->y);
			free(lp);
			}	


		if( !par->silent ) printf("\n");

		if( parameter->shrinking )
			{
			shrinked++;
			if( parameter->h_warning ) h_warning++;
			}
		coverTree->set_model(m,i);
		}
}

void sFastLSVM::save(char *file)
{
	coverTree->save(file);
}

void sFastLSVM::load(char *file)
{
	coverTree = new CoverTree(NULL,-1,false,pred_strategy,1);
	coverTree->load(file);
	par = Malloc(struct lkm_parameter,1);
	par->k = coverTree->getK();
	par->kp = coverTree->getKP();
	par->local_model_selection = 0;
	par->svm_par = coverTree->getParam();	
}

double* sFastLSVM::predict(struct problem *test)
{
	double *res = Malloc(double,test->l); if( res == NULL ) MallocError("res");
	for( int i = 0; i < test->l; i++ )
		{
		struct svm_model *m = coverTree->nn_model(test->x[i]);
		res[i] = svm_predict(m,test->x[i]);
		}
	return res;	
}

double* sFastLSVM::predict_p(struct problem *test)
{
	double *res = Malloc(double,test->l); if( res == NULL ) MallocError("res");
	for( int i = 0; i < test->l; i++ )
		res[i] = predict_p(test->x[i]);
	return res;	
}


double sFastLSVM::predict_p(struct node *test)
{
	struct svm_model *m = coverTree->nn_model(test);

	int nr_class=svm_get_nr_class(m);

	int *labels=(int *) malloc(nr_class*sizeof(int));
	svm_get_labels(m,labels);
	double *prob_estimates = (double *) malloc(nr_class*sizeof(double));

	int ind_an = -1;
	for(int j=0;j<nr_class;j++)
		if( labels[j] == 1 )
			ind_an = j;//printf(" %d",labels[j]);

	free(labels);

	svm_predict_probability(m,test,prob_estimates);
	double res = prob_estimates[ind_an];

	free(prob_estimates);

	return res;
}

double sFastLSVM::predict(struct node *test)
{
	struct svm_model *m = coverTree->nn_model(test);
	return svm_predict(m,test);
}


double* sFastLSVM::predict(struct problem *test, int k)
{
	double *res = Malloc(double,test->l); if( res == NULL ) MallocError("res");

	int pc1 = test->l / 1000;
	int nn = 0;

	double *loc = Malloc(double,k);

	for( int i = 0; i < test->l; i++ )
		{
		struct svm_model **ms = coverTree->knn_models(test->x[i],k);

		for( int j = 0; j < k; j++ )
			loc[j] = svm_predict(ms[j],test->x[i]);

		res[i] = KNN::majorityRuleBin(loc,k);

		if( i == pc1 * nn  ) 
			{
			printf("%d point tested [%.1f%%]\n",i,((double)i)/(double(test->l))*100.0);
			nn++;
			}

		free(ms);
		}
	return res;
}
