#include <math.h>
#include "mex.h"

/* External Interfaces:
 *    http://www.mathworks.com/access/helpdesk/help/techdoc/matlab_external/index.html
 *
 * API Reference:
 *    http://www.mathworks.com/access/helpdesk/help/techdoc/apiref/
 *
 * Handling Sparse Arrays:
 *    http://www.mathworks.com/access/helpdesk/help/techdoc/matlab_external/f24370.html
 *    http://www.ccr.jussieu.fr/ccr/Documentation/Calcul/matlab5v11/docs/00009/009a1.htm
 */


/* If you are using a compiler that equates NaN to be zero, you
 * must compile this example using the flag  -DNAN_EQUALS_ZERO.
 * For example:
 *     mex -DNAN_EQUALS_ZERO fulltosparse.c
 * This will correctly define the IsNonZero macro for your C
 * compiler.  IEEE "NaN" is the only numeric constant that is not
 * equivalent with itself.
 */
#if defined(NAN_EQUALS_ZERO)
	#define isNonZero(d) ((d) != 0.0 || (d)!=(d))
#else
	#define isNonZero(d) ((d) != 0.0)
#endif


/* Metrics: */
double fisher ( int m, int pc, int pi[], double px[],
                       int qc, int qi[], double qx[] );
double pwmi   ( int m, int pc, int pi[], double px[],
                       int qc, int qi[], double qx[] );
double expmi  ( int m, int pc, int pi[], double px[],
                       int qc, int qi[], double qx[] );
double lrt    ( int m, int pc, int pi[], double px[],
                       int qc, int qi[], double qx[] );
double chi2   ( int m, int pc, int pi[], double px[],
                       int qc, int qi[], double qx[] );
double yuleq  ( int m, int pc, int pi[], double px[],
                       int qc, int qi[], double qx[] );
double cooccur( int m, int pc, int pi[], double px[],
                       int qc, int qi[], double qx[] );

typedef double (*Metric)( int m, int pc, int pi[], double px[],
                                 int qc, int qi[], double qx[] );
/* jvd: parameters used by metrics */
double smooth;	/* jvd: expmi, pwmi */

void allPairsColumns    ( int m, int n, Metric metric,
                          int Bi[], int Bj[], double Bx[],
                          int Ai[], int Aj[], double Ax[] );
int* contingencyTable   ( int c[], int m,
                          int pc, int pi[], double px[],
                          int qc, int qi[], double qx[] );


/* jvd: Matlab entry point.  "mexFunction" is the only function that makes
 * (extensive) use of "mex.h" */
