// $Id: individual.cpp,v 1.31 2018/01/03 21:33:03 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 "branch.h"
#include "errhandling.h"
#include "individual.h"
#include "random.h"
#include "registry.h"
#include "stringx.h"
#include "xml_strings.h"
#include "dlmodel.h"  // for GetAllelesFromDLs()

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

using namespace std;

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

StringVec1d Individual::GetAllTipNames() const
{
    StringVec1d names;
    vector<Branch_ptr>::const_iterator tip;

    for(tip = m_ptips.begin(); tip != m_ptips.end(); ++tip)
    {
        assert(boost::dynamic_pointer_cast<TBranch>(*tip));
        names.push_back(boost::dynamic_pointer_cast<TBranch>(*tip)->m_label);
    }

    return names;

} // GetAllTipNames

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

bool Individual::AnyPhaseUnknownSites() const
{
    LongVec2d::const_iterator locus;
    for(locus = m_phasemarkers.begin(); locus < m_phasemarkers.end(); ++locus)
        if (!(locus->empty())) return true;

    return false;

} // AnyPhaseUnknownSites

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

bool Individual::MultipleTraitHaplotypes() const
{
    map<pair<string, long>, Haplotypes >::const_iterator hap;
    for (hap = m_haplotypesmap.begin(); hap != m_haplotypesmap.end(); hap++)
    {
        if (hap->second.MultipleHaplotypes()) return true;
    }
    return false;
}

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

void Individual::PruneSamePhaseUnknownSites()
{
    // remove homozygous sites from list of phase-unknown sites
    unsigned long locus;
    LongVec2d newphases;

    for(locus = 0; locus < m_phasemarkers.size(); ++locus)
    {
        LongVec1d newphase;
        const LongVec1d& locusphase = m_phasemarkers[locus];
        if (locusphase.empty()) continue;
        LongVec1d::const_iterator marker;
        for (marker = locusphase.begin(); marker != locusphase.end(); ++marker)
        {
            Branch_ptr firsttip = m_ptips.front();
            vector<Branch_ptr>::const_iterator tip = m_ptips.begin();
            for(++tip; tip != m_ptips.end(); ++tip)
            {
                if (firsttip->DiffersInDLFrom(*tip, locus, *marker))
                {
                    newphase.push_back(*marker);
                    break;
                }
            }
        }
        newphases.push_back(newphase);
    }

    m_phasemarkers = newphases;

} // Individual::PruneSamePhaseUnknownSites

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

pair<long,long> Individual::PickRandomPhaseMarker(Random& rs) const
{
    long nmarkers(0L);
    LongVec2d::const_iterator locus;
    for(locus = m_phasemarkers.begin(); locus < m_phasemarkers.end(); ++locus)
    {
        nmarkers += locus->size();
    }

    long picked = rs.Long(nmarkers);
    unsigned long pickedlocus;
    for(pickedlocus = 0; pickedlocus < m_phasemarkers.size(); ++pickedlocus)
    {
        picked -= m_phasemarkers[pickedlocus].size();
        if (picked < 0)
        {
            picked += m_phasemarkers[pickedlocus].size();
            break;
        }
    }

    return make_pair(pickedlocus, m_phasemarkers[pickedlocus][picked]);

} // Individual::PickRandomPhaseMarker

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

pair<string,long> Individual::PickRandomHaplotypeMarker() const
{
    assert(m_haplotypesmap.size() > 0); //ChooseAllHaplotypes needs to be run.
    unsigned long whichhap = registry.GetRandom().Long(m_haplotypesmap.size());
    map<pair<string, long>, Haplotypes >::const_iterator hapmap=m_haplotypesmap.begin();
    for (unsigned long hapnum=0;
         hapnum<m_haplotypesmap.size() && hapmap != m_haplotypesmap.end();
         hapnum++, hapmap++)
    {
        if (hapnum == whichhap)
        {
            return hapmap->first;
        }
    }
    assert(false); //should have found one.
    return m_haplotypesmap.begin()->first;

} // Individual::PickRandomHaplotypeMarker

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

void Individual::SetPhaseMarkers(const LongVec2d& pm)
{
    m_phasemarkers = pm;
}

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

void Individual::SetPhaseSites(const LongVec2d& ps)
{
    m_phasesites = ps;
}

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

