// $Id: ui_vars_datapackplus.cpp,v 1.37 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 <set>

#include "datapack.h"
#include "dlmodel.h"
#include "errhandling.h"
#include "region.h"
#include "stringx.h"
#include "ui_vars.h"
#include "ui_vars_datapackplus.h"
#include "ui_strings.h"

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

UIVarsDataPackPlus::UIVarsDataPackPlus(UIVars * myUIVars, DataPack& dp)
    : UIVarsComponent(myUIVars),
      datapack(dp),
      m_effectivePopSizes(),
      m_simulateData()
{
    for (long regnum=0; regnum<GetNumRegions(); regnum++)
    {
        m_effectivePopSizes.push_back(datapack.GetRegion(regnum).GetEffectivePopSize());
        deque<bool> nosims(GetNumLoci(regnum), false);
        m_simulateData.push_back(nosims);
    }
}

UIVarsDataPackPlus::UIVarsDataPackPlus(UIVars * myUIVars, const UIVarsDataPackPlus& dpp)
    : UIVarsComponent(myUIVars),
      datapack(dpp.datapack),
      m_effectivePopSizes(dpp.m_effectivePopSizes),
      m_simulateData(dpp.m_simulateData)
{
}

UIVarsDataPackPlus::~UIVarsDataPackPlus()
{
}

long
UIVarsDataPackPlus::GetMaxLoci() const
{
    long maxL = 0;
    long numRegions = GetNumRegions();
    for(long i=0; i < numRegions; i++)
    {
        long lociCount = GetNumLoci(i);
        if(lociCount > maxL)
        {
            maxL = lociCount;
        }
    }

    return maxL;
}

long
UIVarsDataPackPlus::GetNumRegions() const
{
    return datapack.GetNRegions();
}

long
UIVarsDataPackPlus::GetNumLoci(long regIndex) const
{
    return datapack.GetRegion(regIndex).GetNloci();
}

long
UIVarsDataPackPlus::GetNCrossPartitions() const
{
    return datapack.GetNCrossPartitions();
}

long
UIVarsDataPackPlus::GetNPartitionsByForceType(force_type ft) const
{
    return datapack.GetNPartitionsByForceType(ft);
}

double UIVarsDataPackPlus::GetEffectivePopSize(long region) const
{
    return m_effectivePopSizes[region];
}

long UIVarsDataPackPlus::GetNumIndividualsWithMultipleHaplotypes(long region, string lname) const
{
    long nMultiInd=0;
    const IndVec individuals = datapack.GetRegion(region).GetIndividuals();
    for (unsigned long ind=0; ind<individuals.size(); ind++)
    {
        if (individuals[ind].GetHaplotypesFor(lname, 0).MultipleHaplotypes())
        {
            nMultiInd++;
        }
    }
    return nMultiInd;
}

bool UIVarsDataPackPlus::GetSimulateData(long region, long locus) const
{
    if (region < 0 || locus < 0)
    {
        assert(false); //Don't ask about this! (probably FLAGLONG or the like)
        return false;
    }
    return m_simulateData[region][locus];
}

// LS NOTE:  This should move into the converter eventually.  Probably.
void UIVarsDataPackPlus::SetEffectivePopSize(long region, double size)
{
    if (size <= 0)
    {
        throw data_error("The effective population size must be positive.");
    }
    assert(region < GetNumRegions());
    if (region >= static_cast<long>(m_effectivePopSizes.size()))
    {
        throw implementation_error("Unable to set this region's size--vector not long enough.");
    }
    m_effectivePopSizes[region] = size;

    // may need to update start values at this point
    GetUIVars().forces.FillCalculatedStartValues();
}

void UIVarsDataPackPlus::SetSimulateData(long region, long locus, bool sim)
{
    m_simulateData[region][locus] = sim;
}

void UIVarsDataPackPlus::RevertAllMovingLoci()
{
    for (long reg=0; reg<GetNumRegions(); reg++)
    {
        datapack.GetRegion(reg).RevertMovingLoci();
    }
}

bool
UIVarsDataPackPlus::CanHapArrange() const
{
    return datapack.CanHapArrange();
}

bool
UIVarsDataPackPlus::HasSomeNucleotideData() const
{
    for (long reg=0; reg<GetNumRegions(); reg++)
    {
        for (long loc=0; loc<GetNumLoci(reg); loc++)
        {
            if ((GetDataType(reg, loc) == dtype_DNA) || (GetDataType(reg, loc) == dtype_SNP))
            {
                return true;
            }
        }
    }
    return false;
}

