/*
 * This file is part of Jstacs.
 *
 * Jstacs is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * 
 * Jstacs is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with Jstacs. If not, see <http://www.gnu.org/licenses/>.
 * 
 * For more information on Jstacs, visit http://www.jstacs.de
 */

package de.jstacs.models.utils;

import java.io.IOException;

import de.jstacs.WrongAlphabetException;
import de.jstacs.data.AlphabetContainer;
import de.jstacs.data.Sample;
import de.jstacs.data.Sequence;
import de.jstacs.data.Sample.ElementEnumerator;
import de.jstacs.data.sequences.IntSequence;
import de.jstacs.data.sequences.WrongSequenceTypeException;
import de.jstacs.models.Model;

/**
 * This class is useful for some test for any (discrete) models. It implements
 * several statistics (log-likelihood, Shannon entropy, AIC, BIC, ...) to
 * compare models.
 * 
 * @see de.jstacs.models.AbstractModel
 * 
 * @author Jens Keilwagen
 */
public class ModelTester {
	/**
	 * Returns the Kullback-Leibler-divergence <code>D(p_m1||p_m2)</code>.
	 * 
	 * <br>
	 * <br>
	 * 
	 * Computes <code>\sum_x p(x|m1) * \log \frac{p(x|m1)}{p(x|m2)}</code>.
	 * 
	 * @param m1
	 *            one discrete model
	 * @param m2
	 *            another discrete model
	 * @param length
	 *            the length of the sequence (for inhomogeneous models length
	 *            has to be {@link Model#getLength()})
	 * 
	 * @return the Kullback-Leibler-divergence
	 * 
	 * @throws Exception
	 *             if something went wrong
	 */
	public static double getKLDivergence(Model m1, Model m2, int length)
			throws Exception {
		SeqIterator s = new SeqIterator(m1.getAlphabetContainer(), length);
		Sequence seq;
		double kl = 0, v;
		do {
			seq = s.getSequence();
			v = m1.getLogProbFor(seq);
			kl += Math.exp(v) * (v - m2.getLogProbFor(seq));
		} while (s.next());
		return kl;
	}

	/**
	 * Returns the difference of the Kullback-Leibler-divergences, i.e.
	 * <code>D(p_m1||p_m2) - D(p_m2||p_m1)</code>.
	 * 
	 * <br>
	 * <br>
	 * 
	 * Computes
	 * <code>\sum_x (p(x|m1)-p(x|m2)) * \log \frac{p(x|m1)}{p(x|m2)}</code>.
	 * 
	 * @param m1
	 *            one discrete model
	 * @param m2
	 *            another discrete model
	 * @param length
	 *            the length of the sequence (for inhomogeneous models length
	 *            has to be {@link Model#getLength()})
	 * 
	 * @return the difference of the Kullback-Leibler-divergence
	 * 
	 * @throws Exception
	 *             if something went wrong
	 */
	public static double getSymKLDivergence(Model m1, Model m2, int length)
			throws Exception {
		SeqIterator s = new SeqIterator(m1.getAlphabetContainer(), length);
		Sequence seq;
		double kl = 0, p1, p2;
		do {
			seq = s.getSequence();
			p1 = m1.getProbFor(seq);
			p2 = m2.getProbFor(seq);
			kl += (p1 - p2) * Math.log(p1 / p2);
		} while (s.next());
		return kl;
	}

	/**
	 * Returns the log-likelihood of a {@link Sample} <code>data</code> for a
	 * given model <code>m</code>.
	 * 
	 * @param m
	 *            the given model
	 * @param data
	 *            the {@link Sample}
	 * 
	 * @return the log-likelihood of <code>data</code>
	 * 
	 * @throws Exception
	 *             if something went wrong
	 */
	public static double getLogLikelihood(Model m, Sample data)
			throws Exception {
		return getLogLikelihood(m, data, null);
	}