void mexFunction( int nlhs,       mxArray *plhs[],
                  int nrhs, const mxArray *prhs[] ) {
	int m,n,nzmax,*Ai,*Aj;
	double *Ax;
	int *Bi,*Bj;
	double *Bx;
	Metric metric;
	void *newptr;
	char metname[32];

	/* jvd: provide brief explanation if no input arguments are given */
    if (nrhs < 1){
		mexPrintf(
			"allcolsmetric   Applies the symmetric metric between all pairs of sparse\n"
			"                columns and returns these results as a lower triangular sparse\n"
			"                matrix.  Choices are:\n"
			"                   'fisher'  Fisher distance\n"
			"                   'expmi'   Expected mutual information\n"
            "                                (optional smoothing parameter)\n"
			"                   'pwmi'    Point-wise mutual information\n"
            "                                (optional smoothing parameter)\n"
			"                   'lrt'     Likelihood ratio test\n"
			"                   'chi2'    Chi2 statistic\n"
			"                   'yuleq'   Yule's Q\n"
			"                   'cooccur' Lower tri (df) coocur\n\n"
			"                For example:\n"
			"                   X=sparse((rand(5,5)>.5).*rand(5,5));\n"
			"                   allcolsmetric(X,'expmi',1.0)\n"
		);
		return;
    }

	/* jvd: Check for proper number of input and output arguments.*/    
	if( nrhs < 2 || nrhs>3 )
		mexErrMsgTxt("At least two input arguments required.");

	/* jvd: Check data type of input argument. */
	if( !(mxIsSparse(prhs[0])) )
		mexErrMsgTxt("First input argument must be sparse.");
	if( !(mxIsDouble(prhs[0])) )
		mexErrMsgTxt("First input argument must be of type double.");
	if( mxGetNumberOfDimensions(prhs[0])!=2 )
		mexErrMsgTxt("First input argument must be two dimensional.");
	if( mxGetPi(prhs[0])!=NULL )
		mexErrMsgTxt("First input argument must be real.");
	if( !mxIsChar(prhs[1]) )
		mexErrMsgTxt("Second input argument must be name of valid metric.");

	mxGetString(prhs[1],metname,31);
	if( strncmp(metname,"fisher",32)==0 ) {
		metric = (Metric)fisher;
	} else if( strncmp(metname,"expmi",32)==0 ) {
		if( nrhs==3 ) {
			if( mxGetN(prhs[2])!=1 || mxGetM(prhs[2])!=1 || mxGetPi(prhs[2])!=NULL ) {
				mexErrMsgTxt("Optional paramter must be a real-valued scalar.");
				return;
			}
			smooth = mxGetScalar(prhs[2]);
		} else {
			smooth = 0.0;
		}
		metric = (Metric)expmi;
	} else if( strncmp(metname,"pwmi",32)==0 ) {
		if( nrhs==3 ) {
			if( mxGetN(prhs[2])!=1 || mxGetM(prhs[2])!=1 || mxGetPi(prhs[2])!=NULL ) {
				mexErrMsgTxt("Optional paramter must be a real-valued scalar.");
				return;
			}
			smooth = mxGetScalar(prhs[2]);
		} else {
			smooth = 0.0;
		}
		metric = (Metric)pwmi;
	} else if( strncmp(metname,"lrt",32)==0 ) {
		metric = (Metric)lrt;
	} else if( strncmp(metname,"chi2",32)==0 ) {
		metric = (Metric)chi2;
	} else if( strncmp(metname,"yuleq",32)==0 ) {
		metric = (Metric)yuleq;
	} else if( strncmp(metname,"cooccur",32)==0 ) {
		metric = (Metric)cooccur;
	} else {
		mexErrMsgTxt("Unrecognized metric.");
	}
	
	/* jvd: Get the size and pointers to input data. */
	m  = mxGetM ( prhs[0] );
	n  = mxGetN ( prhs[0] );
	Ax = mxGetPr( prhs[0] );
	Ai = mxGetIr( prhs[0] );
	Aj = mxGetJc( prhs[0] );
	/*cmplx = (Ai == NULL ? 0 : 1);*/

	/* jvd: Allocate space for upper triangular sparse matrix. */
	nzmax = n*(n+1)/2;
	plhs[0] = mxCreateSparse(n,n,nzmax,mxREAL);
	if( plhs[0]==NULL ) {
		mexErrMsgTxt("Unable to allocate sufficient heap memory for output.  "
				"Try reducing input size.");
		return;
	}

	/* jvd: Get the pointers to output data. */
	Bx = mxGetPr( plhs[0] );
	Bi = mxGetIr( plhs[0] );
	Bj = mxGetJc( plhs[0] );

	allPairsColumns(m,n,metric,Bi,Bj,Bx,Ai,Aj,Ax);
	
	/* jvd: cleanup unused memory.  Bj[n] is the actual number of non-zeros */
	if( Bj[n]==nzmax ) {
		/* jvd: do nothing */
	} else if( Bj[n]==0 ) {
		mxDestroyArray( plhs[0] );
		plhs[0] = mxCreateSparse(n,n,0,mxREAL);
	} else if( Bj[n]<nzmax ) {
		newptr = mxRealloc(Bi, sizeof(*Bi)*Bj[n]);
		mxSetIr(plhs[0], newptr);
		newptr = mxRealloc(Bx, sizeof(*Bx)*Bj[n]);
		mxSetPr(plhs[0], newptr);
	} else {
		mexErrMsgTxt("Fatal Error: More elements created than "
				"space allocated.");
	}

}


