// $Id: arrangervec.cpp,v 1.35 2018/01/03 21:33:03 mkkuhner Exp $

/*
Copyright 2002-2005 Mary K. Kuhner, Peter Beerli, and Joseph Felsenstein

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

*/

/* Authors: Peter Beerli, Mary Kuhner, Jon Yamato and Joseph Felsenstein */


#include <cassert>
#include <iostream>
#include <map>
#include <vector>

#include "local_build.h"

#include "arranger.h"
#include "arranger_types.h"
#include "arrangervec.h"
#include "constants.h"
#include "registry.h"
#include "stringx.h"

//------------------------------------------------------------------------------------

ArrangerVec::ArrangerVec(double dropTiming, double sizeTiming,
                         double hapTiming, double probhapTiming,
                         double bayesTiming, double locusTiming,
                         double zilchTiming, double stairTiming,
                         double epochsizeTiming, double epochnudgeTiming)
{
    assert(dropTiming >= 0.0);
    assert(sizeTiming >= 0.0);
    assert(hapTiming >= 0.0);
    assert(probhapTiming >= 0.0);
    assert(bayesTiming >= 0.0);
    assert(locusTiming >= 0.0);
    assert(zilchTiming >= 0.0);
    assert(stairTiming >= 0.0);
    assert(epochsizeTiming >= 0.0);
    assert(epochnudgeTiming >= 0.0);
    double denovoTiming = 0.0;

#ifdef STATIONARIES
#ifdef ALL_ARRANGERS_DENOVO // Stationaries with denovo arranger in ALL cases.
    denovoTiming = dropTiming;
    dropTiming = 0.0;
#else // Stationaries with denovo arranger only in Bayesian case.
    if (bayesTiming > 0)
    {
        denovoTiming = dropTiming;
        dropTiming = 0.0;
    }
#endif // Stationaries: Bayesian versus all cases.
#endif // STATIONARIES

    const ForceSummary & fs(registry.GetForceSummary());

    arrangers[arrangerstrings::DENO]     = new DenovoArranger(fs, denovoTiming);
    arrangers[arrangerstrings::DROP]     = new DropArranger(fs, dropTiming);
    arrangers[arrangerstrings::TREESIZE] = new TreeSizeArranger(fs, sizeTiming);
    arrangers[arrangerstrings::HAP]      = new HapArranger(hapTiming);
    arrangers[arrangerstrings::PROBHAP]  = new ProbHapArranger(probhapTiming);
    arrangers[arrangerstrings::BAYES]    = new BayesArranger(bayesTiming);
    arrangers[arrangerstrings::LOCUS]    = new LocusArranger(locusTiming);
    arrangers[arrangerstrings::ZILCH]    = new ZilchArranger(zilchTiming);
    arrangers[arrangerstrings::STICK]    = new StairArranger(stairTiming);
    arrangers[arrangerstrings::EPOCHSIZE] = new EpochSizeArranger(fs, epochsizeTiming);
//    commented out because implementation is incomplete  JY 2012/11/06
//    arrangers[arrangerstrings::EPOCHNUDGE] = new EpochNudgeArranger(fs, epochnudgeTiming);

    Normalize();
}

//------------------------------------------------------------------------------------

ArrangerVec::~ArrangerVec()
{
    ClearAll();
}

//------------------------------------------------------------------------------------

ArrangerVec::ArrangerVec(const ArrangerVec& src)
{
    CopyAllMembers(src);
} // copy ctor

//------------------------------------------------------------------------------------

ArrangerVec& ArrangerVec::operator=(const ArrangerVec& src)
{
    CopyAllMembers(src);
    return *this;
} // op=

//------------------------------------------------------------------------------------

double ArrangerVec::GetArrangerTiming(const string & atype) const
{
    if(arrangers.find(atype) != arrangers.end())
    {
        const_arrangerit it = arrangers.find(atype);
        return it->second->GetTiming();
    }
    else
    {
        throw implementation_error("Missing arranger of type " + atype);
        return -1.0;
    }
}

//------------------------------------------------------------------------------------

void ArrangerVec::CopyAllMembers(const ArrangerVec & cp)
{
    if (this != &cp)
    {
        ClearAll();
        const_arrangerit it;
        for(it = cp.arrangers.begin(); it != cp.arrangers.end(); it++)
        {
            Arranger* arr = it->second->Clone();
            arrangers.insert(std::make_pair(arr->GetName(), arr));
        }
    }
}

//------------------------------------------------------------------------------------

