/*
 * PiecewiseLinearPopulation.java
 *
 * Copyright (C) 2002-2008 Alexei Drummond and Andrew Rambaut
 *
 * This file is part of BEAST.
 * See the NOTICE file distributed with this work for additional
 * information regarding copyright ownership and licensing.
 *
 * BEAST is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as
 * published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 *  BEAST 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 Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with BEAST; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA  02110-1301  USA
 */

package dr.evolution.coalescent;

/**
 * @author Alexei Drummond
 * @author Andrew Rambaut
 * @version $Id: PiecewiseLinearPopulation.java,v 1.7 2005/05/24 20:25:56 rambaut Exp $
 */
public class PiecewiseLinearPopulation extends PiecewiseConstantPopulation {

    /**
     * Construct demographic model with default settings
     */
    public PiecewiseLinearPopulation(double[] intervals, double[] thetas, Type units) {

        super(intervals, thetas, units);
    }

    // **************************************************************
    // Implementation of abstract methods
    // **************************************************************

    /**
     * @return the value of the demographic function for the given epoch and time relative to start of epoch.
     */
    protected final double getDemographic(int epoch, double t) {
        // if in last epoch then the population is flat.
        if (epoch == (thetas.length - 1)) {
            return getEpochDemographic(epoch);
        }

        final double popSize1 = getEpochDemographic(epoch);
        final double popSize2 = getEpochDemographic(epoch + 1);

        final double width = getEpochDuration(epoch);

        assert 0 <= t && t <= width;

        return popSize1 + (t / width) * (popSize2 - popSize1);
    }

    public DemographicFunction getCopy() {
        PiecewiseLinearPopulation df = new PiecewiseLinearPopulation(new double[intervals.length], new double[thetas.length], getUnits());
        System.arraycopy(intervals, 0, df.intervals, 0, intervals.length);
        System.arraycopy(thetas, 0, df.thetas, 0, thetas.length);

        return df;
    }

    /**
     * @return the value of the intensity function for the given epoch.
     */
    protected final double getIntensity(int epoch) {

        final double N1 = getEpochDemographic(epoch);
        final double N2 = getEpochDemographic(epoch + 1);
        final double w = getEpochDuration(epoch);

        return w * Math.log(N2 / N1) / (N2 - N1);
    }

    /**
     * @return the value of the intensity function for the given epoch and time relative to start of epoch.
     */
    protected final double getIntensity(int epoch, double relativeTime) {
        assert relativeTime <= getEpochDuration(epoch);

        final double N1 = getEpochDemographic(epoch);
        final double N2 = getEpochDemographic(epoch + 1);
        final double w = getEpochDuration(epoch);

        return w * Math.log(N1 * w / (N2 * relativeTime + N1 * (w - relativeTime))) / (N1 - N2);
    }

    /**
     * @param targetI the intensity
     * @return the time corresponding to the given target intensity
     */
    public final double getInverseIntensity(double targetI) {

        int epoch = 0;
        double cI = 0;
        double eI = getIntensity(epoch);
        double time = 0.0;
        while (cI + eI < targetI) {
            cI += eI;
            time += getEpochDuration(epoch);

            epoch += 1;
            eI = getIntensity(epoch);
        }

        time += getInverseIntensity(epoch, targetI - cI);

        return time;
    }

    /**
     * @param epoch the epoch
     * @param I     intensity corresponding to relative time within this epoch
     * @return the relative time within the given epoch to accumulate the given residual intensity
     */
    private double getInverseIntensity(int epoch, double I) {

        final double N1 = getEpochDemographic(epoch);
        final double N2 = getEpochDemographic(epoch + 1);
        final double w = getEpochDuration(epoch);

        double time = -N1 * w * (Math.exp(I * (N2 - N1) / w) - 1.0) / (N1 - N2);

        return time;
    }
}

@@ -2922,32 +2922,61 @@
 ration(epoch);%0A%0A
+        if( N1 != N2 ) %7B%0A    
         return w
@@ -3005,24 +3005,76 @@
  (N2 - N1);%0A
+        %7D else %7B%0A            return w/N1;%0A        %7D%0A
     %7D%0A%0A    /
@@ -3478,32 +3478,59 @@
 ration(epoch);%0A%0A
+        if( N1 != N2 ) %7B%0A  
         return w
@@ -3610,16 +3610,79 @@
  - N2);%0A
+        %7D else %7B%0A            return relativeTime/N1;%0A        %7D%0A
     %7D%0A%0A 
@@ -4703,23 +4703,100 @@
     
-double time = -
+final double dn = N2 - N1;%0A        double time;%0A        if( dn != 0.0 ) %7B%0A           time = 
 N1 *
@@ -4814,25 +4814,18 @@
 exp(I * 
-(N2 - N1)
+dn
  / w) - 
@@ -4831,26 +4831,73 @@
  1.0) / 
-(N1 - N2);
+dn;%0A        %7D else %7B%0A            time = I * N1;%0A        %7D
 %0A%0A      

Add handling of constant intervals in Piecewise linear. Those can happen demographic is explicitly set for simulation purposes.