// $Id: forceparam.cpp,v 1.43 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 */


#include <cassert>
#include <iostream>

#include "forceparam.h"
#include "forcesummary.h"
#include "force.h"
#include "mathx.h"
#include "region.h"
#include "sumfilehandler.h"
#include "vectorx.h"                    // for LongSquareRootOfLong
#include "xmlsum_strings.h"             // for xml sumfile handling

#ifdef DMALLOC_FUNC_CHECK
#include "/usr/local/include/dmalloc.h"
#endif

using namespace std;

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

ForceParameters::ForceParameters(long region)
    : m_region(region),
      m_space(known_region),
      m_forceTags(registry.GetForceSummary().GetForceTags()),
      m_forceSizes(registry.GetForceSummary().GetForceSizes()),
      m_epochptr(registry.GetForceSummary().GetEpochs())
{
    assert (region >= 0);
}

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

ForceParameters::ForceParameters(param_space space)
    : m_region(FLAGLONG),
      m_space(space),
      m_forceTags(registry.GetForceSummary().GetForceTags()),
      m_forceSizes(registry.GetForceSummary().GetForceSizes())
{
    assert (m_space != known_region); //can be either global or unknown
}

//------------------------------------------------------------------------------------
//Need this constructor for the registry, when we don't have a forcesummary yet.

ForceParameters::ForceParameters(param_space space, ForceTypeVec1d types, LongVec1d sizes)
    : m_region(FLAGLONG),
      m_space(space),
      m_forceTags(types),
      m_forceSizes(sizes)
{
    assert (m_space != known_region);   // Can be either global or unknown.
}

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

ForceParameters::ForceParameters(const ForceParameters &fp, long region)
    : m_region(region),
      m_space(known_region),
      m_forceTags(fp.m_forceTags),
      m_forceSizes(fp.m_forceSizes)
{
    assert (region >= 0);
    switch (fp.GetParamSpace())
    {
        case unknown_region:
            CopyRegionalMembers(fp);
            FillGlobalParamsFromRegionalParams();
            break;
        case known_region:
            if (fp.m_region != region)
            {
                assert(false); //Why are we assigning fps for different regions to each other?
            }
            //else fall through to:
        case global_region:
            CopyGlobalMembers(fp);
            FillRegionalParamsFromGlobalParams();
            break;
    }
}

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

void ForceParameters::CopyGlobalMembers(const ForceParameters& fp)
{
    m_globalThetas = fp.m_globalThetas;
    m_migrates = fp.m_migrates;
    m_recrates = fp.m_recrates;
    m_growths  = fp.m_growths;
    m_diseases = fp.m_diseases;
    m_logisticSelectionCoefficient = fp.m_logisticSelectionCoefficient;
    m_epochtimes = fp.m_epochtimes;
} // ForceParameters::CopyGlobalMembers

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

void ForceParameters::CopyRegionalMembers(const ForceParameters& fp)
{
    m_regionalThetas = fp.m_regionalThetas;
    m_migrates = fp.m_migrates;
    m_recrates = fp.m_recrates;
    m_growths  = fp.m_growths;
    m_diseases = fp.m_diseases;
    m_logisticSelectionCoefficient = fp.m_logisticSelectionCoefficient;
    m_epochtimes = fp.m_epochtimes;
} // ForceParameters::CopyRegionalMembers

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

void ForceParameters::SetGlobalParametersByTag(force_type tag, const DoubleVec1d& v)
{
    assert (m_space != unknown_region);
    string msg = "ForceParameters::SetGlobalParametersByTag() received tag ";
    switch (tag)
    {
        case force_COAL:
            SetGlobalThetas(v);
            return;
        case force_MIG:
        case force_DIVMIG:
            SetMigRates(v);
            return;
        case force_REC:
            SetRecRates(v);
            return;
        case force_EXPGROWSTICK:
        case force_GROW:
            SetGrowthRates(v);
            return;
        case force_LOGSELECTSTICK:
        case force_LOGISTICSELECTION:
            SetLogisticSelectionCoefficient(v);
            return;
        case force_DISEASE:
            SetDiseaseRates(v);
            return;
        case force_DIVERGENCE:
            SetEpochTimes(v);
            return;
        case force_REGION_GAMMA:
            msg += "\"" + ToString(tag) + ",\" which should never happen. (This is a ";
            msg += "pseudo-force that\'s only treated as a force within certain ";
            msg += "contexts of the program.)";
            throw implementation_error(msg);
        case force_NONE:
            msg += "\"" + ToString(tag) + ",\" which should never happen.";
            throw implementation_error(msg);
    }

    msg += "\"" + ToString(tag) + ",\" but it does not know how to use it.";
    throw implementation_error(msg);
} // SetGlobalParametersByTag

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