void Individual::AddHaplotype(long regnum, string lname, long marker,
                              const StringVec1d& alleles, double penetrance)
{
    map<pair<string, long>, Haplotypes>::iterator hap;
    hap = m_haplotypesmap.find(make_pair(lname, marker));
    if (hap == m_haplotypesmap.end())
    {
        Haplotypes haplo(regnum, lname);
        m_haplotypesmap.insert(make_pair(make_pair(lname, marker),haplo));
        hap = m_haplotypesmap.find(make_pair(lname, marker));
    }
    hap->second.AddHaplotype(alleles, penetrance);
}

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

StringVec1d Individual::GetAllelesFor(string lname, long marker) const
{
    map<pair<string, long>, Haplotypes>::const_iterator hap;
    hap = m_haplotypesmap.find(make_pair(lname, marker));

    if (hap == m_haplotypesmap.end())
    {
        StringVec1d emptyvec;
        return emptyvec;
    }

    return hap->second.GetAlleles();
}

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

vector<LocusCell> Individual::GetLocusCellsFor(string lname, long marker) const
{
    map<pair<string, long>, vector<LocusCell> >::const_iterator hap;
    hap = m_currentHapsMap.find(make_pair(lname, marker));
    if (hap == m_currentHapsMap.end())
    {
        assert(false); //Sholdn't ask for a haplotype we haven't chosen.
        // This is probably another phase 1/phase 2 issue
        string msg = "Can't find haplotypes for locus " + lname
            + ", marker " + ToString(marker) + ", individual " + m_name
            + ".";
        throw implementation_error(msg);
    }
    return hap->second;
}

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

Haplotypes Individual::GetHaplotypesFor(string lname, long marker) const
{
    map<pair<string, long>, Haplotypes >::const_iterator hap;
    hap = m_haplotypesmap.find(make_pair(lname, marker));
    if (hap == m_haplotypesmap.end())
    {
        assert(false); //Sholdn't ask for a haplotype we haven't chosen.
        // This is probably another phase 1/phase 2 issue
        string msg = "Can't find haplotypes for segment " + lname
            + ", marker " + ToString(marker) + ", individual " + m_name
            + ".";
        throw implementation_error(msg);
    }
    return hap->second;
}

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

string Individual::GetMarkerDataFor(string lname, long marker) const
{
    map<pair<string, long>, Haplotypes>::const_iterator hap;
    hap = m_haplotypesmap.find(make_pair(lname, marker));
    if (hap == m_haplotypesmap.end())
    {
        return "";
    }
    return hap->second.GetMarkerData();
}

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

void Individual::ChooseNewHaplotypesFor(string lname, long marker)
{
    map<pair<string, long>, Haplotypes>::iterator hap;
    pair<string, long> tag = make_pair(lname, marker);
    hap = m_haplotypesmap.find(tag);
    if (hap == m_haplotypesmap.end())
    {
        assert(false);                  // We shouldn't ask about haplotypes for which we don't have data.
        throw implementation_error("Can't choose a haplotype for locus" +
                                   lname + ".");
        //If we want to handle this anyway, the proper thing to do would
        // be to make some DLCells with the equivalent of ?s in them.
    }
    vector<LocusCell> cells = hap->second.ChooseNewHaplotypes();
    map<pair<string, long>, vector<LocusCell> >::iterator oldhap;
    oldhap = m_currentHapsMap.find(tag);
    bool newhap = true;
    if (oldhap != m_currentHapsMap.end())
    {
        if (oldhap->second == cells)
        {
            newhap = false;
        }
    }
    m_currentHapsMap.erase(tag);
    m_currentHapsMap.insert(make_pair(tag, cells));
    assert(newhap); //Should always pick new site DLs.
}

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

bool Individual::ChooseRandomHaplotypesFor(string lname, long marker)
{
    map<pair<string, long>, Haplotypes>::iterator hap;
    pair<string, long> tag = make_pair(lname, marker);
    hap = m_haplotypesmap.find(tag);
    if (hap == m_haplotypesmap.end())
    {
        assert(false);                  // We shouldn't ask about haplotypes for which we don't have data.
        throw implementation_error("Can't choose a haplotype for locus" + lname + ".");
        //If we want to handle this anyway, the proper thing to do would
        // be to make some DLCells with the equivalent of ?s in them.
    }
    vector<LocusCell> cells = hap->second.ChooseRandomHaplotypes();
    map<pair<string, long>, vector<LocusCell> >::iterator oldhap;
    oldhap = m_currentHapsMap.find(tag);
    bool newhap = true;
    if (oldhap != m_currentHapsMap.end())
    {
        if (oldhap->second == cells)
        {
            newhap = false;
        }
    }
    m_currentHapsMap.erase(tag);
    m_currentHapsMap.insert(make_pair(tag, cells));
    return newhap;
}

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