	/**
	 * Returns the log-likelihood of a {@link Sample} <code>data</code> for a
	 * given model <code>m</code>.
	 * 
	 * @param m
	 *            the given model
	 * @param data
	 *            the {@link Sample}
	 * @param weights
	 *            the weight for each element of the {@link Sample}
	 * 
	 * @return the log-likelihood of <code>data</code>
	 * 
	 * @throws Exception
	 *             if something went wrong
	 */
	public static double getLogLikelihood(Model m, Sample data, double[] weights)
			throws Exception {
		int counter, d = data.getNumberOfElements();
		double erg = 0;
		ElementEnumerator ei = new ElementEnumerator(data);
		if (weights == null) {
			for (counter = 0; counter < d; counter++) {
				erg += m.getLogProbFor(ei.nextElement());
			}
		} else if (d != weights.length) {
			throw new IllegalArgumentException(
					"The weights and the sample does not match.");
		} else {
			for (counter = 0; counter < d; counter++) {
				erg += weights[counter] * m.getLogProbFor(ei.nextElement());
			}
		}
		return erg;
	}

	/**
	 * This method computes the marginal distribution for any discrete model
	 * <code>m</code> and all sequences that fulfill the <code>constraint</code>
	 * , if possible.
	 * 
	 * @param m
	 *            a discrete model
	 * @param constraint
	 *            <code>constraint[i] < 0</code> stands for an irrelevant
	 *            position, <code>constraint[i] = c</code> with
	 * 
	 *            <code>0 <= c < m.getAlphabets()[(m.getLength==0)?0:i].getAlphabetLength()</code>
	 *            is the encoded character of position <code>i</code>
	 * 
	 * @return the marginal distribution for a discrete model
	 * 
	 * @throws Exception
	 *             if something went wrong
	 */
	public static double getMarginalDistribution(Model m, int[] constraint)
			throws Exception {
		if (m.getLength() != 0 && m.getLength() != constraint.length) {
			throw new IOException(
					"This model can only classify sequences of length "
							+ m.getLength() + ".");
		}
		double erg = 0;
		SeqIterator s = new SeqIterator(m.getAlphabetContainer(),
				constraint.length);
		do {
			if (s.isSatisfied(constraint)) {
				erg += m.getProbFor(s.getSequence());
			}
		} while (s.next());
		return erg;
	}

	/**
	 * This method computes the maximum deviation between the probabilities for
	 * all sequences of <code>length</code> for discrete models <code>m1</code>
	 * and <code>m2</code>.
	 * 
	 * @param m1
	 *            one discrete model
	 * @param m2
	 *            another discrete model
	 * @param length
	 *            the length of the sequence (for inhomogeneous models length
	 *            has to be {@link Model#getLength()})
	 * 
	 * @return the maximum deviation between the probabilities
	 * 
	 * @throws Exception
	 *             if something went wrong
	 */
	public static double getMaxOfDeviation(Model m1, Model m2, int length)
			throws Exception {
		if (m1.getLength() != 0 && m1.getLength() != length) {
			throw new IOException(
					"The model m1 can only classify sequences of length "
							+ m1.getLength() + ".");
		}
		if (m2.getLength() != 0 && m2.getLength() != length) {
			throw new IOException(
					"This model m2 can only classify sequences of length "
							+ m2.getLength() + ".");
		}
		if (!m1.getAlphabetContainer().checkConsistency(
				m2.getAlphabetContainer())) {
			throw new IOException(
					"The models are training on different alphabets.");
		}
		double p, max = 0;
		SeqIterator s = new SeqIterator(m1.getAlphabetContainer(), length);
		Sequence seq;
		do {
			seq = s.getSequence();
			p = Math.abs(m1.getProbFor(seq) - m2.getProbFor(seq, 0, s.last));
			if (p > max) {
				max = p;
			}
		} while (s.next());
		return max;
	}

	/**
	 * Returns one most probable sequence for the discrete model <code>m</code>.
	 * (Maybe there are more than one most probable sequences. In this case only
	 * one of them is returned.)
	 * 
	 * <br>
	 * <br>
	 * 
	 * This is only a standard implementation. For some special models like
	 * Markov models it is possible to compute the probabilities of the
	 * sequences much faster by using a dynamic-programming-algorithm.
	 * 
	 * @param m
	 *            the discrete model
	 * @param length
	 *            the length of the sequence (for inhomogeneous models length
	 *            has to be {@link Model#getLength()})
	 * 
	 * @return one most probable sequence
	 * 
	 * @throws Exception
	 *             if something went wrong
	 */
	public static Sequence getMostProbableSequence(Model m, int length)
			throws Exception {
		SeqIterator s = new SeqIterator(m.getAlphabetContainer(), length);
		Sequence current, seq = s.getSequence();
		double pmax = m.getLogProbFor(seq), p;
		while (s.next()) {
			current = s.getSequence();
			p = m.getLogProbFor(current);
			if (p > pmax) {
				pmax = p;
				seq = current;
			}
		}
		return seq;
	}

