// $Id: ui_vars_chainparams.cpp,v 1.40 2018/01/03 21:33:05 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 "defaults.h"
#include "timex.h"      // for GetTime()
#include "ui_vars.h"

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

UIVarsChainParameters::UIVarsChainParameters(UIVars * myUIVars)
    :
    UIVarsComponent(myUIVars),
    adaptiveTemperatures            (defaults::useAdaptiveTemperatures),
    chainTemperatures               (defaults::chainTemperatures()),
    dropArrangerRelativeTiming      (defaults::dropArrangerTiming),
    sizeArrangerRelativeTiming      (defaults::sizeArrangerTiming),
    haplotypeArrangerRelativeTiming (defaults::haplotypeArrangerTiming),
    probhapArrangerRelativeTiming   (defaults::probhapArrangerTiming),
    bayesianArrangerRelativeTiming  (FLAGDOUBLE),
    locusArrangerRelativeTiming     (FLAGDOUBLE),
    zilchArrangerRelativeTiming     (0.0),
    stairArrangerRelativeTiming     (0.0),
    epochSizeArrangerRelativeTiming (0.0),
    doBayesianAnalysis              (defaults::useBayesianAnalysis),
    finalChainSamplingInterval      (defaults::finalInterval),
    finalNumberOfChains             (defaults::finalNChains),
    finalNumberOfChainsToDiscard    (defaults::finalDiscard),
    finalNumberOfSamples            (defaults::finalNSamples),
    initialChainSamplingInterval    (defaults::initInterval),
    initialNumberOfChains           (defaults::initNChains),
    initialNumberOfChainsToDiscard  (defaults::initDiscard),
    initialNumberOfSamples          (defaults::initNSamples),
    numberOfReplicates              (defaults::replicates),
    temperatureInterval             (defaults::temperatureInterval)
{
}

UIVarsChainParameters::UIVarsChainParameters(UIVars * myUIVars, const UIVarsChainParameters& chainparams)
    : UIVarsComponent(myUIVars),
      adaptiveTemperatures            (chainparams.adaptiveTemperatures),
      chainTemperatures               (chainparams.chainTemperatures),
      dropArrangerRelativeTiming      (chainparams.dropArrangerRelativeTiming),
      sizeArrangerRelativeTiming      (chainparams.sizeArrangerRelativeTiming),
      haplotypeArrangerRelativeTiming (chainparams.haplotypeArrangerRelativeTiming),
      probhapArrangerRelativeTiming   (chainparams.probhapArrangerRelativeTiming),
      bayesianArrangerRelativeTiming  (chainparams.bayesianArrangerRelativeTiming),
      locusArrangerRelativeTiming     (chainparams.locusArrangerRelativeTiming),
      zilchArrangerRelativeTiming     (chainparams.zilchArrangerRelativeTiming),
      stairArrangerRelativeTiming     (chainparams.stairArrangerRelativeTiming),
      epochSizeArrangerRelativeTiming (chainparams.epochSizeArrangerRelativeTiming),
      doBayesianAnalysis              (chainparams.doBayesianAnalysis),
      finalChainSamplingInterval      (chainparams.finalChainSamplingInterval),
      finalNumberOfChains             (chainparams.finalNumberOfChains),
      finalNumberOfChainsToDiscard    (chainparams.finalNumberOfChainsToDiscard),
      finalNumberOfSamples            (chainparams.finalNumberOfSamples),
      initialChainSamplingInterval    (chainparams.initialChainSamplingInterval),
      initialNumberOfChains           (chainparams.initialNumberOfChains),
      initialNumberOfChainsToDiscard  (chainparams.initialNumberOfChainsToDiscard),
      initialNumberOfSamples          (chainparams.initialNumberOfSamples),
      numberOfReplicates              (chainparams.numberOfReplicates),
      temperatureInterval             (chainparams.temperatureInterval)
{
}

UIVarsChainParameters::~UIVarsChainParameters()
{
}

long int UIVarsChainParameters::GetChainCount() const
{
    return chainTemperatures.size();
}

double UIVarsChainParameters::GetChainTemperature(long int chainId) const
{
    assert(chainId < (long int)chainTemperatures.size());
    return chainTemperatures[chainId];
}

DoubleVec1d UIVarsChainParameters::GetChainTemperatures() const
{
    return chainTemperatures;
}