void ForceParameters::SetRegionalParametersByTag(force_type tag, const DoubleVec1d& v)
{
    assert (m_space != global_region);
    string msg = "ForceParameters::SetRegionalParametersByTag() received tag ";
    switch (tag)
    {
        case force_COAL:
            SetRegionalThetas(v);
            return;
        case force_MIG:
        case force_DIVMIG:
            SetMigRates(v);
            return;
        case force_REC:
            SetRecRates(v);
            return;
        case force_EXPGROWSTICK:
        case force_GROW:
            SetGrowthRates(v);
            return;
        case force_LOGSELECTSTICK:
        case force_LOGISTICSELECTION:
            SetLogisticSelectionCoefficient(v);
            return;
        case force_DISEASE:
            SetDiseaseRates(v);
            return;
        case force_DIVERGENCE:
            SetEpochTimes(v);
            return;
        case force_REGION_GAMMA:
            msg += "\"" + ToString(tag) + ",\" which should never happen. (This is a ";
            msg += "pseudo-force that\'s only treated as a force within certain ";
            msg += "contexts of the program.)";
            throw implementation_error(msg);
        case force_NONE:
            msg += "\"" + ToString(tag) + ",\" which should never happen.";
            throw implementation_error(msg);
    }

    msg += "\"" + ToString(tag) + ",\" but it does not know how to use it.";
    throw implementation_error(msg);
} // SetRegionalParametersByTag

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

void ForceParameters::SetGlobalThetas(const DoubleVec1d& v)
{
    assert(m_space != unknown_region);
    m_globalThetas = v;
    if (m_space == known_region)
    {
        FillRegionalParamsFromGlobalParams();
    }
} // SetThetas

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

void ForceParameters::SetRegionalThetas(const DoubleVec1d& v)
{
    assert(m_space != global_region);
    m_regionalThetas = v;
    if (m_space == known_region)
    {
        FillGlobalParamsFromRegionalParams();
    }
} // SetThetas

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

void ForceParameters::SetMigRates(const DoubleVec1d& v)
{
    long rowsize = LongSquareRootOfLong(v.size());

    // put the given values into migrates, but never put a
    // non-zero value into the diagonal entries!

    long i, j;
    long index = 0;
    m_migrates.clear();
    for (i = 0; i < rowsize; ++i)
    {
        for (j = 0; j < rowsize; ++j)
        {
            if (i == j)
            {
                m_migrates.push_back(0.0);
                assert(v[index] == 0.0);
            }
            else m_migrates.push_back(v[index]);
            ++index;
        }
    }
} // SetMigRates

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

void ForceParameters::SetRecRates(const DoubleVec1d& v)
{
    assert(static_cast<long> (v.size()) == 1); // Program supports only 1 recombination rate.
    m_recrates = v;
    //LS NOTE:  if this gets split into SetRegionalRecRates and SetGlobalRecRates
    // we will need to also fill the appropriate parallel vector if m_space is
    // known_region (see SetGlobalThetas)
} // SetRecRates

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

void ForceParameters::SetGrowthRates(const DoubleVec1d& v)
{
    m_growths = v;
} // SetGrowths

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

void ForceParameters::SetLogisticSelectionCoefficient(const DoubleVec1d& v)
{
    if (1 != v.size())
    {
        string msg = "ForceParameters::SetLogisticSelectionCoefficient() received ";
        msg += "a vector of " + ToString(v.size()) + " elements.  This vector must ";
        msg += "contain exactly one element, corresponding to the selection coefficient.";
        throw implementation_error(msg);
    }

    m_logisticSelectionCoefficient = v;
} // SetLogisticSelectionCoefficient

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

void ForceParameters::SetDiseaseRates(const DoubleVec1d& v)
{
    long rowsize = LongSquareRootOfLong(v.size());

    // put the given values into disease rates, but never put a
    // non-zero value into the diagonal entries!

    long i, j;
    long index = 0;
    m_diseases.clear();
    for (i = 0; i < rowsize; ++i)
    {
        for (j = 0; j < rowsize; ++j)
        {
            if (i == j)
            {
                m_diseases.push_back(0.0);
                assert(v[index] == 0.0);
            }
            else
            {
                m_diseases.push_back(v[index]);
            }
            ++index;
        }
    }
} // SetDiseaseRates

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