void ArrangerVec::Normalize()
{
    const_arrangerit ait;
    double totalTime = 0.0;
    for(ait = arrangers.begin(); ait != arrangers.end(); ++ait)
    {
        double thisTime = ((*ait).second)->GetTiming();
        assert(thisTime >= 0.0);
        totalTime += thisTime;
    }
    assert(totalTime > 0.0);
    for(ait = arrangers.begin(); ait != arrangers.end(); ++ait)
    {
        double thisTime = ((*ait).second)->GetTiming();
        ((*ait).second)->SetTiming(thisTime / totalTime);
    }
}

//------------------------------------------------------------------------------------

unsigned long ArrangerVec::size() const
{
    return arrangers.size();
}

//------------------------------------------------------------------------------------

bool ArrangerVec::empty() const
{
    return arrangers.empty();
}

//------------------------------------------------------------------------------------

Arranger* ArrangerVec::GetDenovoArranger() const
{
    const_arrangerit it = arrangers.find(arrangerstrings::DENO);
    if (it == arrangers.end())
    {
        impossible_error e("Arranger needed to make Denovo tree is missing");
        throw e;
    }
    return it->second;
} // GetDenovoArranger

//------------------------------------------------------------------------------------

Arranger* ArrangerVec::GetRandomArranger(double rand) const
{
#if 0 // start JREMOVE
    static long int denovo;
    if (denovo == 32) denovo = 33;
    else denovo = 32;

    const_arrangerit it;
    if (denovo == 32) it = arrangers.find(arrangerstrings::DENO);
    else it = arrangers.find(arrangerstrings::EPOCHSIZE);

    return it->second;

#endif // end JREMOVE

#if 1 // JRESTORE
    const_arrangerit it;
    double timing;

    assert(arrangers.size() > 0);

    for (it = arrangers.begin(); it != arrangers.end(); ++it)
    {
        timing = it->second->GetTiming();
        if (timing > rand) return it->second;
        rand -= timing;
    }

    // in case of rounding error, we might flow through to here
    for (it = arrangers.begin(); it != arrangers.end(); ++it)
    {
        if (it->second->GetTiming() > 0.0) return it->second;
    }
#endif

    // no arrangers found!
    assert(false);
    return NULL;
} // GetRandomArranger

//------------------------------------------------------------------------------------

StringVec1d ArrangerVec::GetAllStringsForActiveArrangers() const
{
    StringVec1d retvec;
    for (const_arrangerit arr=arrangers.begin(); arr != arrangers.end(); arr++)
    {
        if ((*arr).second->GetTiming() > 0)
        {
            retvec.push_back((*arr).first);
        }
    }
    return retvec;
}

//------------------------------------------------------------------------------------

void ArrangerVec::ClearAll()
{
    arrangerit it;

    for (it = arrangers.begin(); it != arrangers.end(); ++it)
    {
        delete it->second;
    }
    arrangers.clear();
} // ClearAll

//------------------------------------------------------------------------------------

void ArrangerVec::ZeroHapArranger()
{
    if (arrangers[arrangerstrings::HAP]->GetTiming() != 0.0)
    {
        for (arrangerit arr = arrangers.begin(); arr != arrangers.end(); ++arr)
        {
            (*arr).second->SetSaveTiming();
        }
        arrangers[arrangerstrings::HAP]->SetTiming(0.0);
        Normalize();
    }
}

//------------------------------------------------------------------------------------

void ArrangerVec::ZeroProbHapArranger()
{
    if (arrangers[arrangerstrings::PROBHAP]->GetTiming() != 0.0)
    {
        for (arrangerit arr = arrangers.begin(); arr != arrangers.end(); ++arr)
        {
            (*arr).second->SetSaveTiming();
        }
        arrangers[arrangerstrings::PROBHAP]->SetTiming(0.0);
        Normalize();
    }
}

//------------------------------------------------------------------------------------

void ArrangerVec::ZeroLocusArranger()
{
    if (arrangers[arrangerstrings::LOCUS]->GetTiming() != 0.0)
    {
        for (arrangerit arr = arrangers.begin(); arr != arrangers.end(); ++arr)
        {
            (*arr).second->SetSaveTiming();
        }
        arrangers[arrangerstrings::LOCUS]->SetTiming(0.0);
        Normalize();
    }
}

//------------------------------------------------------------------------------------

void ArrangerVec::RestoreTiming()
{
    for (arrangerit arr = arrangers.begin(); arr != arrangers.end(); ++arr)
    {
        (*arr).second->RestoreTiming();
    }
    Normalize();
}

//------------------------------------------------------------------------------------

void ArrangerVec::PrintArrangers() const
{
    for (arrangermap::const_iterator arr = arrangers.begin(); arr != arrangers.end(); ++arr)
    {
        std::cerr << (*arr).first << ": " << ToString((*arr).second->GetTiming()) << std::endl;
    }
}

//____________________________________________________________________________________