bool
UIVarsDataPackPlus::HasSomeMsatData() const
{
    for (long reg=0; reg<GetNumRegions(); reg++)
    {
        for (long loc=0; loc<GetNumLoci(reg); loc++)
        {
            if (GetDataType(reg, loc) == dtype_msat)
            {
                return true;
            }
        }
    }
    return false;
}

bool
UIVarsDataPackPlus::HasSomeKAlleleData() const
{
    for (long reg=0; reg<GetNumRegions(); reg++)
    {
        for (long loc=0; loc<GetNumLoci(reg); loc++)
        {
            if (GetDataType(reg, loc) == dtype_kallele)
            {
                return true;
            }
        }
    }
    return false;
}

string
UIVarsDataPackPlus::GetCrossPartitionName(long index) const
{
    return GetCrossPartitionNames()[index];
}

StringVec1d
UIVarsDataPackPlus::GetCrossPartitionNames() const
{
    return datapack.GetAllCrossPartitionNames();
}

string
UIVarsDataPackPlus::GetForcePartitionName(force_type ft, long index) const
{
    StringVec1d names = GetForcePartitionNames(ft);
    return names[index];
}

StringVec1d
UIVarsDataPackPlus::GetForcePartitionNames(force_type ft) const
{
    return datapack.GetAllPartitionNames(ft);
}

string
UIVarsDataPackPlus::GetLocusName(long regionIndex, long locusIndex) const
{
    string name = datapack.GetRegion(regionIndex).GetLocus(locusIndex).GetName();
    string locID;
    //LS DEBUG MAPPING:  When IsMovable is no longer sufficient to tell if
    // something is a trait or not, (i.e. if it's being used to partition data)
    // this will need to be changed.
    if (datapack.GetRegion(regionIndex).GetLocus(locusIndex).IsMovable())
    {
        locID = ", trait ";
    }
    else
    {
        if (datapack.GetRegion(regionIndex).GetNloci() == 1)
        {
            //We assume that if there is only one locus in a region, the thing
            // that has better name recognition is the region, not the locus.
            locID = ": ";
            name = datapack.GetRegion(regionIndex).GetRegionName();
        }
        else
        {
            locID = ", segment " + indexToKey(locusIndex) + " : ";
        }
    }

    string fullname = "Region "+indexToKey(regionIndex) + locID + name;
    return fullname;
}

long
UIVarsDataPackPlus::GetLocusIndex(long regionIndex, string locusName) const
{
    return datapack.GetRegion(regionIndex).GetLocusIndex(locusName);
}

bool
UIVarsDataPackPlus::HasLocus(long regionIndex, string locusName) const
{
    return datapack.GetRegion(regionIndex).HasLocus(locusName);
}

string
UIVarsDataPackPlus::GetRegionName(long index) const
{
    string name = datapack.GetRegion(index).GetRegionName();
    string fullname = "Region "+indexToKey(index)+" : "+name;
    return fullname;
}

string
UIVarsDataPackPlus::GetSimpleRegionName(long index) const
{
    return datapack.GetRegion(index).GetRegionName();
}

