// SkylinePlot.java
//
// (c) 1999-2000 PAL Development Core Team
//
// This package may be distributed under the
// terms of the GNU General Public License (GPL)

 
package pal.coalescent;

import pal.tree.*;
import pal.misc.*;
import pal.math.*;
import pal.statistics.*;
import pal.io.*;

import java.io.*;

/**
 * Generalized skyline plot derived from a tree or a coalescent interval.
 * See K. Strimmer and O.G. Pybus. 2001 for more details.
 *
 * This method is a successor of the "classic" skyline plot (Pybus et al. 2000)
 * which is implemented in pal.coalescent.ClassicSkylinePlot
 * 
 * @version $Id: SkylinePlot.java,v 1.13 2001/04/03 11:50:29 korbinian Exp $
 *
 * @author Korbinian Strimmer
 */
public class SkylinePlot implements Report, Units, Serializable
{
	//
	// Public stuff
	//

	/**
	 * Construct skyline plot from a tree.  If a corresponding
	 * coalescent interval <= minSize it is pooled with the neighbor interval
	 * (minSize > 0 is a smoothing parameter)
	 */
	public SkylinePlot(Tree tree, double minSize)
	{
		CoalescentIntervals ci = IntervalsExtractor.extractFromClockTree(tree, minSize);
		
		numBootstraps = 0;
		
		constructSkylinePlot(ci, null);
	}

	/**
	 * Construct skyline plot plus bias, error, and confidence intervals by supplying
	 * a binary reference tree and a set of binary bootstrap trees.  If an interval
	 * is <= minSize in the reference tree or in average in the bootstrap trees
	 * the corresponding interval is pooled with the neighbor interval
	 * in all of the trees (to keep coalescent intervals consistent)
	 */
	public SkylinePlot(Tree tree, Tree[] btree, double minSize)
	{
		CoalescentIntervals ci = IntervalsExtractor.extractFromClockTree(tree);
		numBootstraps = btree.length;
		
		if(!ci.isBinaryCoalescent())
		{
			throw new IllegalArgumentException("Reference tree is not binary");
		}
			
		CoalescentIntervals[] bci = new CoalescentIntervals[numBootstraps];
		for (int i = 0; i < numBootstraps; i++)
		{
			bci[i] = IntervalsExtractor.extractFromClockTree(btree[i]);
			
			if(!bci[i].isBinaryCoalescent())
			{
				throw new IllegalArgumentException("Bootstrao tree " + i + " is not binary");
			}	
		}
				
		int numIntervals = ci.getIntervalCount();
		boolean[] smallInterval = new boolean[numIntervals];
		boolean[] testInterval = new boolean[numIntervals];
		
		
		// find intervals that are small in the reference tree
		ci.getSmallIntervals(minSize, testInterval);
		for (int k = 0; k < numIntervals; k++)
		{
			smallInterval[k] = testInterval[k];
		}	
		
		// find intervals that are small in all of the bootstrap trees
		//int[] stat = new int[numIntervals];		
		/*for (int i = 0; i < numBootstraps; i++)
		{
			bci[i].getSmallIntervals(minSize, testInterval);
			for (int k = 0; k < numIntervals; k++)
			{
				if (testInterval[k] == true)
				{
					stat[k]++;
				}
			}
		}
		for (int k = 0; k < numIntervals; k++)
		{
			if (stat[k] == numBootstraps) smallInterval[k] = true;
		}*/

		// find intervals that are small on average over all bootstrap trees		
		for (int k = 0; k < numIntervals; k++)
		{
			double average = 0.0;
			for (int i = 0; i < numBootstraps; i++)
			{
				average += bci[i].getInterval(k);
			}
			average = average / (double) numBootstraps;
			
			if (average <= minSize)
			{
				//System.out.println("small average " + k);
				smallInterval[k] = true;
			}
		}
							
		ci.poolIntervals(smallInterval);
		for (int i = 0; i < numBootstraps; i++)
		{
			bci[i].poolIntervals(smallInterval);
		}
						
		constructSkylinePlot(ci, bci);
	}

	
	/**
	 * Construct skyline plot from a set of given coalescent intervals
	 */
	public SkylinePlot(CoalescentIntervals ci)
	{
		numBootstraps = 0;
		constructSkylinePlot(ci, null);
	}


