// $Id: ui_vars_traitmodels.cpp,v 1.19 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 "locus.h"
#include "ui_regid.h"
#include "ui_vars.h"
#include "ui_vars_traitmodels.h"

#define DEFAULTMLOC mloc_mapfloat
#define MAXMULTIHAP 20

using namespace std;

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

string ToString(mloc_type type)
{
    switch(type)
    {
        case mloc_data:
            return "use as data";
        case mloc_mapjump:
            return "mapping (jump)";
        case mloc_mapfloat:
            return "mapping (float)";
        case mloc_partition:
            return "partition the data";
    }
    assert(false);
    throw implementation_error("Uncaught moving locus analysis type.");
}

string ToXMLString(mloc_type type)
{
    switch(type)
    {
        case mloc_data:
            return "data";
        case mloc_mapjump:
            return "jump";
        case mloc_mapfloat:
            return "float";
        case mloc_partition:
            return "partition";
    }
    assert(false);
    throw implementation_error("Uncaught moving locus analysis type.");
}

mloc_type ProduceMlocTypeOrBarf(const string& in)
{
    string st = in;
    LowerCase(st);
    if (st == "data") { return mloc_data; };
    if (st == "jump") { return mloc_mapjump; };
    if (st == "float") { return mloc_mapfloat; };
    if (st == "partition") { return mloc_partition; };
    throw data_error("Illegal trait analysis setting \""+in+"\"");
}

UIVarsSingleTraitModel::UIVarsSingleTraitModel(UIRegId regionId, string name,
                                               rangeset mrange,rangepair fullr,
                                               const Locus* locus,
                                               long multihapnum)
    : m_region(regionId.GetRegion()),
      m_locus(regionId.GetLocus()),
      m_datatype(regionId.GetDataType()),
      m_type(DEFAULTMLOC),
      m_range(mrange),
      m_fullrange(fullr),
      m_name(name),
      m_phenotypes(regionId.GetRegion(), locus->GetName()),
      m_multihapnum(multihapnum)
{
    assert(m_datatype == dtype_kallele); //We don't have models for other types.
    assert(m_range.size() > 0);
    if (m_range.size() == 1)
    {
        if (m_range.begin()->first +1 == m_range.begin()->second)
        {
            m_type = mloc_data;
        }
    }
    if ((m_type != mloc_data) && (m_multihapnum > MAXMULTIHAP))
    {
        m_type = mloc_mapjump;
    }
}

UIVarsSingleTraitModel::~UIVarsSingleTraitModel()
{
}

void UIVarsSingleTraitModel::SetAnalysisType(mloc_type type)
{
    m_type = type;
}

void UIVarsSingleTraitModel::SetRange(rangeset range)
{
    m_range = range;
}

void UIVarsSingleTraitModel::AddPhenotype(StringVec1d& alleles, string name, double penetrance)
{
    m_phenotypes.AddPhenotype(alleles, name, penetrance);
}

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

UIVarsTraitModels::UIVarsTraitModels(UIVars* myUIVars)
    : UIVarsComponent(myUIVars),
      m_individualModels()
{
    //Requires that datapackplus be set up in the uivars.
    long nregions = GetConstUIVars().datapackplus.GetNumRegions();
    for(long regionIndex=0; regionIndex < nregions; regionIndex++)
    {
        long numLoci = GetConstUIVars().datapackplus.GetNumLoci(regionIndex);
        for(long locusIndex=0; locusIndex < numLoci; locusIndex++)
        {
            if (GetConstUIVars().datapackplus.IsMovable(regionIndex, locusIndex))
            {
                UIRegId regID(regionIndex, locusIndex, GetConstUIVars());
                AddTraitModel(regID);
            }
        }
    }
}

UIVarsTraitModels::UIVarsTraitModels(UIVars* myUIVars, const UIVarsTraitModels& clone)
    : UIVarsComponent(myUIVars),
      m_individualModels(clone.m_individualModels)
{
}

UIVarsTraitModels::~UIVarsTraitModels()
{
}