/* ------------------------------------------------------------------------- *
 * The following functions make no use of "mex.h"                            *
 *    except mexWarnMsgTxt mexErrMsgTxt and possibly mexPrintf
 * ------------------------------------------------------------------------- */


/* jvd: "all" pairs assumes the metric applied is symmetric */
/* jvd: the result of this function may leave unused memory at the tail of Ai
 *      and Ax */
void allPairsColumns( int m, int n, Metric metric,
                      int Bi[], int Bj[], double Bx[],
                      int Ai[], int Aj[], double Ax[]) {
	int c1,c2,pos=0;
	double dist;

	Bj[0] = 0;
	for( c1=0;c1<n;c1++ ) {
		/* jvd: update col length as we go. this is the maximum number of
		 *      elements in this column */
		Bj[c1+1] = Bj[c1] + n - c1;
		
		for( c2=c1;c2<n;c2++ ) {
			/* jvd: using the index of the beginning of each column,
			 *      hence Aj[] as indexer */
			/* jvd: TODO: would it be more efficient to pass first the list
			 * with the smaller ending row index?  or would it be faster to
			 * first pass the list with fewer non-zero entries (ie the less
			 * dense list)? */
			dist = metric( m, Aj[c1+1]-Aj[c1], &(Ai[Aj[c1]]), &(Ax[Aj[c1]]),
						      Aj[c2+1]-Aj[c2], &(Ai[Aj[c2]]), &(Ax[Aj[c2]])  );

			/* jvd: to maintain sparsity, we must ensure the result is only
			 * stored if it is nonzero */
			if( isNonZero(dist) ) {
				Bx[pos] = dist;
				Bi[pos] = c2;
				++pos;
			} else {
				/* jvd: there is now one less in this column than we
				 *      anticipated, make correction */
				--Bj[c1+1];
			}
		}
	}
}



int* contingencyTable( int c[], int m,
                       int pc, int pi[], double px[],
                       int qc, int qi[], double qx[] ) {
    int i=0, j=0;
	while( i<pc && j<qc ) {
        if( pi[i] < qi[j] ) {
			/* jvd: case 1,0 */
			if( px[i]!=0 ) {
				c[2]+=1;
			} else {
				mexErrMsgTxt("contingencyTable(): shouldn't be here (1).");
				return NULL;
			}
            ++i;
            continue;
        }
		/* jvd: case 1,1 */
		if( pi[i]==qi[j] && px[i]!=0 && qx[j]!=0 ) {
			c[3]+=1;
			++i;
		/* jvd: case 0,1 */
		} else if( qx[j]!=0 ) {
			c[1]+=1;
		} else {
			mexErrMsgTxt("contingencyTable(): shouldn't be here (2).");
			return NULL;
		}
		++j;
	}
	/* jvd: finish up the remaining elements */
	c[1] += qc-j;	/* jvd: case 0,1 */
	c[2] += pc-i;	/* jvd: case 1,0 */
	
	/* jvd: case 0,0 */
	c[0] = (m - c[1] - c[2] - c[3]);

	/*mexPrintf("%d\t%d\t%d\t%d\n",c[0],c[1],c[2],c[3]);*/
	
	return c;
}


#define MI(PX,PY,PXY) ( ((PXY)==0.0 || (PX)==0.0 || (PY)==0.0)?0.0:log((PXY)/((PX)*(PY)))/log(2.0) )

