package sb2tests;

import static org.junit.Assert.assertEquals;

import java.util.ArrayList;
import java.util.List;

import org.junit.Test;

import beast.evolution.alignment.TaxonSet;
import beast.util.TreeParser;
import starbeast2.GeneTree;
import starbeast2.MultispeciesCoalescent;
import starbeast2.MultispeciesPopulationModel;
import starbeast2.SpeciesTree;

abstract class PopulationTestHelper {
    String newickSpeciesTree;
    List<String> newickGeneTrees = new ArrayList<>();

    TreeParser speciesTree;
    List<TreeParser> geneTrees = new ArrayList<>();
    
    SpeciesTree speciesTreeWrapper;
    List<GeneTree> geneTreeWrappers = new ArrayList<>();

    MultispeciesCoalescent msc;

    double ploidy;
    double popSize;
    double expectedLogP;
    int nSpecies;

    final double allowedError = 10e-6;

    abstract public TaxonSet generateSuperset() throws Exception;
    abstract public MultispeciesPopulationModel generatePopulationModel() throws Exception;

    @Test
    public void testLogP() throws Exception {
        TaxonSet speciesSuperset = generateSuperset();
        initializeSpeciesTree(speciesSuperset);
        initializeGeneTrees();

        final int nBranches = (nSpecies * 2) - 1;
        final MultispeciesPopulationModel populationModel = generatePopulationModel();
        populationModel.initPopSizes(nBranches);
        populationModel.initPopSizes(popSize);

        msc = new MultispeciesCoalescent();
        msc.initByName("speciesTree", speciesTreeWrapper, "geneTree", geneTreeWrappers, "populationModel", populationModel);

        double calculatedLogP = msc.calculateLogP();
        assertEquals(expectedLogP, calculatedLogP, allowedError);
    }

    public void initializeSpeciesTree(TaxonSet speciesSuperset) throws Exception {
        speciesTree = new TreeParser();
        speciesTree.initByName("newick", newickSpeciesTree, "IsLabelledNewick", true, "taxonset", speciesSuperset);
        speciesTreeWrapper = new SpeciesTree();
        speciesTreeWrapper.initByName("tree", speciesTree);
    }

    public void initializeGeneTrees() throws Exception {
        for (String geneTreeNewick: newickGeneTrees) {
            TreeParser geneTree = new TreeParser();
            geneTree.initByName("newick", geneTreeNewick, "IsLabelledNewick", true);
            geneTrees.add(geneTree);

            GeneTree geneTreeWrapper = new GeneTree();
            geneTreeWrapper.initByName("tree", geneTree, "ploidy", ploidy, "speciesTree", speciesTreeWrapper);
            geneTreeWrappers.add(geneTreeWrapper);
        }
    }
}

@@ -275,18 +275,24 @@
 alescent
+Single
 ;%0A
-
 import s
@@ -668,41 +668,8 @@
 );%0A%0A
-    MultispeciesCoalescent msc;%0A%0A
     
@@ -1379,16 +1379,132 @@
         
+double calculatedLogP = 0.0;%0A        for (GeneTree gt: geneTreeWrappers) %7B%0A            MultispeciesCoalescentSingle 
 msc = ne
@@ -1527,16 +1527,22 @@
 alescent
+Single
 ();%0A    
@@ -1545,16 +1545,20 @@
         
+ 
+   
 msc.init
@@ -1612,31 +1612,17 @@
 Tree%22, g
-eneTreeWrappers
+t
 , %22popul
@@ -1652,25 +1652,26 @@
 Model);%0A
-%0A
+  
         
 double c
@@ -1654,38 +1654,33 @@
 del);%0A          
-double
+ 
  calculatedLogP 
@@ -1679,16 +1679,17 @@
 tedLogP 
++
 = msc.ca
@@ -1699,24 +1699,35 @@
 lateLogP();%0A
+        %7D%0A%0A
         asse

Update unit test for MultispeciesCoalescentSingle