// $Id: timemanager.h,v 1.9 2018/01/03 21:32:58 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 */


#ifndef TIMEMANAGER_H
#define TIMEMANAGER_H

#include <iostream>                     // for debug function PrintStickThetasToFile() and related debugs
#include <map>                          // for MakeStickUsingBranches()

#include "vectorx.h"
#include "constants.h"
#include "stair.h"

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

class ForceParameters;
class Random;
class Force;
class TreeSummary;

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

class TimeManager
{
  public:
    // creation, destruction, copying semantics
    TimeManager();
    TimeManager(const TimeManager& src);
    virtual TimeManager& operator=(const TimeManager& src);
    virtual ~TimeManager() {};

    // SizeAtTimeT functionality
    virtual DoubleVec1d XpartThetasAtT(double t, const ForceParameters& fp) =0;
    virtual DoubleVec1d PartitionThetasAtT(double t, force_type force, const ForceParameters& fp) =0;

    virtual double TimeOfActiveCoal(double tstart, const LongVec1d& lineages,
                                    const ForceParameters& fp, long& targetxpart, double maxtime) =0;

    virtual double TimeOfInactiveCoal(double tstart, const LongVec1d& activeline,
                                      const LongVec1d& inactivelines, const ForceParameters& fp,
                                      long& targetxpart, double maxtime) = 0;

    virtual double TimeOfTraitMutation(double tstart, const LongVec1d& lineages,
                                       const ForceParameters& fp, long& tiptrait,
                                       long& roottrait, double maxtime) = 0;

    // To avoid dynamic_cast, these are defined on the base class even
    // if not useful there.  Most of them are no-ops on the base class.
    virtual bool UsingStick() const;
    virtual void CopyStick(const TimeManager& src);
    virtual void ClearStick();
    virtual void SetStickParameters(const ForceParameters& fp);
    virtual void ScoreStick(TreeSummary& treesum) const;
    virtual void ChopOffStickAt(double lasttime);
    virtual void MakeStickUsingBranches(const ForceParameters& fp,
                                        const std::vector<std::pair<double,LongVec1d> >& lpcounts);
    virtual void MakeStickTilTime(const ForceParameters& fp, double);

    // Debugging functions.
    virtual void PrintStickThetasToFile(std::ofstream& of) const;
    virtual void PrintStickFreqsToFile(std::ofstream& of) const;
    virtual void PrintStickFreqsToFileAtTime(std::ofstream& of, double time) const;
    virtual void PrintStickThetasToFileForJoint300(std::ofstream& of) const;
    virtual void PrintStickFreqsToFileForJoint300(std::ofstream& of) const;
    virtual long GetStairSize() const { return 0L; };
    virtual void PrintStickToFile(std::ofstream& of) const;

  protected:
    // convenience reference to random number generator
    Random& randomSource;
};

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

class ConstantTimeManager : public TimeManager
{
  public:
    // SizeAtTimeT functionality
    virtual DoubleVec1d XpartThetasAtT(double t, const ForceParameters& fp);
    virtual DoubleVec1d PartitionThetasAtT(double t, force_type force,
                                           const ForceParameters& fp);

    virtual double TimeOfActiveCoal(double tstart, const LongVec1d& lineages,
                                    const ForceParameters& fp, long& targetxpart, double);

    virtual double TimeOfInactiveCoal(double tstart, const LongVec1d&
                                      activelines, const LongVec1d& inactivelines, const ForceParameters& fp,
                                      long& targetxpart, double);

    virtual double TimeOfTraitMutation(double tstart, const LongVec1d& lineages,
                                       const ForceParameters& fp, long& tiptrait, long& roottrait, double);
};

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

class ExpGrowTimeManager : public TimeManager
{
  public:
    // SizeAtTimeT functionality
    virtual DoubleVec1d XpartThetasAtT(double t, const ForceParameters& fp);
    virtual DoubleVec1d PartitionThetasAtT(double t, force_type force,
                                           const ForceParameters& fp);

    virtual double TimeOfActiveCoal(double tstart, const LongVec1d& lineages,
                                    const ForceParameters& fp, long& targetxpart, double);

    virtual double TimeOfInactiveCoal(double tstart, const LongVec1d&
                                      activelines, const LongVec1d& inactivelines, const ForceParameters& fp,
                                      long& targetxpart, double);

    virtual double TimeOfTraitMutation(double tstart, const LongVec1d& lineages,
                                       const ForceParameters& fp, long& tiptrait, long& roottrait, double);
};

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

class LogSelectTimeManager : public TimeManager
{

  public:
    // SizeAtTimeT functionality
    virtual DoubleVec1d XpartThetasAtT(double t, const ForceParameters& fp);
    virtual DoubleVec1d PartitionThetasAtT(double t, force_type force,
                                           const ForceParameters& fp);

    virtual double TimeOfActiveCoal(double tstart, const LongVec1d& lineages,
                                    const ForceParameters& fp, long& targetxpart, double);

    virtual double TimeOfInactiveCoal(double tstart, const LongVec1d&
                                      activelines, const LongVec1d& inactivelines, const ForceParameters& fp,
                                      long& targetxpart, double);

    virtual double TimeOfTraitMutation(double tstart, const LongVec1d& lineages,
                                       const ForceParameters& fp, long& tiptrait, long& roottrait, double);

  private:
    double LogisticSelectionPickTime(const double& theta_A0, const double& theta_a0,
                                     const double& s, const double& t_s, const double& dtau_A,
                                     const double& dtau_a, long* pChosenPopulation);
};

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