double UIVarsChainParameters::GetBayesianArrangerRelativeTiming() const
{
    if (GetDoBayesianAnalysis())
    {
        if(bayesianArrangerRelativeTiming <= 0.0)
        {
            assert(bayesianArrangerRelativeTiming == FLAGDOUBLE);
            return GetDropArrangerRelativeTiming();
        }
        else
        {
            return bayesianArrangerRelativeTiming;
        }
    }
    else
    {
        return 0.0;
    }
}

double UIVarsChainParameters::GetLocusArrangerRelativeTiming() const
{
    if (GetConstUIVars().traitmodels.AnyJumpingAnalyses() > 0)
    {
        if(locusArrangerRelativeTiming <= 0.0)
        {
            assert(locusArrangerRelativeTiming == FLAGDOUBLE);
            return dropArrangerRelativeTiming * .2;
        }
        else
        {
            return locusArrangerRelativeTiming;
        }
    }
    else
    {
        return 0.0;
    }
}

double UIVarsChainParameters::GetDropArrangerRelativeTiming() const
{
    if (zilchArrangerRelativeTiming > 0)
    {
        return 0.0;
    }
    return dropArrangerRelativeTiming;
}

double UIVarsChainParameters::GetStairArrangerRelativeTiming() const
{
    return stairArrangerRelativeTiming;
}

double UIVarsChainParameters::GetEpochSizeArrangerRelativeTiming() const
{
    return epochSizeArrangerRelativeTiming;
}

double UIVarsChainParameters::GetSizeArrangerRelativeTiming() const
{
    return sizeArrangerRelativeTiming;
}

double UIVarsChainParameters::GetZilchArrangerRelativeTiming() const
{
    if (GetConstUIVars().datapackplus.AnySimulation())
    {
        return zilchArrangerRelativeTiming;
    }
    else
    {
        return 0.0;
    }
}

bool UIVarsChainParameters::GetHaplotypeArrangerPossible() const
{
    return GetConstUIVars().datapackplus.CanHapArrange();
}

double UIVarsChainParameters::GetHaplotypeArrangerRelativeTiming() const
{
    return haplotypeArrangerRelativeTiming;
}

bool UIVarsChainParameters::GetProbHapArrangerPossible() const
{
    return (GetConstUIVars().traitmodels.AnyJumpingAnalyses() &&
            GetConstUIVars().datapackplus.AnyRelativeHaplotypes());
}

double UIVarsChainParameters::GetProbHapArrangerRelativeTiming() const
{
    return probhapArrangerRelativeTiming;
}

long int UIVarsChainParameters::GetTemperatureInterval() const
{
    return temperatureInterval;
}

bool UIVarsChainParameters::GetAdaptiveTemperatures() const
{
    return adaptiveTemperatures;
}

long int UIVarsChainParameters::GetNumberOfReplicates() const
{
    return numberOfReplicates;
}

long int UIVarsChainParameters::GetInitialNumberOfChains() const
{
    return initialNumberOfChains;
}

long int UIVarsChainParameters::GetInitialNumberOfSamples() const
{
    return initialNumberOfSamples;
}

long int UIVarsChainParameters::GetInitialChainSamplingInterval() const
{
    return initialChainSamplingInterval;
}

long int UIVarsChainParameters::GetInitialNumberOfChainsToDiscard() const
{
    return initialNumberOfChainsToDiscard;
}

long int UIVarsChainParameters::GetFinalNumberOfChains() const
{
    return finalNumberOfChains;
}

long int UIVarsChainParameters::GetFinalNumberOfSamples() const
{
    return finalNumberOfSamples;
}

long int UIVarsChainParameters::GetFinalChainSamplingInterval() const
{
    return finalChainSamplingInterval;
}

long int UIVarsChainParameters::GetFinalNumberOfChainsToDiscard() const
{
    return finalNumberOfChainsToDiscard;
}

bool UIVarsChainParameters::GetDoBayesianAnalysis() const
{
    return doBayesianAnalysis;
}