string
UIVarsDataPackPlus::GetParamName(force_type ftype, long index,bool doLongName) const
{
    string pname = "";
    long fromIndex;
    long toIndex;
    long nstates;
    switch(ftype)
    {
        case force_COAL:
            if(doLongName)
            {
                pname += uistr::theta + " for ";
                pname += GetCrossPartitionName(index);
            }
            else
            {
                pname += "Theta";
                pname += indexToKey(index);
            }
            break;
        case force_DIVERGENCE:
            if(doLongName)
            {
                pname += uistr::divergenceEpochBoundaryTime + " for ";
                pname += ToString(index);
            }
            else
            {
                pname += "Epoch " + ToString(index);
            }
            break;
        case force_DISEASE:
            nstates = GetNPartitionsByForceType(ftype);
            toIndex = index / nstates;
            fromIndex = index % nstates;
            if(doLongName)
            {
                pname += uistr::muRate + " from ";
                pname += GetForcePartitionName(ftype,fromIndex);
                pname += " to ";
                pname += GetForcePartitionName(ftype,toIndex);
            }
            else
            {
                pname += "D";
                pname += indexToKey(fromIndex);
                pname += indexToKey(toIndex);
            }
            break;
        case force_EXPGROWSTICK:
        case force_GROW:
            if(doLongName)
            {
                pname += uistr::growthRate + " for ";
                pname += GetCrossPartitionName(index);
            }
            else
            {
                pname += "Growth";
                pname += indexToKey(index);
            }
            break;
        case force_MIG:
        case force_DIVMIG:
            nstates = GetNPartitionsByForceType(ftype);
            toIndex = index / nstates;
            fromIndex = index % nstates;
            if(doLongName)
            {
                pname += uistr::migrationByID;
                pname += GetForcePartitionName(ftype,toIndex);
                pname += uistr::migrationByID2;
                pname += GetForcePartitionName(ftype,fromIndex);
            }
            else
            {
                pname += "M";
                pname += indexToKey(fromIndex);
                pname += indexToKey(toIndex);
            }
            break;
        case force_REC:
            if(doLongName)
            {
                pname += uistr::recRate;
            }
            else
            {
                pname += "Rec";
            }
            break;
        case force_REGION_GAMMA:
            if(doLongName)
            {
                pname += uistr::regGammaShape;
            }
            else
            {
                pname += "ShapeParam";
            }
            break;
        case force_LOGSELECTSTICK:
        case force_LOGISTICSELECTION:
            if(doLongName)
            {
                pname += uistr::logisticSelectionCoefficient;
            }
            else
            {
                pname += "Sel. Coeff.";
            }
            break;
        case force_NONE:
            assert(false);
            pname += "No force";
            break;
    }
    return pname;
}

string
UIVarsDataPackPlus::GetParamNameOfForce(force_type ftype) const
{
    switch(ftype)
    {
        case force_COAL:
            return "Theta";
            break;
        case force_DISEASE:
            return "Disease Mutation Rate";
            break;
        case force_GROW:
            return "Growth Rate";
            break;
        case force_DIVERGENCE:
            return "Epoch Boundary Time";
            break;
        case force_DIVMIG:
        case force_MIG:
            return "Migration Rate";
            break;
        case force_REC:
            return "Recombination Rate";
            break;
        case force_REGION_GAMMA:
            return "Gamma Over Regions";
            break;
        case force_EXPGROWSTICK:
            return "Exponential Growth via stick";
            break;
        case force_LOGISTICSELECTION:
            return "Logistic Selection Coefficient";
            break;
        case force_LOGSELECTSTICK:
            return "Logistic Selection via stick";
            break;
        case force_NONE:
            assert(false);
            throw implementation_error("UIVarsDataPackPlus::GetParamNameOfForce for force_NONE!");
            break;
    };
    assert(false);
    throw implementation_error("UIVarsDataPackPlus::GetParamNameOfForce missing switch case");
}

data_type
UIVarsDataPackPlus::GetDataType(long regionId, long locusId) const
{
    const Region & thisRegion = datapack.GetRegion(regionId);
    const Locus & thisLocus = thisRegion.GetLocus(locusId);
    return thisLocus.GetDataType();
}

bool UIVarsDataPackPlus::IsMovable(long regionId, long locusId) const
{
    const Region & thisRegion = datapack.GetRegion(regionId);
    const Locus & thisLocus = thisRegion.GetLocus(locusId);
    return thisLocus.IsMovable();
}

string UIVarsDataPackPlus::GetName(long regionId, long locusId) const
{
    const Region & thisRegion = datapack.GetRegion(regionId);
    const Locus & thisLocus = thisRegion.GetLocus(locusId);
    return thisLocus.GetName();
}

rangeset UIVarsDataPackPlus::GetRange(long regionId, long locusId) const
{
    const Region & thisRegion = datapack.GetRegion(regionId);
    const Locus & thisLocus = thisRegion.GetLocus(locusId);
    return thisLocus.GetAllowedRange();
}

rangepair UIVarsDataPackPlus::GetRegionSiteSpan(long regionId) const
{
    return datapack.GetRegion(regionId).GetSiteSpan();
}

long UIVarsDataPackPlus::GetNumSites(long regionId) const
{
    const Region & thisRegion = datapack.GetRegion(regionId);
    return thisRegion.GetNumSites();
}

long UIVarsDataPackPlus::GetNumSites(long regionId, long locusId) const
{
    const Region & thisRegion = datapack.GetRegion(regionId);
    const Locus & thisLocus = thisRegion.GetLocus(locusId);
    return thisLocus.GetNsites();
}