	/**
	 * This method computes the Shannon entropy for any discrete model
	 * <code>m</code> and all sequences of <code>length</code>, if possible.
	 * 
	 * @param m
	 *            the discrete model
	 * @param length
	 *            the length of the sequence (for inhomogeneous models length
	 *            has to be {@link Model#getLength()})
	 * 
	 * @return the Shannon entropy for a discrete model
	 * 
	 * @throws Exception
	 *             if something went wrong
	 */
	public static double getShannonEntropy(Model m, int length)
			throws Exception {
		if (m.getLength() != 0 && m.getLength() != length) {
			throw new IOException(
					"This model can only classify sequences of length "
							+ m.getLength() + ".");
		}
		double p, erg = 0;
		SeqIterator s = new SeqIterator(m.getAlphabetContainer(), length);
		do {
			p = m.getProbFor(s.getSequence());
			if (p != 0) {
				erg -= p * Math.log(p);
			}
			if (p < 0 || p > 1) {
				throw new IOException("The probability of sequence "
						+ s.getSequence() + " is not correct (" + p + ").");
			}
		} while (s.next());
		return erg;
	}

	/**
	 * This method computes the Shannon entropy in bits for any discrete model
	 * <code>m</code> and all sequences of <code>length</code>, if possible.
	 * 
	 * @param m
	 *            the discrete model
	 * @param length
	 *            the length of the sequence (for inhomogeneous models length
	 *            has to be {@link Model#getLength()})
	 * 
	 * @return the Shannon entropy in bits for a discrete model
	 * 
	 * @throws Exception
	 *             if something went wrong
	 */
	public static double getShannonEntropyInBits(Model m, int length)
			throws Exception {
		return (getShannonEntropy(m, length) / Math.log(2));
	}

	/**
	 * This method computes the sum of deviations between the probabilities for
	 * all sequences of <code>length</code> for discrete models <code>m1</code>
	 * and <code>m2</code>.
	 * 
	 * @param m1
	 *            one discrete model
	 * @param m2
	 *            another discrete model
	 * @param length
	 *            the length of the sequence (for inhomogeneous models length
	 *            has to be {@link Model#getLength()})
	 * 
	 * @return the sum of deviations between the probabilities
	 * 
	 * @throws Exception
	 *             if something went wrong
	 */
	public static double getSumOfDeviation(Model m1, Model m2, int length)
			throws Exception {
		if (m1.getLength() != 0 && m1.getLength() != length) {
			throw new IOException(
					"The model m1 can only classify sequences of length "
							+ m1.getLength() + ".");
		}
		if (m2.getLength() != 0 && m2.getLength() != length) {
			throw new IOException(
					"This model m2 can only classify sequences of length "
							+ m2.getLength() + ".");
		}
		if (!m1.getAlphabetContainer().checkConsistency(
				m2.getAlphabetContainer())) {
			throw new IOException(
					"The models are training on different alphabets.");
		}
		double sum = 0;
		SeqIterator s = new SeqIterator(m1.getAlphabetContainer(), length);
		Sequence seq;
		do {
			seq = s.getSequence();
			sum += Math.abs(m1.getProbFor(seq, 0, s.last)
					- m2.getProbFor(seq, 0, s.last));
		} while (s.next());
		return sum;
	}

