package learn.struct.it;

import java.util.BitSet;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import util.math.Permutations;
import dataset.Dataset;

/**
 * A Chi-Square Independence Test. This test is applied when you have two
 * categorical variables from a single population. It is used to determine
 * whether there is a significant association between the two variables.
 * 
 * @author fschluter
 *
 */
public class ChiSquare implements StatisticalIndependenceTest {
	private static double gamser, gammcf, gln;

	private static double FPMIN = 1.0e-30;

	private static int ITMAX = 2000;

	private static double EPS = 3.0e-7;

	private double df;

	/**
	 * Method to compute the P-Value
	 * 
	 * @return P-Value
	 */
	public double computeLogPValue(List<Integer> xList, List<Integer> yList,
			List<Integer> z, Dataset dataset) {

		int x = xList.get(0);
		int y = yList.get(0);
		HashMap<Integer, int[][]> contTables = computeContingencyTables(x, y,
				z, dataset);

		long numValX = dataset.getCardinalities(x);
		long numValY = dataset.getCardinalities(y);
		int numSlices = contTables.size();
		df = (numValX * numValY - 1) * numSlices;

		/** Computes conditional chi square value */
		double chsq = 0.0;

		for (int[][] ct : contTables.values())
			chsq += chiVal(ct);

		double logPValue = chsq == 0 ? 0.0 : ln_gammq(.5 * df, 0.5 * chsq);

		return logPValue;
	}

	/**
	 * Computes the Chi-Value
	 * 
	 * @param matrix
	 * @return Chi-Value
	 */
	protected double chiVal(int[][] matrix) {
		int nrows, ncols, row, col;
		double[] rtotal, ctotal;
		double expect = 0, chival = 0, n = 0;

		nrows = matrix.length;
		ncols = matrix[0].length;
		rtotal = new double[nrows];
		ctotal = new double[ncols];
		for (row = 0; row < nrows; row++) {
			for (col = 0; col < ncols; col++) {
				rtotal[row] += matrix[row][col];
				ctotal[col] += matrix[row][col];
				n += matrix[row][col];
			}
		}

		chival = 0.0;
		for (row = 0; row < nrows; row++) {
			for (col = 0; col < ncols; col++) {
				if (ctotal[col] == 0 || rtotal[row] == 0) {
					df--;
					continue;
				}
				expect = (ctotal[col] * rtotal[row]) / n;

				if (expect != 0)
					chival += computeCell(matrix[row][col], expect);
			}
		}

		return chival;
	}

	// Oij is the frequency count and Eij is the null-hypthesis (independence)
	// model expected value
	/**
	 * Method to compute the Cell
	 * 
	 * @param Oij
	 *            is the frequency count
	 * @param Eij
	 *            is the null-hypthesis (independence)
	 * @return Cell
	 */
	protected double computeCell(double Oij, double Eij) {
		return Math.abs(Oij - Eij) * Math.abs(Oij - Eij) / Eij;
	}

	// TODO: completar javadoc
	/**
	 * 
	 * @param df
	 * @param x
	 * @return
	 */
	public static double gammq(double df, double x) {
		if (x < 0.0 || df <= 0.0) {
			System.err.println("Invalid arguments in routine gammq()");
			(new Exception()).printStackTrace();
			System.exit(0);
		}

		if (x < (df + 1.0)) {
			if (x == 0)
				return 1;

			gser(df, x);
			return 1.0 - gamser;
		} else {
			gcf(df, x);
			return gammcf;
		}
	}

	// TODO: completar javadoc
	/**
	 * 
	 * @param df
	 * @param x
	 */
	static void gser(double df, double x) {
		int n;
		double sum, del, ap;

		gln = gammln(df);
		if (x < 0.0) {
			System.err.println("x less than 0 in routine gser()");
			return;
		} else {
			ap = df;
			del = sum = 1.0 / df;
			for (n = 1; n <= ITMAX; n++) {
				++ap;
				del *= x / ap;
				sum += del;
				if (Math.abs(del) < Math.abs(sum) * EPS) {
					gamser = sum * Math.exp(-x + df * Math.log(x) - (gln));
					return;
				}
			}
			System.err
					.println("a too large, ITMAX too small in routine gser()");
			return;
		}
	}

	// TODO: completar javadoc
	/**
	 * 
	 * @param a
	 * @param x
	 */
	static void gcf(double a, double x) {
		int i;
		double an, b, c, d, del, h;

		gln = gammln(a);
		b = x + 1.0 - a;
		c = 1.0 / FPMIN;
		d = 1.0 / b;
		h = d;
		for (i = 1; i <= ITMAX; i++) {
			an = -i * (i - a);
			b += 2.0;
			d = an * d + b;
			if (Math.abs(d) < FPMIN)
				d = FPMIN;
			c = b + an / c;
			if (Math.abs(c) < FPMIN)
				c = FPMIN;
			d = 1.0 / d;
			del = d * c;
			h *= del;
			if (Math.abs(del - 1.0) < EPS)
				break;
		}
		if (i > ITMAX)
			System.err.println("a too large, ITMAX too small in gcf()");
		gammcf = Math.exp(-x + a * Math.log(x) - (gln)) * h;
	}

