/*
 The code can be release under the GNU Lesser General Pulic License (LGPL)
 Please do not use fastem without understanding how it works. Please see the reference for more details
  
 Yi Zhang (University of California, Santa Cruz)
 Wei Xu (NEC Research America)
 Reference:
 Yi Zhang, Wei Xu, Fast Exact Maximum Likelihood Estimation for Mixture of Language Models, Poster in Proceedings of the 30st Annual International ACM SIGIR Conference on Research and Development in Information Retrieval, Amsterdam, Netherland

*/#include <windows.h>
#include <stdio.h>
#include <algorithm>

double get_time()
{
	__int64 count;
	static __int64 freq=0,base=0;
	if (freq==0) {
		QueryPerformanceFrequency((LARGE_INTEGER*)&freq);
		QueryPerformanceCounter((LARGE_INTEGER*)&base);
	}
	QueryPerformanceCounter((LARGE_INTEGER*)&count);
	return (double)1000.0*(double)(count-base)/(double)freq;
}

void fastmix(int n,const double* p,const double* f,double* q,double alpha,double beta)
{
	double gamma=beta/alpha;
	int i;
	int t=-1;
	double Sf=0,Sp=0;
	int start=0,end=n;
	for (i=0;i<n;i++) {
		q[i]=i;
	}
	do {
		int s=(int)q[(start+end)/2];
		int a=start;
		int b=end-1;
		double Sf_new=Sf+f[s],Sp_new=Sp+p[s];
		std::swap(q[a],q[(start+end)/2]);
		int j=start+1;
		while(j<=b) {
			int i=(int)q[j];
			if (f[i]*p[s]>f[s]*p[i]) {
				Sf_new+=f[i];
				Sp_new+=p[i];
				std::swap(q[a],q[j]);
				a++;
				j++;
			}
			else if (f[i]*p[s]<f[s]*p[i]) {
				std::swap(q[b],q[j]);
				b--;
			}
			else {
				Sf_new+=f[i];
				Sp_new+=p[i];
				j++;
			}
		}
		if ((gamma+Sp_new)*f[s]-Sf_new*p[s]>0) {
			start=b+1;
			Sp=Sp_new;
			Sf=Sf_new;
			t=s;
		}
		else {
			end=a;
		}
	} while(end>start);
	gamma=1/gamma;
	double lambda=(1+gamma*Sp)/Sf;
	for (i=0;i<n;i++) {
		if (f[i]*p[t]>=f[t]*p[i]) q[i]=lambda*f[i]-gamma*p[i];
		else q[i]=0;
	}
}

void em(int n,const double* p,const double* f,double* q,double alpha,double beta)
{
	int i=0;
	double Sq=0;
	for (i=0;i<n;i++) {
		Sq+=q[i]=f[i]*beta*q[i]/(alpha*p[i]+beta*q[i]);
	}
	Sq=1/Sq;
	for (i=0;i<n;i++) {
		q[i]*=Sq;
	}
}

void main()
{
	const int n=5;
	double p[n],f[n],q[n],qq[n];
	int i;
	double t_em=0,t_fast=0;

	int n_try=100;
	for (int tr=0;tr<n_try;tr++) {
		double Sp=0,Sf=0;
		for (i=0;i<n;i++) {
			Sp+=p[i]=rand();
			Sf+=f[i]=rand();
		}
		Sp=1/Sp;
		Sf=1/Sf;
		for (i=0;i<n;i++) {
			p[i]*=Sp;
			f[i]*=Sf;
			q[i]=f[i];
		}
		double t0;
		t0=get_time();
		fastmix(n,p,f,qq,0.99,0.01);
		t_fast=get_time()-t0;

		for (int j=0;j<1;j++) {
			t0=get_time();
			em(n,p,f,q,0.99,0.01);
			t_em=get_time()-t0;
			double d=0;
			for (i=0;i<n;i++){
				d+=(q[i]-qq[i])*(q[i]-qq[i]);
			}
//			printf("%d %lg\n",j,d);
		}


	}
	printf("n=%d\n",n);
	printf("EM: %lg\n",t_em/n_try);
	printf("FAST EM: %lg\n",t_fast/n_try);
}