double pwmi( int m, int pc, int pi[], double px[],
                    int qc, int qi[], double qx[] ) {
	int contTable[] = {0,0,0,0}; /* jvd: contingency table (histogram) */
	/* jvd: jt density has four events, hence add 4*c to normalization */
	double Z = m + 4.0*smooth;
	double P_X1, P_Y1, P_XY[4];

	/* jvd: 0:XY=(0,0), 1:XY=(0,1), 2:XY=(1,0), 3:XY=(1,1)  */
	contingencyTable( contTable, m, pc, pi, px, qc, qi, qx );

	/* jvd: normalize the joint */
	P_XY[0] = (double)(contTable[0]+smooth)/Z;
	P_XY[1] = (double)(contTable[1]+smooth)/Z;
	P_XY[2] = (double)(contTable[2]+smooth)/Z;
	P_XY[3] = (double)(contTable[3]+smooth)/Z;

	/* jvd: compute the marginals (superfluous but convenient) */
	P_X1 = P_XY[2] + P_XY[3];
	P_Y1 = P_XY[1] + P_XY[3];

	/* jvd: make base 2 */
	return MI(P_X1,P_Y1,P_XY[3]);	/* 3: X=1, Y=1 */
}



double expmi( int m, int pc, int pi[], double px[],
                     int qc, int qi[], double qx[] ) {
	int contTable[] = {0,0,0,0}; /* jvd: contingency table (histogram) */
	/* jvd: jt density has four events, hence add 4*c to normalization */
	double Z = m + 4.0*smooth;
	double P_X[2], P_Y[2], P_XY[4];
	double sum = 0.0;

	/* jvd: 0:XY=(0,0), 1:XY=(0,1), 2:XY=(1,0), 3:XY=(1,1)  */
	contingencyTable( contTable, m, pc, pi, px, qc, qi, qx );

	/* jvd: normalize the joint */
	P_XY[0] = (double)(contTable[0]+smooth)/Z;
	P_XY[1] = (double)(contTable[1]+smooth)/Z;
	P_XY[2] = (double)(contTable[2]+smooth)/Z;
	P_XY[3] = (double)(contTable[3]+smooth)/Z;

	/* jvd: compute the marginals (superfluous but convenient) */
	P_X[0] = P_XY[0] + P_XY[1];
	P_X[1] = P_XY[2] + P_XY[3];
	P_Y[0] = P_XY[0] + P_XY[2];
	P_Y[1] = P_XY[1] + P_XY[3];

	sum = 0.0;
	sum += P_XY[0] * MI(P_X[0],P_Y[0],P_XY[0]); /* 0: X=0, Y=0 */
	sum += P_XY[1] * MI(P_X[0],P_Y[1],P_XY[1]); /* 1: X=0, Y=1 */
	sum += P_XY[2] * MI(P_X[1],P_Y[0],P_XY[2]); /* 2: X=1, Y=0 */
	sum += P_XY[3] * MI(P_X[1],P_Y[1],P_XY[3]); /* 3: X=1, Y=1 */
		
	return sum;
}



double cooccur( int m, int pc, int pi[], double px[],
                       int qc, int qi[], double qx[] ) {
	int contTable[] = {0,0,0,0}; /* jvd: contingency table (histogram) */
	/* jvd: 0:XY=(0,0), 1:XY=(0,1), 2:XY=(1,0), 3:XY=(1,1)  */
	contingencyTable( contTable, m, pc, pi, px, qc, qi, qx );
	return contTable[3];
}