void UIVarsChainParameters::SetChainCount(long int chainCount)
{
    if(chainCount < 1)
    {
        throw data_error("There must be a positive number of simultaneous searches");
    }
    if (chainCount > defaults::maxNumHeatedChains)
    {
        string err = "There must be less than "
            + ToString(defaults::maxNumHeatedChains)
            + " multiple searches for a given run. "
            + " (A reasonably high number is 5.)";
        throw data_error(err);
    }
    if (chainCount > 5)
    {
        GetConstUIVars().GetUI()->AddWarning("Warning:  a high number of simultaneous heated chains can cause "
                                             "LAMARC to run out of memory.  Five chains is probably a reasonably high upper limit.");
    }
    while((long int)chainTemperatures.size() < chainCount)
    {
        chainTemperatures.push_back(MakeNextChainTemperature(chainTemperatures));
    }
    chainTemperatures.resize(chainCount);
}

void UIVarsChainParameters::SetChainTemperature(double val, long int chainId)
{
    assert(chainId < (long int)chainTemperatures.size());
    if (val <= 0)
    {
        throw data_error("All chain temperatures must be positive.");
    }
    chainTemperatures[chainId] = val;
}

void UIVarsChainParameters::RescaleDefaultSizeArranger()
{
    if (dropArrangerRelativeTiming > 0)
    {
        sizeArrangerRelativeTiming = sizeArrangerRelativeTiming * dropArrangerRelativeTiming;
    }
}

void UIVarsChainParameters::SetBayesianArrangerRelativeTiming(double val)
{
    if (val < 0.0)
    {
        throw data_error("The bayesian rearranger frequency must be positive.");
    }
    if (val == 0.0)
    {
        if (GetDoBayesianAnalysis() == false) return;
        throw data_error("The Bayesian rearranger frequency must be positive when a Bayesian analysis is on.\n"
                         "If you don't want a Bayesian analysis, turn it off in the 'Search Strategy' menu.");
    }
    if (GetDoBayesianAnalysis() == false)
    {
        throw data_error("You are not performing a Bayesian analysis, and therefore may not turn on the\n"
                         "Bayesian arranger (the 'bayesian' tag within the 'strategy' tag)");
    }
    bayesianArrangerRelativeTiming = val;
}

void UIVarsChainParameters::SetLocusArrangerRelativeTiming(double val)
{
    if (val <= 0.0)
    {
        throw data_error("The trait location rearranger frequency must be greater than zero.");
    }
    locusArrangerRelativeTiming = val;
}

void UIVarsChainParameters::SetDropArrangerRelativeTiming(double val)
{
    // shut this test off for newick testing 8/10 JRM
    if (val <= 0)
    {
        throw data_error("The topology rearranger frequency must be greater than zero.");
    }
    zilchArrangerRelativeTiming = 0.0;
    dropArrangerRelativeTiming = val;
}

void UIVarsChainParameters::SetStairArrangerRelativeTiming(double val)
{
    if (val < 0)
    {
        throw data_error("The stair rearranger frequency must be positive.");
    }
    stairArrangerRelativeTiming = val;
}

void UIVarsChainParameters::SetEpochSizeArrangerRelativeTiming(double val)
{
    if (val < 0)
    {
        throw data_error("The epoch-size rearranger frequency must be positive.");
    }
    epochSizeArrangerRelativeTiming = val;
}

void UIVarsChainParameters::SetSizeArrangerRelativeTiming(double val)
{
    if (val < 0)
    {
        throw data_error("The size rearranger frequency must be positive.");
    }
    sizeArrangerRelativeTiming = val;
}

void UIVarsChainParameters::SetHaplotypeArrangerRelativeTiming(double val)
{
    if (val < 0)
    {
        throw data_error("The haplotype rearranger frequency must be positive.");
    }
    haplotypeArrangerRelativeTiming = val;
}

void UIVarsChainParameters::SetProbHapArrangerRelativeTiming(double val)
{
    if (val < 0)
    {
        throw data_error("The trait haplotype rearranger frequency must be positive.");
    }
    probhapArrangerRelativeTiming = val;
}

void UIVarsChainParameters::SetZilchArrangerRelativeTiming(double val)
{
    if (val < 0.0)
    {
        throw data_error("The do-nothing rearranger frequency must be positive.");
    }
    if (GetConstUIVars().datapackplus.AnySimulation())
    {
        zilchArrangerRelativeTiming = val;
    }
    else
    {
        throw data_error("You may not use the do-nothing arranger if you are not simulating data.");
    }
}