void ForceParameters::SetEpochTimes(const DoubleVec1d& v)
{
    m_epochtimes = v;
}

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

void ForceParameters::SetGlobalParameters(const DoubleVec1d& v)
{
    assert(m_space != unknown_region);
    SetParameters(v, true);
} // SetGlobalParameters

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

void ForceParameters::SetRegionalParameters(const DoubleVec1d& v)
{
    assert(m_space != global_region);
    SetParameters(v, false);
} // SetRegionalParameters

//------------------------------------------------------------------------------------
// SetParameters is private; use the two above functions.

void ForceParameters::SetParameters(const DoubleVec1d& v, bool isGlobal)
{
    for (unsigned long fnum = 0, pnum = 0; fnum < m_forceTags.size(); fnum++)
    {
        DoubleVec1d oneForceVec;
        for (long fpnum = 0; fpnum < m_forceSizes[fnum]; pnum++, fpnum++)
        {
            assert (v.size() > static_cast<unsigned long>(pnum));
            oneForceVec.push_back(v[pnum]);
        }
        if (isGlobal)
        {
            SetGlobalParametersByTag(m_forceTags[fnum], oneForceVec);
        }
        else
        {
            SetRegionalParametersByTag(m_forceTags[fnum], oneForceVec);
        }
    }
}

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

const DoubleVec1d& ForceParameters::GetGlobalThetas()  const
{
    assert(m_space != unknown_region);
    return m_globalThetas;
}

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

const DoubleVec1d& ForceParameters::GetRegionalThetas()const
{
    assert(m_space != global_region);
    return m_regionalThetas;
}

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

DoubleVec1d ForceParameters::GetGlobalLogParameters() const
{
    assert(m_space != unknown_region);
    return GetLogParameters(true);
}

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

DoubleVec1d ForceParameters::GetRegionalLogParameters() const
{
    assert(m_space != global_region);
    return GetLogParameters(false);
}

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

DoubleVec1d ForceParameters::GetOnlyDiseaseRates() const
{
    DoubleVec1d returnvec;
    long nallparams(m_diseases.size());
    long ncontigparams(sqrt(nallparams));
    long ind(1); // skip the first element as it is always invalid
    while (ind < nallparams)
    {
        long startind(ind);
        for(; ind < startind+ncontigparams; ++ind)
            returnvec.push_back(m_diseases[ind]);
        ++ind; // there's always 1 invalid param between sets of valid
        // params
    }
    assert(static_cast<long>(returnvec.size()) == ncontigparams * (ncontigparams-1));
    return returnvec;
}

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

//GetLogParameters is private; use the above two calls in code.
DoubleVec1d ForceParameters::GetLogParameters(bool isGlobal) const
{
    DoubleVec1d tempvec;
    DoubleVec1d resultvec;

    for (unsigned long fnum=0; fnum < m_forceTags.size(); fnum++)
    {
        if (isGlobal)
        {
            tempvec = GetGlobalParametersByTag(m_forceTags[fnum]);
        }
        else
        {
            tempvec = GetRegionalParametersByTag(m_forceTags[fnum]);
        }
        if (m_forceTags[fnum] != force_GROW)
        {
            LogVec0(tempvec, tempvec);
        }
        resultvec.insert(resultvec.end(), tempvec.begin(), tempvec.end());
    }

    return resultvec;
} // GetLogParameters

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

double ForceParameters::GetRegionalLogParameter(long pnum) const
{
    DoubleVec1d tempvec = GetRegionalParameters();
    const ParamVector paramvec(true);
    if (paramvec[pnum].IsForce(force_GROW))
    {
        return tempvec[pnum];
    }
    else
    {
        return SafeLog(tempvec[pnum]);
    }
}

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

DoubleVec1d ForceParameters::GetGlobalParameters() const
{
    assert(m_space != unknown_region);
    return GetParameters(true);
}

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

DoubleVec1d ForceParameters::GetRegionalParameters() const
{
    assert(m_space != global_region);
    return GetParameters(false);
}

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

