package dr.evolution.coalescent;

import dr.evolution.tree.Tree;
import dr.evolution.util.Date;
import dr.evolution.util.Taxa;
import dr.evolution.util.Taxon;
import dr.evolution.util.Units;

import java.io.*;
import java.util.ArrayList;
import java.util.List;

/**
 * @author Alexei Drummond
 */
public class PLCoalescentSimulator {
    private static PrintStream debug = System.err;

    public static void main(String[] arg) throws IOException {

        // READ DEMOGRAPHIC FUNCTION

        String filename = arg[0];
        BufferedReader reader = new BufferedReader(new FileReader(filename));

        double popSizeScale = 1.0;
        double generationTime = 1.0;
        if (arg.length > 2) {
            popSizeScale = Double.parseDouble(arg[2]);
        }
        if (arg.length > 3) {
            generationTime = Double.parseDouble(arg[3]);
        }

        PrintWriter populationFuncLogger = null;
        if (arg.length > 5) {
            String logFileName = arg[5];
            if (logFileName.equals("-")) {
                populationFuncLogger = new PrintWriter(System.out);
            } else {
                populationFuncLogger = new PrintWriter(new FileWriter(logFileName));
            }
        }

        List<Double> times = new ArrayList<Double>();

        String line = reader.readLine();
        String[] tokens = line.trim().split("[\t ]+");
        if (tokens.length < 2) throw new RuntimeException();

        List<Double>[] popSizes = new List[tokens.length - 1];
        for (int i = 0; i < tokens.length - 1; i++) {
            popSizes[i] = new ArrayList<Double>();
        }

        while (line != null) {

            double time = Double.parseDouble(tokens[0]) / generationTime;

            times.add(time);
            for (int i = 1; i < tokens.length; i++) {
                popSizes[i - 1].add(Double.parseDouble(tokens[i]));
            }
            line = reader.readLine();
            if (line != null) {
                tokens = line.trim().split("[\t ]+");
                if (tokens.length != popSizes.length + 1) throw new RuntimeException();
            }
        }

        reader.close();

        // READ SAMPLE TIMES

        String samplesFilename = arg[1];

        reader = new BufferedReader(new FileReader(samplesFilename));

        line = reader.readLine();
        Taxa taxa = new Taxa();
        int id = 0;
        while (line != null) {

            if (!line.startsWith("#")) {

                tokens = line.split("[\t ]+");

                if (tokens.length == 4) {

                    double t0 = Double.parseDouble(tokens[0]);
                    double t1 = Double.parseDouble(tokens[1]);
                    double dt = Double.parseDouble(tokens[2]);
                    int k = Integer.parseInt(tokens[3]);
                    for (double time = t0; time <= t1; time += dt) {

                        double sampleTime = time / generationTime;
                        for (int i = 0; i < k; i++) {
                            Taxon taxon = new Taxon("t" + id);
                            taxon.setAttribute(dr.evolution.util.Date.DATE, new Date(sampleTime, Units.Type.GENERATIONS, true));
                            taxa.addTaxon(taxon);
                            id += 1;
                        }
                    }

                } else {

                    // sample times are in the same units as simulation
                    double sampleTime = Double.parseDouble(tokens[0]) / generationTime;
                    int count = Integer.parseInt(tokens[1]);

                    for (int i = 0; i < count; i++) {
                        Taxon taxon = new Taxon(id + "");
                        taxon.setAttribute(dr.evolution.util.Date.DATE, new Date(sampleTime, Units.Type.GENERATIONS, true));
                        taxa.addTaxon(taxon);
                        id += 1;
                    }
                }
            }
            line = reader.readLine();
        }

        double minTheta = Double.MAX_VALUE;
        double maxTheta = 0.0;

        PrintWriter out;
        if (arg.length > 4) {
            out = new PrintWriter(new FileWriter(arg[4]));
        } else {
            out = new PrintWriter(System.out);
        }

        int pp = 0;
        for (List<Double> popSize : popSizes) {
            double[] thetas = new double[popSize.size()];
            double[] intervals = new double[times.size() - 1];

            if (populationFuncLogger != null) {
                populationFuncLogger.println("# " + pp);
                ++pp;
            }

            // must reverse the direction of the model
            for (int j = intervals.length; j > 0; j--) {
                intervals[intervals.length - j] = times.get(j) - times.get(j - 1);

                final double theta = popSize.get(j) * popSizeScale;
                thetas[intervals.length - j] = theta;
                if (theta < minTheta) {
                    minTheta = theta;
                }
                if (theta > maxTheta) {
                    maxTheta = theta;
                }

                final double t = times.get(intervals.length) - times.get(j);
                if (populationFuncLogger != null) {
                    populationFuncLogger.println(t + "\t" + theta);
                }
            }

            if (debug != null) {
                debug.println("min theta = " + minTheta);
                debug.println("max theta = " + maxTheta);
            }

            PiecewiseLinearPopulation demo = new PiecewiseLinearPopulation(intervals, thetas, Units.Type.GENERATIONS);

            CoalescentSimulator simulator = new CoalescentSimulator();
            Tree tree = simulator.simulateTree(taxa, demo);

            out.println(Tree.Utils.newick(tree));
            if (debug != null) {
                debug.println(Tree.Utils.newick(tree));
            }
        }

        if (populationFuncLogger != null) {
            populationFuncLogger.flush();
            populationFuncLogger.close();
        }

        out.flush();
        out.close();
    }
}