StringVec2d UIVarsDataPackPlus::GetUniqueAlleles(long regionId,
                                                 long locusId) const
{
    StringVec2d retVec;
    const Region& thisRegion = datapack.GetRegion(regionId);
    const Locus& thisLocus = thisRegion.GetLocus(locusId);
    vector<TipData> tips = thisLocus.GetTipData();
    const IndVec individuals = thisRegion.GetIndividuals();
    for (long marker = 0; marker < thisLocus.GetNmarkers(); marker++)
    {
        std::set<string> uniqueAlleles;
        //First, get any alleles from the unknown haplotypes
        for (unsigned long ind=0; ind<individuals.size(); ind++)
        {
            StringVec1d alleles = individuals[ind].GetAllelesFor(thisLocus.GetName(), marker);
            for (unsigned long nallele=0; nallele<alleles.size(); nallele++)
            {
                if (alleles[nallele] != "?")
                {
                    uniqueAlleles.insert(alleles[nallele]);
                }
            }
        }
        //Now get any alleles from the tips
        for (unsigned long tip = 0; tip < tips.size(); ++tip)
        {
            if (!tips[tip].m_nodata)
            {
                string allele = tips[tip].data[marker];
                if (allele != "?")
                {
                    uniqueAlleles.insert(allele);
                }
            }
        }
        //Convert the set to a vector (for simplicity for our clients)
        StringVec1d alleleStrings;
        for (std::set<string>::iterator unique=uniqueAlleles.begin(); unique != uniqueAlleles.end(); unique++)
        {
            alleleStrings.push_back(*unique);
        }
        retVec.push_back(alleleStrings);
    }
    return retVec;

}

std::set<long> UIVarsDataPackPlus::GetPloidies(long region) const
{
    return datapack.GetRegion(region).GetPloidies();
}

const Locus* UIVarsDataPackPlus::GetConstLocusPointer(long region, long locus) const
{
    return &datapack.GetRegion(region).GetLocus(locus);
}

bool UIVarsDataPackPlus::AnySimulation() const
{
    for (unsigned long region=0; region<m_simulateData.size(); region++)
    {
        for (unsigned long locus=0; locus<m_simulateData[region].size(); locus++)
        {
            if (m_simulateData[region][locus]) return true;
        }
    }
    return false;
}

bool UIVarsDataPackPlus::AnyRelativeHaplotypes() const
{
    for (long region=0; region<GetNumRegions(); region++)
    {
        const Region& reg = datapack.GetRegion(region);
        for (long ind=0; ind<reg.GetNIndividuals(); ind++)
        {
            const Individual& indiv = reg.GetIndividual(ind);
            if (indiv.MultipleTraitHaplotypes()) return true;
        }
    }
    return false;
}

bool UIVarsDataPackPlus::AnySNPDataWithDefaultLocations() const
{
    for (long region=0; region<GetNumRegions(); region++)
    {
        if (datapack.GetRegion(region).AnySNPDataWithDefaultLocations())
        {
            return true;
        }
    }
    return false;
}

string UIVarsDataPackPlus::GetAncestorName(long index) const
{
    //const Region & thisRegion = datapack.GetRegion(regionId);
    //const Locus & thisLocus = thisRegion.GetLocus(locusId);
    //uiInterface.GetCurrentVars().forces.AddAncestor(ancname);

    //return thisLocus.GetName();
    //string pname = datapack.forces.UIVarsDivergenceForceGetParent(index).GetName();
    //datapack.GetRegion(index).GetRegionName();
    string pname = uistr::divergenceEpochAncestor;
    pname += ToString(index + 1);

    return pname;
}

string UIVarsDataPackPlus::GetDescendentNames(long index) const
{
    //const Region & thisRegion = datapack.GetRegion(regionId);
    //const Locus & thisLocus = thisRegion.GetLocus(locusId);
    //uiInterface.GetCurrentVars().forces.AddAncestor(ancname);

    //return thisLocus.GetName();
    //string pname = datapack.forces.GetAncestor(index);
    string pname = uistr::divergenceEpochDescendent;
    pname += ToString(index + 1);
    pname += "-1 ";
    pname += uistr::divergenceEpochDescendent;
    pname += ToString(index + 1);
    pname += "-2";

    return pname;
}

//____________________________________________________________________________________