void Individual::RandomizeAllHaplotypes()
{
    map<pair<string, long>, Haplotypes>::iterator hap;
    for (hap = m_haplotypesmap.begin(); hap != m_haplotypesmap.end(); hap++)
    {
        ChooseRandomHaplotypesFor(hap->first.first, hap->first.second);
    }
}

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

void Individual::ChooseFirstHaplotypeFor(string lname, long marker)
{
    map<pair<string, long>, Haplotypes>::iterator hap;
    pair<string, long> tag = make_pair(lname, marker);
    hap = m_haplotypesmap.find(tag);
    if (hap == m_haplotypesmap.end())
    {
        assert(false);                  // We shouldn't ask about haplotypes for which we don't have data.
        string msg("Can't choose a haplotype for segment" + lname + ".");
        throw implementation_error(msg);
    }
    vector<LocusCell> cells = hap->second.ChooseFirstHaplotypes();
    m_currentHapsMap.erase(tag);
    m_currentHapsMap.insert(make_pair(tag, cells));
}

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

bool Individual::ChooseNextHaplotypeFor(string lname, long marker)
{
    map<pair<string, long>, Haplotypes>::iterator hap;
    pair<string, long> tag = make_pair(lname, marker);
    hap = m_haplotypesmap.find(tag);
    if (hap == m_haplotypesmap.end())
    {
        assert(false);                  // We shouldn't ask about haplotypes for which we don't have data.
        string msg("Can't choose a haplotype for segment" + lname + ".");
        throw implementation_error(msg);
    }
    vector<LocusCell> cells = hap->second.ChooseNextHaplotypes();
    if (cells.size() == 0)
    {
        //returning an empty vector means that we've gone through all the haplotypes
        return false;
    }
    m_currentHapsMap.erase(tag);
    m_currentHapsMap.insert(make_pair(tag, cells));
    return true;
}

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

void Individual::SetHaplotypes(string lname, long marker, Haplotypes haps)
{
    map<pair<string, long>, Haplotypes>::iterator hap;
    pair<string, long> whichlocus = make_pair(lname, marker);
    hap = m_haplotypesmap.find(whichlocus);
    if (hap != m_haplotypesmap.end())
    {
        m_haplotypesmap.erase(hap);
    }
    m_haplotypesmap.insert(make_pair(whichlocus, haps));
    //LS DEBUG:  Probably change this back to ChooseRandom--was trying
    //to not use the random number generator here.
    ChooseFirstHaplotypeFor(lname, marker);
}

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

StringVec1d Individual::GetAllelesFromDLs(long locus, long marker, bool moving,
                                          DataModel_ptr model)
{
    StringVec1d alleles;
    for (size_t tip=0; tip<m_ptips.size(); tip++)
    {
        Cell_ptr cell = m_ptips[tip]->GetDLCell(locus, marker, moving);
        string allele = model->CellToData(cell, marker);
        alleles.push_back(allele);
    }
    return alleles;
}

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

bool Individual::IsValidIndividual() const
{
    // all we can check is that the pointers are non-NULL

    if (m_ptips.empty()) return true;
    unsigned long i;
    for (i = 0; i < m_ptips.size(); ++i)
    {
        if (m_ptips[0] == NULL) return false;  // bad branch!
    }

    return true;
} // IsValidIndividual

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

StringVec1d Individual::GetTraitXML(long nspaces) const
{
    StringVec1d retvec;
    if (m_haplotypesmap.size() == 0) return retvec;

    string spaces(nspaces, ' ');
    nspaces += 2;

    map<pair<string, long>, Haplotypes >::const_iterator hap;
    for (hap = m_haplotypesmap.begin(); hap != m_haplotypesmap.end(); hap++)
    {
        retvec.push_back(spaces + MakeTag(xmlstr::XML_TAG_GENOTYPE_RESOLUTIONS));
        retvec.push_back(spaces + "  " + MakeTag(xmlstr::XML_TAG_TRAIT_NAME) + " "
                         + hap->first.first + " " + MakeCloseTag(xmlstr::XML_TAG_TRAIT_NAME));
        StringVec1d haplotypes = hap->second.GetHaplotypesXML(nspaces);
        retvec.insert(retvec.end(), haplotypes.begin(), haplotypes.end());
        retvec.push_back(spaces + MakeCloseTag(xmlstr::XML_TAG_GENOTYPE_RESOLUTIONS));
    }
    return retvec;
}

//____________________________________________________________________________________