void UIVarsTraitModels::AddTraitModel(UIRegId regID)
{
    assert(m_individualModels.find(regID) == m_individualModels.end());
    string name = GetConstUIVars().datapackplus.GetName(regID.GetRegion(), regID.GetLocus());
    rangeset mrange = GetConstUIVars().datapackplus.GetRange(regID.GetRegion(), regID.GetLocus());
    rangepair fullrange = GetConstUIVars().datapackplus.GetRegionSiteSpan(regID.GetRegion());
    long multihapnum = GetConstUIVars().datapackplus.GetNumIndividualsWithMultipleHaplotypes(regID.GetRegion(), name);
    if (mrange.size() == 0)
    {
        //Nothing originally set, so we'll allow the whole thing.
        mrange.insert(fullrange);
    }
    const Locus* locus = GetConstUIVars().datapackplus.GetConstLocusPointer(regID.GetRegion(), regID.GetLocus());
    fullrange.second++; // EWFIX. WHY WHY WHY
    UIVarsSingleTraitModel tmodel(regID, name, mrange, fullrange, locus, multihapnum);
    m_individualModels.insert(make_pair(regID, tmodel));
}

//Setters

void UIVarsTraitModels::SetAnalysisType(UIRegId regID, mloc_type type)
{
    assert(m_individualModels.find(regID) != m_individualModels.end());
    UIVarsSingleTraitModel& model = m_individualModels.find(regID)->second;
    switch(type)
    {
        case mloc_data:
            //We need to make sure there's only one site where it can be placed.
            if (!OneSite(model.GetRange()))
            {
                throw data_error("You may not set the trait analysis type to 'data' without specifying the exact position of your trait.");
            }
            break;
        case mloc_partition:
            throw data_error("We are not set up to use trait data to set partitions yet.  Soon!");
        case mloc_mapfloat:
            if (model.GetMultiHapNum() > MAXMULTIHAP)
            {
                throw data_error("You have too many individuals (" + ToString(model.GetMultiHapNum()) + ", which is more than maximum " + ToString(MAXMULTIHAP) + " allowed) with multiple haplotype resolutions (such as heterozygotes or individuals with a dominant phenotype) to perform a 'floating' analysis--the number of complete likelihood calculations required per tree is 2^N, with N the number of individuals with multiple resolutions.  We recommend either a 'jumping' analysis, or that you remove (at random) individuals from your data set.");
            }
            //Otherwise, fall through to:
        case mloc_mapjump:
            if (OneSite(model.GetRange()))
            {
                GetConstUIVars().GetUI()->AddWarning("You currently have only a single site where your trait is allowed.  The mapping analysis will not have anything to compare it to.  Use the 'A' option to add more sites to the analysis.");
            }
            break;
    }
    model.SetAnalysisType(type);
}

void UIVarsTraitModels::AddRange(UIRegId regID, rangepair addpart)
{
    assert(m_individualModels.find(regID) != m_individualModels.end());
    rangeset rset = m_individualModels.find(regID)->second.GetRange();
    rset = AddPairToRange(addpart, rset);
    SetRange(regID, rset);
}

void UIVarsTraitModels::RemoveRange(UIRegId regID, rangepair removepart)
{
    assert(m_individualModels.find(regID) != m_individualModels.end());
    rangeset rset = m_individualModels.find(regID)->second.GetRange();
    rset = RemovePairFromRange(removepart, rset);
    SetRange(regID, rset);
}

void UIVarsTraitModels::SetRangeToPoint(UIRegId regID, long site)
{
    assert(m_individualModels.find(regID) != m_individualModels.end());
    rangepair rp(make_pair(site, site+1));
    rangeset rset;
    rset.insert(rp);
    SetRange(regID, rset);
}

void UIVarsTraitModels::AddPhenotype(UIRegId regID, StringVec1d& alleles, string name, double penetrance)
{
    assert(m_individualModels.find(regID) != m_individualModels.end());
    m_individualModels.find(regID)->second.AddPhenotype(alleles, name, penetrance);
}

//Getters:
long UIVarsTraitModels::GetNumMovableLoci() const
{
    return m_individualModels.size();
}

mloc_type UIVarsTraitModels::GetAnalysisType(UIRegId regID) const
{
    assert(m_individualModels.find(regID) != m_individualModels.end());
    return m_individualModels.find(regID)->second.GetAnalysisType();
}

rangeset UIVarsTraitModels::GetRange(UIRegId regID) const
{
    assert(m_individualModels.find(regID) != m_individualModels.end());
    return m_individualModels.find(regID)->second.GetRange();
}