DoubleVec1d ForceParameters::GetParameters(bool isGlobal) const
{
    DoubleVec1d tempvec;
    DoubleVec1d resultvec;

    for (unsigned long fnum=0; fnum < m_forceTags.size(); fnum++)
    {
        if (isGlobal)
        {
            tempvec = GetGlobalParametersByTag(m_forceTags[fnum]);
        }
        else
        {
            tempvec = GetRegionalParametersByTag(m_forceTags[fnum]);
        }
        resultvec.insert(resultvec.end(), tempvec.begin(), tempvec.end());
    }

    return resultvec;
} // GetParameters

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

const DoubleVec1d& ForceParameters::GetRegionalParametersByTag(force_type tag) const
{
    assert(m_space != global_region);
    string msg = "ForceParameters::GetRegionalParametersByTag() received tag ";
    switch (tag)
    {
        case force_COAL:
            return GetRegionalThetas();
        case force_MIG:
        case force_DIVMIG:
            return GetMigRates();
        case force_REC:
            return GetRecRates();
        case force_EXPGROWSTICK:
        case force_GROW:
            return GetGrowthRates();
        case force_DISEASE:
            return GetDiseaseRates();
        case force_LOGSELECTSTICK:
        case force_LOGISTICSELECTION:
            return GetLogisticSelectionCoefficient();
        case force_DIVERGENCE:
            return GetEpochTimes();
        case force_REGION_GAMMA:
            msg += "\"" + ToString(tag) + ",\" which should never happen. (This is a ";
            msg += "pseudo-force that\'s only treated as a force within certain ";
            msg += "contexts of the program.)";
            throw implementation_error(msg);
        case force_NONE:
            msg += "\"" + ToString(tag) + ",\" which should never happen.";
            throw implementation_error(msg);
    }

    msg += "\"" + ToString(tag) + ",\" but it does not know how to use it.";
    throw implementation_error(msg);
} // GetParametersByTag

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

const DoubleVec1d& ForceParameters::GetGlobalParametersByTag(force_type tag) const
{
    assert(m_space != unknown_region);
    string msg = "ForceParameters::GetGlobalParametersByTag() received tag ";
    switch (tag)
    {
        case force_COAL:
            return GetGlobalThetas();
        case force_MIG:
        case force_DIVMIG:
            return GetMigRates();
        case force_REC:
            return GetRecRates();
        case force_EXPGROWSTICK:
        case force_GROW:
            return GetGrowthRates();
        case force_LOGSELECTSTICK:
        case force_LOGISTICSELECTION:
            return GetLogisticSelectionCoefficient();
        case force_DISEASE:
            return GetDiseaseRates();
        case force_DIVERGENCE:
            return GetEpochTimes();
        case force_REGION_GAMMA:
            msg += "\"" + ToString(tag) + ",\" which should never happen. (This is a ";
            msg += "pseudo-force that\'s only treated as a force within certain ";
            msg += "contexts of the program.)";
            throw implementation_error(msg);
        case force_NONE:
            msg += "\"" + ToString(tag) + ",\" which should never happen.";
            throw implementation_error(msg);
    }

    msg += "\"" + ToString(tag) + ",\" but it does not know how to use it.";
    throw implementation_error(msg);
} // GetParametersByTag

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

DoubleVec2d ForceParameters::GetGlobal2dRates(force_type tag) const
{
    assert(tag == force_DISEASE || tag == force_MIG || tag == force_DIVMIG);
    return SquareOffVector(GetGlobalParametersByTag(tag));
} // Get2dRates

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

DoubleVec2d ForceParameters::GetRegional2dRates(force_type tag) const
{
    assert(tag == force_DISEASE || tag == force_MIG || tag == force_DIVMIG);
    return SquareOffVector(GetRegionalParametersByTag(tag));
} // Get2dRates

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

void ForceParameters::FillRegionalParamsFromGlobalParams()
{
    if (m_region == FLAGLONG)
    {
        assert(false);
        throw implementation_error("Unable to convert global parameters to regional parameters"
                                   " because we don't have a region to convert to.");
    }
    double effPopSize = registry.GetDataPack().GetRegion(m_region).GetEffectivePopSize();

    m_regionalThetas = m_globalThetas;
    transform(m_globalThetas.begin(),
              m_globalThetas.end(),
              m_regionalThetas.begin(),
              bind2nd(multiplies<double>(),effPopSize));
}

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