	/**
	 * Construct skyline plot from a set of given coalescent intervals
	 * and additionally infer errors, bias, and confidence sets by
	 * supplying bootstrap coalescent intervals.  Note that the bootstrap
	 * intervals must be compatible with the reference confidence interval
	 */
	public SkylinePlot(CoalescentIntervals ci, CoalescentIntervals[] bci)
	{
		int numLineages = ci.getNumLineages(0);
		int numIntervals = ci.getIntervalCount();
		numBootstraps = bci.length;
		for (int i = 0; i < numBootstraps; i++)
		{
			if (bci[i].getNumLineages(0) != numLineages ||
				bci[i].getIntervalCount() != numIntervals)
			{
				throw new IllegalArgumentException("Coalescent intervals not compatible");
			}
		}
		 
		constructSkylinePlot(ci, bci);
	}

	
	public String toString()
	{		
		OutputTarget out = OutputTarget.openString();
		report(out);
		out.close();
		
		return out.getString();
	}
	
	public void report(PrintWriter out)
	{
		out.println("Generalized Skyline Plot (Strimmer and Pybus 2001)");
		out.println();
		out.print("Unit of time: ");
		if (units == GENERATIONS)
		{
			out.print("generations");
		}
		else
		{
			out.print("expected substitutions");
		}
		out.println();
		out.print("Lin.\tTime\tLength\tN(t)\tVAR-C");
		
		if (numBootstraps > 0) out.print("\tVAR-P\tBIAS\tERR");
				
		out.println();
		double total = 0.0;
		for (int i = 0; i < size; i++)
		{
			printLine(out, i);
		}
		out.println();
		
		//out.print("log L=");
		//fo.displayDecimal(out, logL, 5);
		
		out.println();
		out.println("Lin.:    number of lineages at the beginning of the interval");
		out.println("N:       generalized skyline plot estimate of effective population size");
		out.println("VAR-C:   estimate of coalescent variance");
	
		if (numBootstraps > 0)
		{
			out.println("VAR-P:   bootstrap estimate of phylogenetic sample variance");
			out.println("BIAS:    bootstrap estimate of phylogenetic sample bias");
			out.println("ERR:     total squared error");
			out.println("Number of bootstraps: " + numBootstraps);
		}	
	}

	private void printLine(PrintWriter out, int interval)
	{
		double begin = getStartOfInterval(interval);
		double end = getStartOfInterval(interval+1);
		double size = end-begin;
		
		out.print(getNumlineages(interval) + "\t");

		fo.displayDecimal(out, begin, 5);
		out.print("\t");
		fo.displayDecimal(out, size, 5);
		out.print("\t");
		fo.displayDecimal(out, populationSize[interval], 4);
		out.print("\t");
		fo.displayDecimal(out, coalVar[interval], 4);

		if (numBootstraps > 0)
		{
			double tot = coalVar[interval];
		
			out.print("\t");
			fo.displayDecimal(out, phyloVar[interval], 4);
			tot +=  phyloVar[interval];
		
			out.print("\t");
			fo.displayDecimal(out, phyloBias[interval], 4);
			tot +=  phyloBias[interval]*phyloBias[interval];
		
			out.print("\t");
			fo.displayDecimal(out, tot , 4);
		}

		//out.print("\t");
		//fo.displayDecimal(out, logProb[interval], 4);

		
		out.println();
	}


	/**
	 * Returns population size at a given time
	 */
	public double getPopulationSizeAtTime(double time)
	{
		return populationSize[findInterval(time)];
	}