class StickTimeManager : public TimeManager
{
  public:
    virtual TimeManager& operator=(const TimeManager& src);
    virtual bool UsingStick() const;
    virtual void CopyStick(const TimeManager& src);
    virtual void ClearStick();
    virtual void SetStickParameters(const ForceParameters& fp);
    virtual void ScoreStick(TreeSummary& treesum) const;
    virtual void ChopOffStickAt(double lasttime);
    virtual void MakeStickUsingBranches(const ForceParameters& fp,
                                        const std::vector<std::pair<double,LongVec1d> >& lpcounts);
    virtual void MakeStickTilTime(const ForceParameters& fp, double endtime);

    // Debugging functions.
    virtual void PrintStickThetasToFile(std::ofstream& of) const;
    virtual void PrintStickFreqsToFile(std::ofstream& of) const;
    virtual void PrintStickFreqsToFileAtTime(std::ofstream& of, double time) const;
    virtual void PrintStickThetasToFileForJoint300(std::ofstream& of) const;
    virtual void PrintStickFreqsToFileForJoint300(std::ofstream& of) const;
    virtual long GetStairSize() const { return m_stair.size(); };

    // Debugging function.
    // PrintStickToFile() is meant to match the output of TreeSummary::PrintStickSummaryToFile()
    virtual void PrintStickToFile(std::ofstream& of) const;

    virtual double TimeOfActiveCoal(double tstart, const LongVec1d& lineages,
                                    const ForceParameters& fp, long& targetxpart, double maxtime);

    virtual double TimeOfInactiveCoal(double tstart, const LongVec1d&
                                      activelines, const LongVec1d& inactivelines, const ForceParameters& fp,
                                      long& targetxpart, double maxtime);

    virtual double TimeOfTraitMutation(double tstart, const LongVec1d& lineages,
                                       const ForceParameters& fp, long& tiptrait, long& roottrait, double maxtime);
    virtual DoubleVec1d XpartThetasAtT(double t, const ForceParameters& fp);

    virtual DoubleVec1d PartitionThetasAtT(double t, force_type force,
                                           const ForceParameters& fp);

  protected:
    // the stick
    stair m_stair;

    // FindJointAtTime() is non-const because it may have to add joints
    // when the stick is too short
    stair::iterator FindJointAtTime(double time, const ForceParameters& fp);
    stair::iterator NextJointAppendingIfLast(const stair::iterator& start,
                                             const ForceParameters& fp );

    void AppendRiser(const ForceParameters& fp);
    void AppendRiser(const ForceParameters& fp, const DoubleVec1d& discounts);

    virtual double NextRootTime(double timenow, const ForceParameters& fp)
        const = 0;

    virtual DoubleVec1d NextThetas(double nexttime, const ForceParameters&
                                   fp) const = 0;

    virtual DoubleVec1d NextThetas(double nexttime, const ForceParameters&
                                   fp, const DoubleVec1d& discounts) const = 0;

    // debug function
    void PrintStick() const;

  private:
    bool InactiveCoalImpossible(const LongVec1d& alines, const LongVec1d& ilines)
        const;

    bool MutationImpossible(const LongVec1d& lineages, const DoubleVec1d& mus)
        const;

};

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

class ExpGrowStickTimeManager : public StickTimeManager
{
  public:
    ExpGrowStickTimeManager();
    // we accept StickTimeManager's defaults for the public interface
    // --TimeOfActiveCoal,TimeOfInactiveCoal,XpartThetasAtT,
    //   PartitionThetasAtT,TimeOfTraitMutation

  protected:
    virtual double NextRootTime(double timenow, const ForceParameters& fp)
        const;

    virtual DoubleVec1d NextThetas(double nexttime, const ForceParameters& fp)
        const;

    virtual DoubleVec1d NextThetas(double nexttime, const ForceParameters& fp,
                                   const DoubleVec1d& discounts) const;

  private:
    double m_percentchange;   // how much change allowed in theta per joint
    double m_negln, m_posln;  // used in calculating new times for stick joints
};

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

class SelectStickTimeManager : public StickTimeManager
{
  public:
    // we accept StickTimeManager's defaults for the public interface
    // --TimeOfActiveCoal,TimeOfInactiveCoal,XpartThetasAtT,
    //   PartitionThetasAtT,TimeOfTraitMutation

  protected:
    virtual double NextRootTime(double timenow, const ForceParameters& fp)
        const;

    virtual DoubleVec1d NextThetas(double nexttime, const ForceParameters& fp)
        const;

    virtual DoubleVec1d NextThetas(double nexttime, const ForceParameters& fp,
                                   const DoubleVec1d& discounts) const;

  private:
    double Mean_delta_x(double x, double s, double mu, double nu,
                        double delta_t) const;

    double SD_delta_x(double x, double theta, double delta_t) const;

    DoubleVec1d CalculateThetasWithPrevAndNext(stair::const_iterator prevjoint,
                                               stair::const_iterator chosen, stair::const_iterator nextjoint,
                                               const ForceParameters& fp) const;

    DoubleVec1d CalculateThetasWithPrevOnly(stair::const_iterator prevjoint,
                                            const ForceParameters& fp) const;

    DoubleVec1d CalculateThetasWithPrevOnly(stair::const_iterator prevjoint,
                                            const ForceParameters& fp, const DoubleVec1d& discounts) const;

    DoubleVec1d PickNewThetasWithMeanAndSD(double mean, double sd,
                                           const ForceParameters& fp) const;
};

#endif // TIMEMANAGER_H

//____________________________________________________________________________________