void UIVarsChainParameters::SetTemperatureInterval(long int val)
{
    if (val < 1)
    {
        throw data_error("The swapping interval must be one or greater.");
    }
    temperatureInterval = val;
}

void UIVarsChainParameters::SetAdaptiveTemperatures(bool val)
{
    adaptiveTemperatures = val;
    // NOTE: We allow lamarc to use the current (possibly
    // user-set) temperatures as the starting point for
    // adaptive temperatures. Jon says this is fine.
}

void UIVarsChainParameters::SetNumberOfReplicates(long int val)
{
    if (val < 0)
    {
        throw data_error("The number of replicates must be one or greater.");
    }
    if (val == 0)
    {
        val = 1;
        GetConstUIVars().GetUI()->AddWarning("The minimum number of replicates is one, meaning, 'Do one analysis'.  "
                                             "Setting the number of replicates to one, instead of zero.");
    }
    numberOfReplicates = val;
}

void UIVarsChainParameters::SetInitialNumberOfChains(long int val)
{
    if (val < 0)
    {
        throw data_error("You may not have a negative number of initial chains.");
    }
    initialNumberOfChains = val;
    WarnIfNoSamples();
}

void UIVarsChainParameters::SetInitialNumberOfSamples(long int val)
{
    if (val < 0)
    {
        throw data_error("You may not have a negative number of initial samples.");
    }
    initialNumberOfSamples = val;
    WarnIfNoSamples();
}

void UIVarsChainParameters::SetInitialChainSamplingInterval(long int val)
{
    if (val <= 0)
    {
        throw data_error("You must have a positive sampling interval.");
    }
    initialChainSamplingInterval = val;
}

void UIVarsChainParameters::SetInitialNumberOfChainsToDiscard(long int val)
{
    if (val < 0)
    {
        throw data_error("You may not have a negative number of discarded genealogies.");
    }
    initialNumberOfChainsToDiscard = val;
}

void UIVarsChainParameters::SetFinalNumberOfChains(long int val)
{
    if (val < 0)
    {
        throw data_error("You may not have a negative number of final chains.");
    }
    finalNumberOfChains = val;
    WarnIfNoSamples();
}

void UIVarsChainParameters::SetFinalNumberOfSamples(long int val)
{
    if (val < 0)
    {
        throw data_error("You may not have a negative number of final samples.");
    }
    finalNumberOfSamples = val;
    WarnIfNoSamples();
}

void UIVarsChainParameters::SetFinalChainSamplingInterval(long int val)
{
    if (val <= 0)
    {
        throw data_error("You must have a positive sampling interval.");
    }
    finalChainSamplingInterval = val;
}

void UIVarsChainParameters::SetFinalNumberOfChainsToDiscard(long int val)
{
    if (val < 0)
    {
        throw data_error("You may not have a negative number of discarded genealogies.");
    }
    finalNumberOfChainsToDiscard = val;
}

void UIVarsChainParameters::WarnIfNoSamples()
{
    if (((initialNumberOfChains == 0) || (initialNumberOfSamples == 0))
        && ((finalNumberOfChains == 0) || (finalNumberOfSamples == 0)))
    {
        GetConstUIVars().GetUI()->AddWarning("Warning:  LAMARC will never sample any trees.");
    }
}

void UIVarsChainParameters::SetDoBayesianAnalysis(bool val)
{
    if (val && GetConstUIVars().forces.GetForceOnOff(force_REGION_GAMMA))
    {
        throw data_error("Cannot do a Bayesian analysis while attempting to estimate Gamma.");
    }
    if (!val && GetConstUIVars().forces.GetForceOnOff(force_DIVERGENCE))
    {
        throw data_error("Must do a Bayesian analysis while modeling Divergence.");
    }
    doBayesianAnalysis = val;
}

double MakeNextChainTemperature(const vector<double> temperatures)
// this assigns temperatures with constant difference between
// successive temperatures.
{
    if(temperatures.empty())
    {
        assert(false); //How did we get an empty temperature vec?
        return defaults::minTemperature;
    }
    unsigned long int size = temperatures.size();
    if(size == 1)
    {
        return defaults::secondTemperature;
    }
    double diff = temperatures[size-1] - temperatures[size-2];
    return temperatures[size-1] + diff;
}

//____________________________________________________________________________________