	/**
	 * This method computes the marginal distribution for any discrete model
	 * <code>m</code> and all sequences of <code>length</code>, if possible. So
	 * this method can be used to give a hint whether a model is a distribution
	 * or if some mistakes are in the implementation.
	 * 
	 * <br>
	 * <br>
	 * 
	 * It is expected that this method delivers the value 1.0, but because of
	 * the limited precision in Java the value 1.0 is unrealistic.
	 * 
	 * <br>
	 * <br>
	 * 
	 * <code>Math.abs( 1.0d - getSumOfDistribution( m, length )</code> should be
	 * smaller than <code>1E-10</code>.
	 * 
	 * @param m
	 *            the discrete model
	 * @param length
	 *            the length of the sequence (for inhomogeneous models length
	 *            has to be {@link Model#getLength()})
	 * 
	 * @return the marginal distribution for a discrete model
	 * 
	 * @throws Exception
	 *             if something went wrong
	 */
	public static double getSumOfDistribution(Model m, int length)
			throws Exception {
		if (m.getLength() != 0 && m.getLength() != length) {
			throw new IOException(
					"This model can only classify sequences of length "
							+ m.getLength() + ".");
		}
		double erg = 0, p;
		SeqIterator s = new SeqIterator(m.getAlphabetContainer(), length);
		do {
			p = m.getProbFor(s.getSequence());
			if (p < 0 || p > 1) {
				throw new IOException("The probability (" + p
						+ ") for sequence \"" + s.getSequence()
						+ "\" is not in [0,1].");
			}
			erg += p;
		} while (s.next());
		return erg;
	}

	/**
	 * This method computes the value of Akaikes Information Criterion (AIC). It
	 * uses the formula: AIC = <code>2 * log L(t,x) - 2*k</code>, where
	 * <code>L(t,x)</code> is the likelihood of the {@link Sample} and
	 * <code>k</code> is the number of parameters in the model.
	 * 
	 * <br>
	 * <br>
	 * 
	 * The value of the AIC can be used for model selection.
	 * 
	 * @param m
	 *            a trained model
	 * @param s
	 *            the {@link Sample} for the test
	 * @param k
	 *            the number of parameters of the model <code>m</code>
	 * 
	 * @return the value of AIC
	 * 
	 * @throws Exception
	 *             if something went wrong
	 */
	public static double getValueOfAIC(Model m, Sample s, int k)
			throws Exception {
		return 2 * getLogLikelihood(m, s) - 2 * k;
	}

	/**
	 * This method computes the value of the Bayesian Information Criterion
	 * (BIC). It uses the formula: BIC = <code>2 * log L(t,x) - k *
	 * log n</code>, where <code>L(t,x)</code> is the likelihood of the
	 * {@link Sample}, <code>k</code> is the number of parameters in the model
	 * and <code>n</code> is the number of sequences in the {@link Sample}.
	 * 
	 * <br>
	 * <br>
	 * 
	 * The value of the BIC can be used for model selection.
	 * 
	 * @param m
	 *            a trained model
	 * @param s
	 *            the {@link Sample} for the test
	 * @param k
	 *            the number of parameters of the model <code>m</code>
	 * 
	 * @return value of AIC
	 * 
	 * @throws Exception
	 *             if something went wrong
	 */
	public static double getValueOfBIC(Model m, Sample s, int k)
			throws Exception {
		return 2 * getLogLikelihood(m, s) - k
				* StrictMath.log(s.getNumberOfElements());
	}

	private static class SeqIterator {
		// the sequence
		private int[] seq;

		// simple or not
		private boolean simple;

		// the alphabet.length
		private int[] a;

		// length of the seq
		private int l, last;

		// the alphabets
		private AlphabetContainer abc;

		public SeqIterator(AlphabetContainer abc, int length)
				throws IllegalArgumentException {
			if (!abc.isDiscrete()) {
				throw new IllegalArgumentException("The model is not discrete.");
			}
			this.abc = abc;
			simple = abc.isSimple();
			a = new int[(simple ? 1 : length) + 1];
			int i = 0;
			for (; i < a.length - 1; i++) {
				a[i] = (int) abc.getAlphabetLengthAt(i) - 1;
			}
			a[i] = 1;
			l = length;
			last = l - 1;
			seq = new int[length + 1];
		}

		public boolean next() {
			int s_index = 0;
			while (seq[s_index] == a[simple ? 0 : s_index]) {
				seq[s_index++] = 0;
			}
			seq[s_index]++;
			return seq[l] == 0;
		}

		private boolean isSatisfied(int[] constr) {
			int i = 0;
			while (i < constr.length
					&& (constr[i] == -1 || constr[i] == seq[i])) {
				i++;
			}
			return (i == constr.length);
		}

		public Sequence getSequence() throws WrongAlphabetException,
				WrongSequenceTypeException {
			return new IntSequence(abc, seq, 0, l);
		}
	}
}