long UIVarsTraitModels::GetInitialMapPosition(UIRegId regID) const
{
    //Just use the leftmost valid point.
    return GetRange(regID).begin()->first;
}

string UIVarsTraitModels::GetName(UIRegId regID) const
{
    assert(m_individualModels.find(regID) != m_individualModels.end());
    return m_individualModels.find(regID)->second.GetName();
}

Phenotypes UIVarsTraitModels::GetPhenotypes(UIRegId regID) const
{
    assert(m_individualModels.find(regID) != m_individualModels.end());
    return m_individualModels.find(regID)->second.GetPhenotypes();
}

vector<UIRegId> UIVarsTraitModels::GetRegIDs() const
{
    vector<UIRegId> retvec;
    for (map<UIRegId, UIVarsSingleTraitModel>::const_iterator model = m_individualModels.begin(); model != m_individualModels.end(); model++)
    {
        retvec.push_back(model->first);
    }
    return retvec;
}

bool UIVarsTraitModels::AnyJumpingAnalyses() const
{
    for (map<UIRegId, UIVarsSingleTraitModel>::const_iterator model = m_individualModels.begin(); model != m_individualModels.end(); model++)
    {
        if (model->second.GetAnalysisType() == mloc_mapjump) return true;
    }
    return false;
}

bool UIVarsTraitModels::AnyMappingAnalyses() const
{
    for (map<UIRegId, UIVarsSingleTraitModel>::const_iterator model = m_individualModels.begin(); model != m_individualModels.end(); model++)
    {
        if (model->second.GetAnalysisType() == mloc_mapjump) return true;
        if (model->second.GetAnalysisType() == mloc_mapfloat) return true;
    }
    return false;
}

void UIVarsTraitModels::SetRange(UIRegId regID, rangeset rset)
{
    UIVarsSingleTraitModel& model = m_individualModels.find(regID)->second;

    rangepair fullrange = model.GetFullRange();
    if (rset.size() == 0)
    {
        throw data_error("You must leave at least one allowable site for this trait.  The sites must be chosen from the range " + ToStringUserUnits(fullrange) + ".");
    }

    rangepair below = make_pair(rset.begin()->first-1, fullrange.first);
    rangepair above = make_pair(fullrange.second, rset.rbegin()->second);
    rangeset truerange = rset;

    if (below.first < below.second)
    {
        truerange = RemovePairFromRange(below, truerange);
    }
    if (above.first < above.second)
    {
        truerange = RemovePairFromRange(above, truerange);
    }

    if (truerange.size() == 0)
    {
        throw data_error("You must leave at least one allowable site for this trait.  The sites must be chosen from the range "
                         + ToStringUserUnits(fullrange)
                         + ".");
    }
    if (truerange != rset)
    {
        GetConstUIVars().GetUI()->AddWarning("One or more of the added sites for this trait are outside of the range of the segments in this region.  Truncating the added range.");
    }

    switch (model.GetAnalysisType())
    {
        case mloc_data:
        case mloc_partition:
            if (!OneSite(truerange))
            {
                model.SetAnalysisType(mloc_mapfloat);
                GetConstUIVars().GetUI()->AddWarning("If you don't know where your trait is located, you may perform a mapping analysis,"
                                                     " but you may not use the trait as data or use it to partition your data.  "
                                                     "Changing the analysis type to map the trait after collecting trees ('float').");
            }
            break;
        case mloc_mapjump:
        case mloc_mapfloat:
            if (OneSite(truerange))
            {
                //LS DEBUG MAPPING: change after implementation
#ifdef NDEBUG
                GetConstUIVars().GetUI()->AddWarning("You currently have only a single site where your trait is allowed.  "
                                                     "The mapping analysis will not have anything to compare it to.  "
                                                     "Use the 'A' option to add more sites.");
#else
                GetConstUIVars().GetUI()->AddWarning("You currently have only a single site where your trait is allowed.  "
                                                     "The mapping analysis will not have anything to compare it to.  "
                                                     "Use the 'A' option to add more sites, or the 'U' or 'P' options to change the type of analysis.");
#endif
            }
    }
    model.SetRange(truerange);
}

//Private functions

bool UIVarsTraitModels::OneSite(rangeset rset)
{
    assert(rset.size() > 0);
    if (rset.size() > 1) return false;
    if (rset.begin()->first+1 == rset.begin()->second) return true;
    return false;
}

//____________________________________________________________________________________