double lrt( int m, int pc, int pi[], double px[],
                   int qc, int qi[], double qx[] ) {
	/* jvd: Hypothesis Test Variables */
	int    c1,c2, c12;
	double p, p1, p2;
	double loglambda;
	int    N=m;
	
	/* jvd: misc variables */
	int contTable[] = {0,0,0,0}; /* jvd: contingency table (histogram) */
	double tmp = 0.0;
	
	/* jvd: 0:XY=(0,0), 1:XY=(0,1), 2:XY=(1,0), 3:XY=(1,1)  */
	contingencyTable( contTable, m, pc, pi, px, qc, qi, qx );

	/* Hypothesis Test Variables */
	c1  = contTable[2]+contTable[3];  /* jvd: c1  := XY=(1,0)+XY=(1,1) */
	c2  = contTable[1]+contTable[3];  /* jvd: c2  := XY=(0,1)+XY=(1,1) */
	c12 = contTable[3];	              /* jvd: c12 := XY=(1,1) */
	p   = (double)c2/N;
	p1  = (double)c12/c1;
	p2  = (double)(c2-c12)/(N-c1);

	/* log{ L(k,n,x) } = log{ pow(x,k) * pow(1-x,n-k) }  */
	/*                 = k*log(x) + (n-k)*log(1-x)       */
	#define logL(k,n,x) ( ((x)!=(x)||(x)==0.0||(x)==1.0)?0.0:(double)(k)*log(x) + (double)((n)-(k))*log(1.0-(x)) )

	loglambda = logL(c12,c1,p) + logL(c2-c12,N-c1,p) - logL(c12,c1,p1) - logL(c2-c12,N-c1,p2);

	return -2.0*loglambda;
}



double fisher( int m, int pc, int pi[], double px[],
                      int qc, int qi[], double qx[] ) {
	/* jvd: this is the pre acos identity value */
	double sum=0.0;
    int i=0,j=0;
	
	if( pc==0 || qc==0 ) return acos( sum );

    /* jvd: NOTE: assuming valid probability vector */
    while( i<pc && j<qc ) {
        if( pi[i] < qi[j] ) {
            ++i;
            continue;
        }
        if( pi[i] == qi[j] ) {
			sum += sqrt( px[i]*qx[j] );
		}
        ++j;
    }

	/* jvd: handles problems incurred from machine (in)precision (and deals
	 *      with invalid probability vectors) */
    if( sum>1.0 ) {
        sum = 1.0;
    } else if( sum<0.0 ) {
        sum = 0.0;
	} else {
		/* jvd: do nothing */
	}

	return acos(sum);
}



double chi2 ( int m, int pc, int pi[], double px[],
                     int qc, int qi[], double qx[] ) {
    int contTable[] = {0,0,0,0}; /* jvd: contingency table (histogram) */
    /* jvd: jt density has four events, hence add 4*c to normalization */
    double a,b,c,d;
    double n;
	double num,den;

    /* jvd: 0:XY=(0,0), 1:XY=(0,1), 2:XY=(1,0), 3:XY=(1,1)  */
    contingencyTable( contTable, m, pc, pi, px, qc, qi, qx );

    /* Yi: convential way of naming variables */
    a = contTable[3]; b = contTable[2];
    c = contTable[1]; d = contTable[0];
    n = a + b + c + d;	/* jvd: equivalent to m */
    /* mexPrintf("%d %d %d %d %d\n", a,b,c,d,n); */

    /* Yi: chi2 statistics */
	num = n*(a*d-c*b)*(a*d-c*b);
	den = (a+c)*(b+d)*(a+b)*(c+d);
	if( den==0 ) {
		mexWarnMsgTxt("Warning: divide by zero.");
	}
    return num/den;
}



double yuleq( int m, int pc, int pi[], double px[],
                     int qc, int qi[], double qx[] ) {
	int contTable[] = {0,0,0,0}; /* jvd: contingency table (histogram) */
	/* jvd: jt density has four events, hence add 4*c to normalization */
	double a,b,c,d;
	double num,den;

	/* jvd: 0:XY=(0,0), 1:XY=(0,1), 2:XY=(1,0), 3:XY=(1,1)  */
	contingencyTable( contTable, m, pc, pi, px, qc, qi, qx );

	/* Yi: convential way of naming variables */
	a = contTable[3]; b = contTable[2];
	c = contTable[1]; d = contTable[0];
	/* mexPrintf("%d %d %d %d %d\n", a,b,c,d,n); */

	/* Yi: Yule's Q statistics */
	num = a*d-b*c;
	den = a*d+b*c;
	if( den==0 ) {
		mexWarnMsgTxt("Warning: divide by zero.");
	}
	return num/den;
}