void ForceParameters::FillGlobalParamsFromRegionalParams()
{
    if (m_region == FLAGLONG)
    {
        assert(false);
        throw implementation_error("Unable to convert regional parameters to global parameters"
                                   " because we don't know what region we're converting from.");
    }

    double effPopSize = registry.GetDataPack().GetRegion(m_region).GetEffectivePopSize();

    m_globalThetas = m_regionalThetas;
    transform(m_regionalThetas.begin(),
              m_regionalThetas.end(),
              m_globalThetas.begin(),
              bind2nd(divides<double>(),effPopSize));
}

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

DoubleVec1d ForceParameters::GetRegionalScalars() const
{
    DoubleVec1d resultvec;
    for (unsigned long force=0; force<m_forceTags.size(); force++)
    {
        double val = 1.0;
        if (m_forceTags[force] == force_COAL)
        {
            val = registry.GetDataPack().GetRegion(m_region).GetEffectivePopSize();
        }
        if (force_REGION_GAMMA == m_forceTags[force])
            continue; // should never hit this, but let's put it here for safety
        DoubleVec1d tempvec(m_forceSizes[force], val);
        resultvec.insert(resultvec.end(), tempvec.begin(), tempvec.end());
    }

    return resultvec;
} // GetRegionalScalars

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

void ForceParameters::WriteForceParameters( ofstream& sumout, long numtabs) const
{
    string tabs(numtabs, '\t');
    if ( sumout.is_open() )
    {
        vector<double> vd;

        sumout << tabs << xmlsum::ESTIMATES_START << endl;

        vd = GetGlobalThetas();
        if (vd.size() > 0)
        {
            sumout << tabs << "\t" << xmlsum::THETAS_START << " ";
            SumFileHandler::WriteVec1D( sumout, vd );
            sumout << xmlsum::THETAS_END << endl;
        }

        vd = GetMigRates();
        if (vd.size() > 0)
        {
            string startstring(xmlsum::MIGRATES_START),
                   endstring(xmlsum::MIGRATES_END);
            if (GetEpochTimes().size() > 0) {
               startstring = xmlsum::DIVMIGRATES_START;
               endstring = xmlsum::DIVMIGRATES_END;
            }
            sumout << tabs << "\t" << startstring << " ";
            SumFileHandler::WriteVec1D( sumout, vd );
            sumout << endstring << endl;
        }

        vd = GetRecRates();
        if (vd.size() > 0)
        {
            sumout << tabs << "\t" << xmlsum::RECRATES_START << " ";
            SumFileHandler::WriteVec1D( sumout, vd );
            sumout << xmlsum::RECRATES_END << endl;
        }
        vd = GetGrowthRates();
        if (vd.size() > 0)
        {
            sumout << tabs << "\t" << xmlsum::GROWTHRATES_START << " ";
            SumFileHandler::WriteVec1D( sumout, vd );
            sumout << xmlsum::GROWTHRATES_END << endl;
        }
        vd = GetLogisticSelectionCoefficient();
        if (vd.size() > 0)
        {
            sumout << tabs << "\t" << xmlsum::LOGISTICSELECTION_START << " ";
            SumFileHandler::WriteVec1D( sumout, vd );
            sumout << xmlsum::LOGISTICSELECTION_END << endl;
        }
        vd = GetDiseaseRates();
        if (vd.size() > 0)
        {
            sumout << tabs << "\t" << xmlsum::DISEASERATES_START << " ";
            SumFileHandler::WriteVec1D( sumout, vd );
            sumout << xmlsum::DISEASERATES_END << endl;
        }
        vd = GetEpochTimes();
        if (vd.size() > 0)
        {
            sumout << tabs << "\t" << xmlsum::EPOCHTIMES_START << " ";
            SumFileHandler::WriteVec1D( sumout, vd );
            sumout << xmlsum::EPOCHTIMES_END << endl;
        }

        if (global_region == m_space)
        {
            const RegionGammaInfo *pRegionGammaInfo = registry.GetRegionGammaInfo();
            if (pRegionGammaInfo)
            {
                vd.clear();
                vd.push_back(pRegionGammaInfo->GetMLE());
                sumout << tabs << "\t" << xmlsum::GAMMAOVERREGIONS_START << " ";
                SumFileHandler::WriteVec1D( sumout, vd );
                sumout << xmlsum::GAMMAOVERREGIONS_END << endl;
            }
        }
        sumout << tabs << xmlsum::ESTIMATES_END << endl;
    }
    else
        SumFileHandler::HandleSumOutFileFormatError("WriteForceParameters");
} // WriteForceParameters

//____________________________________________________________________________________
