/*
 * License: FreeBSD (Berkeley Software Distribution)
 * Copyright (c) 2013, Sara Sheehan, Kelley Harris, Yun Song
 * 
 * based on:
 * Copyright (c) 2011, Joshua Paul, Matthias SteinrŸcken, Yun Song
 */

package edu.berkeley.utility;

import Jama.EigenvalueDecomposition;
import Jama.Matrix;

// stationary probabilities for finite-states, finite alleles haplotype, for a given array of transition matrices
public class StationaryProbability {
	double[][] stationaryProbabilityTable;
	
	public StationaryProbability(Matrix[] mutationMatrices) {
		this.stationaryProbabilityTable = new double[mutationMatrices.length][];
		
		for (int i = 0; i < mutationMatrices.length; i++) {
			this.stationaryProbabilityTable[i] = getStationaryDistribution(mutationMatrices[i].getArray());
		}
	}
	
	public double getStationaryProbability(int timeIdx, int allele) {
		return this.stationaryProbabilityTable[timeIdx][allele];
	}
	
	public static double[] getStationaryDistribution(double[][] stochasticMatrix) {
		Matrix sMatrix = new Matrix(stochasticMatrix);
		EigenvalueDecomposition eigenDecomp = sMatrix.transpose().eig();
		
		int largestEigenvalueIndex = -1;
		for (int ev = 0; ev < eigenDecomp.getRealEigenvalues().length; ev++) {
			if (largestEigenvalueIndex == -1 || eigenDecomp.getRealEigenvalues()[ev] > eigenDecomp.getRealEigenvalues()[largestEigenvalueIndex]) largestEigenvalueIndex = ev;
		}
		
		Matrix uStationaryDist = eigenDecomp.getV().getMatrix(0, eigenDecomp.getV().getRowDimension() - 1, largestEigenvalueIndex, largestEigenvalueIndex);
		
		double colSum = 0;
		for (int i = 0; i < uStationaryDist.getRowDimension(); i++) colSum += uStationaryDist.get(i, 0);
		
		return uStationaryDist.timesEquals(1 / colSum).getColumnPackedCopy();
	}
	
	public static double[] getLogStationaryDistribution(double[][] stochasticMatrix) {
		double[] stationaryDist = getStationaryDistribution(stochasticMatrix);
		double[] logStationaryDist = new double[stationaryDist.length];
		for (int a=0; a < stationaryDist.length; a++) {
    		logStationaryDist[a] = Math.log(stationaryDist[a]);
    	}
		return logStationaryDist;
	}
	
	// function for determining whether the mutation matrix is symmetric
	public static boolean isSymmetric(double[][] stochasticMatrix) {
		int numAlleles = stochasticMatrix.length;
		assert numAlleles > 0;
		assert stochasticMatrix[0].length == numAlleles;
		for (int i=0; i < numAlleles; i++) {
			for (int j=i+1; j < numAlleles; j++) {
				if (stochasticMatrix[i][j] != stochasticMatrix[j][i]) {
					return false;
				}
			}
		}
		return true;
	}
}