	/**
	 * Returns population size in interval i
	 */
	public double getPopulationSize(int i)
	{
		return populationSize[i];
	}


	/**
	 * Returns phylogenetic sample variance in interval i
	 */
	public double getPhylogeneticVar(int i)
	{
		return phyloVar[i];
	}

	/**
	 * Returns phylogenetic sample bias in interval i
	 */
	public double getPhylogeneticBias(int i)
	{
		return phyloBias[i];
	}


	/**
	 * Returns coalescent variance in interval i
	 */
	public double getCoalescentVar(int i)
	{
		return coalVar[i];
	}


	/**
	 * Returns log-probability of interval length
	 */
	public double getLogProbability(int i)
	{
		return logProb[i];
	}


	/**
	 * Returns log-likelihood
	 */
	public double getLogLikelihood()
	{
		return logL;
	}


	/**
	 * Find interval corresponding to a specific time
	 */
	public int findInterval(double time)
	{
		if (time < 0) throw new IllegalArgumentException("Negative values for time are not allowed");
		
		for (int i = 0; i < size-1; i++)
		{
			if (time >= cis[i] && time < cis[i+1]) return i;
		}
		
		return size-1;
	}
	 
	/**
	 * Returns the largest value of time defined in this plot (smallest is zero)
	 */
	public double getMaxTime()
	{
		return maxTime;
	}



	/**
	 * Returns the largest estimate of population size.
	 */
	public double getMaxPopulationSize()
	{
		double max = 0.0;
		for (int i = 0; i < size; i++) {
			if (populationSize[i] > max) {
				max = populationSize[i];
			}
		}

		return max;
	}

	/**
	 * Returns the number of intervals in this skyline plot.
	 */
	public int getSize()
	{
		return size;
	}

	/**
	 * Returns unit of time.
	 */
	public int getUnits()
	{
		return units;
	}

	/**
	 * Returns start time of an interval
	 */
	public double getStartOfInterval(int i)
	{
		if (i < size) return cis[i];
		else return maxTime;
	}

	/**
	 * Returns number of bootstraps
	 */
	public double getBootstrapCount()
	{
		return numBootstraps;
	}

	/**
	 * Returns number of lineages in an interval
	 */
	public int getNumlineages(int i)
	{
		return ci.getNumLineages(i);
	}

	/** 
	 * Returns the coalescent intervals for this skyline plot.
	 */
	public CoalescentIntervals getIntervals() {
		return ci;
	}


	// private
	
	private FormattedOutput fo;
	private CoalescentIntervals ci;
	private int size;
	private Binomial binom;
	private double maxTime;
	private int units;
	
	/** cummulative interval sizes */
	private double[] cis;
	
	/** estimated population size in a coalescent interval */
	private double[] populationSize;


	/** phylogenetic sample variance */
	private double[] phyloVar;

	/** phylogenetic sample bias */
	private double[] phyloBias;


	/** coalescent variance */
	private double[] coalVar;

	/** log-probability of pop. size */
	private double[] logProb;
	
	/** log likelihood */
	private double logL;


	private int numBootstraps;

