/*---------------------------------------------------------------------------*
 *                                   IT++			             *
 *---------------------------------------------------------------------------*
 * Copyright (c) 1995-2003 by Tony Ottosson, Thomas Eriksson, Pål Frenger,   *
 * Tobias Ringström, and Jonas Samuelsson.                                   *
 *                                                                           *
 * Permission to use, copy, modify, and distribute this software and its     *
 * documentation under the terms of the GNU General Public License is hereby *
 * granted. No representations are made about the suitability of this        *
 * software for any purpose. It is provided "as is" without expressed or     *
 * implied warranty. See the GNU General Public License for more details.    *
 *---------------------------------------------------------------------------*/

/*! 
  \file 
  \brief Vector quantizer training
  \author Thomas Eriksson
*/

#include "srccode/vqtrain.h"
#include "base/random.h"
#include "base/timing.h"
#include "base/elmatfunc.h"
#include "base/stat.h"
#include "base/matfunc.h"

namespace itpp {

// the cols contains codevectors
double kmeansiter(Array<vec> &DB, mat &codebook)
{
	int				DIM=DB(0).length(),SIZE=codebook.cols(),T=DB.length();
	vec				x,xnum(SIZE);
	mat				xsum(DIM,SIZE);
	int				n,MinIndex,i,j,k;
	double			MinS,S,D,Dold,*xp,*cp;
	
	xsum.clear();
	xnum.clear();

	n=0;
	D=1E20;
	Dold=D;
	D=0;
	for (k=0;k<T;k++) {
		x=DB(k);
		xp=x._data();
		MinS=1E20;
		MinIndex=0;
		for (i=0;i<SIZE;i++) {
			cp=&codebook(0,i);
			S=sqr(xp[0]-cp[0]);
			for (j=1;j<DIM;j++) {
				S+=sqr(xp[j]-cp[j]);
				if (S>=MinS) goto sune;
			}
			MinS=S;
			MinIndex=i;
sune:			i=i;
		}
		D+=MinS;
		cp=&xsum(0,MinIndex);
		for (j=0;j<DIM;j++) {
			cp[j]+=xp[j];
		}
		xnum(MinIndex)++;
	}
	for (i=0;i<SIZE;i++) {
		for (j=0;j<DIM;j++) {
			codebook(j,i)=xsum(j,i)/xnum(i);
		}
	}
	return D;
}

mat kmeans(Array<vec> &DB, int SIZE, int NOITER, bool VERBOSE)
{
	int				DIM=DB(0).length(),T=DB.length();
	mat				codebook(DIM,SIZE);
	int				n,i,j;
	double			D,Dold;
	ivec			ind(SIZE);
	
	for (i=0;i<SIZE;i++) {
		ind(i)=randi(0,T-1);
		j=0;
		while (j<i) {
			if (ind(j)==ind(i)) {
				ind(i)=randi(0,T-1);
				j=0;
			}
			j++;
		}
		codebook.set_col(i,DB(ind(i)));
	}
	

	if (VERBOSE) cout << "Training VQ..." << endl ;

	D=1E20;
	for (n=0;n<NOITER;n++) {
		Dold=D;
		D=kmeansiter(DB,codebook);
		if (VERBOSE) cout << n << ": " << D/T << " ";
		if (std::abs((D-Dold)/D) < 1e-4) break;
	}
	return codebook;
}

mat lbg(Array<vec> &DB, int SIZE, int NOITER, bool VERBOSE)
{
	int		S=1,DIM=DB(0).length(),T=DB.length(),i,n;
	mat		cb;
	vec		delta=0.001*randn(DIM),x;
	double	D,Dold;

	x=zeros(DIM);
	for (i=0;i<T;i++) {
		x+=DB(i);
	}
	x/=T;
	cb.set_size(DIM,1);
	cb.set_col(0,x);
	while (cb.cols()<SIZE) {
		S=cb.cols();
		if (2*S<=SIZE) cb.set_size(DIM,2*S,true);
		else cb.set_size(DIM,2*S,true);
		for (i=S;i<cb.cols();i++) {
			cb.set_col(i,cb.get_col(i-S)+delta);
		}
		D=1E20;
		for (n=0;n<NOITER;n++) {
			Dold=D;
			D=kmeansiter(DB,cb);
			if (VERBOSE) cout << n << ": " << D/T << " ";
			if (std::abs((D-Dold)/D) < 1e-4) break;
		}
	}

	return cb;
}

mat vqtrain(Array<vec> &DB, int SIZE, int NOITER, double STARTSTEP, bool VERBOSE)
{
	int				DIM=DB(0).length();
	vec				x;
	vec				codebook(DIM*SIZE);
	int				n,MinIndex,i,j;
	double			MinS,S,D,step,*xp,*cp;
	
	for (i=0;i<SIZE;i++) {
		codebook.replace_mid(i*DIM,DB(randi(0,DB.length()-1)));
	}
	if (VERBOSE) cout << "Training VQ..." << endl ;

res: D=0; 
	for (n=0;n<NOITER;n++) {
		step=STARTSTEP*(1.0-double(n)/NOITER);if (step<0) step=0;
		x=DB(randi(0,DB.length()-1)); // seems unnecessary! Check it up.
		xp=x._data();

		MinS=1E20;
		MinIndex=0;
		for (i=0;i<SIZE;i++) {
			cp=&codebook(i*DIM);
			S=sqr(xp[0]-cp[0]);
			for (j=1;j<DIM;j++) {
				S+=sqr(xp[j]-cp[j]);
				if (S>=MinS) goto sune;
			}
			MinS=S;
			MinIndex=i;
sune:		i=i;
		}
		D+=MinS;
		cp=&codebook(MinIndex*DIM);
		for (j=0;j<DIM;j++) {
			cp[j]+=step*(xp[j]-cp[j]);
		}
		if ((n%20000==0) && (n>1)) {
			if (VERBOSE) cout << n << ": " << D/20000 << " ";
			D=0;
		}
	}

// checking training result
	vec	dist(SIZE),num(SIZE);

	dist.clear();num.clear();
	for (n=0;n<DB.length();n++) {
		x=DB(n);
		xp=x._data();
		MinS=1E20;
		MinIndex=0;
		for (i=0;i<SIZE;i++) {
			cp=&codebook(i*DIM);
			S=sqr(xp[0]-cp[0]);
			for (j=1;j<DIM;j++) {
				S+=sqr(xp[j]-cp[j]);
				if (S>=MinS) goto sune2;
			}
			MinS=S;
			MinIndex=i;
sune2:		i=i;
		}
		dist(MinIndex)+=MinS;
		num(MinIndex)+=1;
	}
	dist=10*log10(dist*length(dist)/sum(dist));
	if (VERBOSE) cout << endl << "Distortion contribution: " << dist << endl ;
	if (VERBOSE) cout << "Num spread: " << num/DB.length()*100 << " %" << endl << endl ;
	if (min(dist)<-30) {
		cout << "Points without entries! Retraining" << endl ;
		j=min_index(dist);
		i=max_index(num);
		codebook.replace_mid(j*DIM,codebook.mid(i*DIM,DIM));
		goto res;
	}

	mat	cb(DIM,SIZE);
	for (i=0;i<SIZE;i++) {
		cb.set_col(i,codebook.mid(i*DIM,DIM));
	}
	return cb;
}
  


} // namespace itpp