	// TODO: completar javadoc
	/**
	 * 
	 * @param xx
	 * @return
	 */
	static double gammln(double xx) {
		double x, y, tmp, ser;
		double cof[] = { 76.18009172947146, -86.50532032941677,
				24.01409824083091, -1.231739572450155, 0.1208650973866179e-2,
				-0.5395239384953e-5 };
		int j;

		y = x = xx;
		tmp = x + 5.5;
		tmp -= (x + 0.5) * Math.log(tmp);
		ser = 1.000000000190015;
		for (j = 0; j <= 5; j++)
			ser += cof[j] / ++y;
		return -tmp + Math.log(2.5066282746310005 * ser / x);
	}

	// TODO: completar javadoc
	/**
	 * 
	 * @param df
	 * @param x
	 * @return
	 */
	public static double ln_gammq(double df, double x) {

		if (x < 0.0 || df <= 0.0) {
			System.err.println("Invalid arguments in routine gammq()");
		}
		if (x < (df + 1.0)) {
			if (x == 0)
				return 0;

			return Math.log(-Math.expm1(ln_gser(df, x)));
		} else {
			return ln_gcf(df, x);
		}
	}

	// TODO: completar javadoc
	/**
	 * 
	 * @param df
	 * @param x
	 * @return
	 */
	static double ln_gser(double df, double x) {
		int n;
		double sum, del, ap;

		gln = gammln(df);
		if (x < 0.0) {
			System.err.println("x less than 0 in routine gser()");
			return -1;
		} else {
			ap = df;
			del = sum = 1.0 / df;
			for (n = 1; n <= ITMAX; n++) {
				++ap;
				del *= x / ap;
				sum += del;
				if (Math.abs(del) < Math.abs(sum) * EPS) {
					// gamser = sum * Math.exp(-x + df * Math.log(x) - (gln));
					return Math.log(sum) - x + df * Math.log(x) - (gln);
				}
			}
			System.err
					.println("a too large, ITMAX too small in routine gser()");
			return -1;
		}
	}

	// TODO: completar javadoc
	/**
	 * 
	 * @param a
	 * @param x
	 * @return
	 */
	static double ln_gcf(double a, double x) {
		int i;
		double an, b, c, d, del, h;

		gln = gammln(a);
		b = x + 1.0 - a;
		c = 1.0 / FPMIN;
		d = 1.0 / b;
		h = d;

		for (i = 1; i <= ITMAX; i++) {
			an = -i * (i - a);
			b += 2.0;
			d = an * d + b;
			if (Math.abs(d) < FPMIN)
				d = FPMIN;
			c = b + an / c;
			if (Math.abs(c) < FPMIN)
				c = FPMIN;
			d = 1.0 / d;
			del = d * c;
			h *= del;
			if (Math.abs(del - 1.0) < EPS)
				break;
		}

		if (i > ITMAX)
			System.err.println("a too large, ITMAX too small in gcf()");

		return -x + a * Math.log(x) - (gln) + Math.log(h);
	}

	// This method computes the contingency tables for a
	// triplet. If the parameter context is different to null, then
	// TODO: completar javadoc
	/**
	 * This method computes the contingency tables for a triplet
	 * 
	 * @param x
	 * @param y
	 * @param z
	 * @return
	 */
	private HashMap<Integer, int[][]> computeContingencyTables(int x, int y,
			List<Integer> z, Dataset dataset) {
		int[] counts;
		BitSet vars = new BitSet(dataset.getNumberOfVariables());
		vars.set(x, true);
		vars.set(y, true);

		List<int[]> permutationsOfZ = Permutations.getPermutations(z,
				dataset.getCardinalities());
		HashMap<Integer, int[][]> contTables = new HashMap<Integer, int[][]>();

		for (int[] permutationOfZ : permutationsOfZ) {
			Map<Integer, Integer> map = new HashMap<Integer, Integer>();
			for (int i = 0; i < z.size(); i++) {
				vars.set(z.get(i), true);
				map.put(z.get(i), permutationOfZ[i]);
			}

			counts = dataset.computeSufficientStatistics(vars, map, null);
			int[][] contTable = new int[dataset.getCardinalities(x)][dataset
					.getCardinalities(x)];
			contTable[0][0] = counts[0];
			contTable[0][1] = counts[1];
			contTable[1][0] = counts[2];
			contTable[1][1] = counts[3];

			// active garbage collector to drop counts
			counts = null;
			contTables.put(permutationOfZ.toString().hashCode(), contTable);
		}

		return contTables;
	}

	@Override
	public boolean independent(Dataset dataset, List<Integer> x,
			List<Integer> y, List<Integer> z) {
		// TODO Auto-generated method stub
		return false;
	}

	@Override
	public boolean independent(Dataset dataset, IndependenceTriplet triplet) {
		// TODO Auto-generated method stub
		return false;
	}

	@Override
	public double[] getLoglikelihoods() {
		// TODO Auto-generated method stub
		return null;
	}

	@Override
	public double[] getLogProbs() {
		throw new RuntimeException("Chi square only computes log likelihood. It does not compute posterior probabilities. Use the getLogLikelihoods method");
	}

}