	private double cut(double n)
	{
		if (n < 0.0)
			return 0.0;
		else
			return n;
	}

	
	/**
	 * Construct skyline plot from given coalescent intervals
	 */
	private void constructSkylinePlot(CoalescentIntervals ci, CoalescentIntervals[] bci)
	{		
		this.ci = ci;
		
		if (bci == null)
		{
			numBootstraps = 0;
		}
		else
		{
			numBootstraps = bci.length;
		}
		
		fo = FormattedOutput.getInstance();
		units = ci.getUnits();
		binom = new Binomial();
		binom.setMax(ci.getNumLineages(0));
		
		/* interval count */
		size = ci.getIntervalCount();
		
		/* cummulative interval sizes */
		cis = new double[size];
		maxTime = 0.0;
		for (int i = 0; i < size; i++)
		{
			cis[i] = maxTime;
			maxTime += ci.getInterval(i);
		}
		
		populationSize = new double[size];
		for (int i = 0; i < size; i++)
		{					
			populationSize[i] =  estimatePopSize(ci, i);
		}

		coalVar = new double[size];
		for (int i = 0; i < size; i++)
		{					
			coalVar[i] =
				populationSize[i]*
				populationSize[i]*
				computeFactor(ci.getNumLineages(i),
					ci.getCoalescentEvents(i));
		}
		
		logProb = new double[size];
		logL = 0.0;
		for (int i = 0; i < size; i++)
		{
			logProb[i] = Math.log(
				computePr(
					ci.getInterval(i),
					ci.getNumLineages(i),
					ci.getCoalescentEvents(i),
					populationSize[i]));
			logL += logProb[i]; 
		}

		
		if (numBootstraps > 0)
		{
			double[][] resampledPopSize = new double[size][numBootstraps];
			for (int i = 0; i < size; i++)
			{
				for (int j = 0; j < numBootstraps; j++)
				{
					resampledPopSize[i][j] = estimatePopSize(bci[j], i);
				}
			}

			phyloVar = new double[size];
			for (int i = 0; i < size; i++)
			{
				phyloVar[i] = BootstrapStatistics.computeVAR(resampledPopSize[i]);
			}

			phyloBias = new double[size];
			for (int i = 0; i < size; i++)
			{
				phyloBias[i] = BootstrapStatistics.computeBIAS(populationSize[i],resampledPopSize[i]);
			}

		
		}		
	}


	/** generalized skyline plot estimate */
	private double estimatePopSize(CoalescentIntervals ci, int interval)
	{
		int lineages = ci.getNumLineages(interval);
		int diffLineages = ci.getCoalescentEvents(interval);	
		
		return ci.getInterval(interval)/computeSum(lineages, diffLineages, -1.0);
	}
		
	
	/** compute \sum_{i=1}^k {l-i+1 \choose 2}^p */
	private double computeSum(int l, int k, double p)
	{
		double sum = 0;
		for (int i = 1; i <= k; i++)
		{
			sum += Math.pow(binom.getNChoose2(l-i+1), p);
		}
		
		return sum;
	}
	
	/** compute sum(l,k,-2)/sum(l,k,-1)^2 */
	private double computeFactor(int l, int k)
	{
		if (k == 1)
			return 1.0;
		else
			return computeSum(l, k,-2.0)/Math.pow(computeSum(l, k,-1.0) ,2.0);
	}
	
	/** compute c_{i,k} */
	private double computeC(int i, int l, int k)
	{
		if (k == 1)
		{
			//System.out.println(1.0);
		
			return 1.0;
		}
		else
		{
			double mult = 1.0;
			for (int j = 1; j <= k; j++)
			{
				if (i != j)
				{
					//double lambdai = 1.0/binom.getNChoose2(l-i+1);
					//double lambdaj = 1.0/binom.getNChoose2(l-j+1);
					
					//mult *= lambdaj/(lambdaj-lambdai);

					double bi = binom.getNChoose2(l-i+1);
					double bj = binom.getNChoose2(l-j+1);
					
					mult *= bi/(bi-bj);
				}
			}
			
			//System.out.println(mult);
			
			return mult;
		}	
	}
	
	/** compute Pr(w_{l,k}) */
	private double computePr(double w, int l, int k, double m)
	{
		double sum = 0.0;
		for (int i = 1; i <= k; i++)
		{
			double lambdai = m/binom.getNChoose2(l-i+1);
			
			sum += computeC(i,l,k)*lambdai*Math.exp(-lambdai*w);
		}
		
		if (sum <= 0) // should not happen!
		{
			System.out.println("w=" + w);
			System.out.println("l=" + l);
			System.out.println("k=" + k);
			System.out.println("m=" + m);
			System.out.println("pr=" +sum);
		}
		
		return sum;
	}
}
