// $Id: haplotypes.cpp,v 1.14 2018/01/03 21:32:57 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: Lucian Smith, Mary Kuhner, Jon Yamato and Joseph Felsenstein */


#include <cassert>
#include <iostream>                     // debug only

#include "errhandling.h"
#include "haplotypes.h"
#include "locus.h"
#include "locuscell.h"
#include "registry.h"
#include "region.h"
#include "stringx.h"
#include "xml_strings.h"

using std::multiset;
using std::cerr;
using std::endl;

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

Haplotypes::Haplotypes(long regnum, string lname)
    : m_regionnum(regnum),
      m_locusname(lname),
      m_haplotype_alleles(),
      m_penetrances(),
      m_haplotype_dlcells(),
      m_current_hapindex(FLAGLONG)
{
    // intentionally blank
}

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

Haplotypes::Haplotypes(Haplotypes oldhaps, bool clear)
    : m_regionnum(oldhaps.m_regionnum),
      m_locusname(oldhaps.m_locusname),
      m_haplotype_alleles(),
      m_penetrances(),
      m_haplotype_dlcells(),
      m_current_hapindex(FLAGLONG)
{
    assert(clear == true);
}

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

void Haplotypes::ConvertAllelesToDLCells()
{
    assert(m_haplotype_alleles.size() > 0);
    assert(m_haplotype_alleles.size() == m_penetrances.size());
    assert(m_haplotype_dlcells.size() == 0);

    for (unsigned long hap=0; hap<m_haplotype_alleles.size(); hap++)
    {
        vector<LocusCell> resolution;
        for (unsigned long cell=0; cell<m_haplotype_alleles[hap].size(); cell++)
        {
            StringVec1d onecell;
            onecell.push_back(m_haplotype_alleles[hap][cell]);
            const Locus& locus=registry.GetDataPack().GetRegion(m_regionnum).GetLocus(m_locusname);
            resolution.push_back(locus.GetDataTypePtr()->CreateInitializedDLCell(locus, onecell));
        }
        //Now we multiply final LocusCell by the appropriate penetrance.  We
        // could choose any of them if we wanted to, but choose the last so that
        // the CollapseHaplotypeDLs() routine doesn't have to worry about
        // penetrances.
        resolution[resolution.size()-1] *= m_penetrances[hap];
        m_haplotype_dlcells.push_back(resolution);
    }
    // Combine haplotype resolutions which share their identity at a tip
    //  by simply adding the DLCells at the other tip.
    CollapseHaplotypeDLs();
}

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

void Haplotypes::CollapseHaplotypeDLs()
{
    vector<vector<LocusCell> >::iterator dlList = m_haplotype_dlcells.begin();
    for (;dlList != m_haplotype_dlcells.end();dlList++)
    {
        vector<vector<LocusCell> >::iterator dlComp = dlList;
        dlComp++;
        for (;dlComp != m_haplotype_dlcells.end();)
        {
            bool allbutlastmatch = true;
            //Actually, the last can match, too, but that should never happen.
            for (size_t allele=0; allele<((*dlList).size()-1); allele++)
            {
                if (!((*dlList)[allele] == (*dlComp)[allele]))
                {
                    allbutlastmatch = false;
                }
            }
            if (allbutlastmatch)
            {
                //We can collapse them into one DLcell
                (*dlList)[(*dlList).size()-1] += (*dlComp)[(*dlComp).size()-1];
                //Now delete the compared one.
                dlComp = m_haplotype_dlcells.erase(dlComp);
            }
            else
            {
                dlComp++;
            }
        }
    }
}

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

void Haplotypes::AddHaplotype(StringVec1d alleles, double penetrance)
{
    if (m_haplotype_alleles.size() > 0)
    {
        if (alleles.size() != m_haplotype_alleles[0].size())
        {
            string msg = "The haplotype resolution \"";
            for (unsigned long i=0; i<alleles.size(); i++)
            {
                msg += alleles[i] + " ";
            }
            msg += "\" has a different number of alleles than the haplotype "
                "resolution \"";
            for (unsigned long i=0; i<m_haplotype_alleles[0].size(); i++)
            {
                msg += m_haplotype_alleles[0][i] + " ";
            }
            msg += "\".  Remember that spaces are not allowed in allele names.  Also, if you have samples"
                " with multiple ploidies (if you have samples from an X chromosome, say) each phenotype must match"
                " a set of genotypes of the same ploidy -- even if the phenotype of 'X0' matches the phenotype"
                " of 'XX', they must be defined separately.";
            throw data_error(msg);
        }
    }
    for (size_t hap=0; hap<m_haplotype_alleles.size(); hap++)
    {
        if (m_haplotype_alleles[hap] == alleles)
        {
            return;
        }
    }
    m_haplotype_alleles.push_back(alleles);
    m_penetrances.push_back(penetrance);
}

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

void Haplotypes::AddHaplotype(multiset<string> alleles, double penetrance)
{
    StringVec2d allAlleles = SetToVecs(alleles);
    for (size_t alleleVec=0; alleleVec != allAlleles.size(); alleleVec++)
    {
        AddHaplotype(allAlleles[alleleVec], penetrance);
    }
}

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

vector<LocusCell> Haplotypes::ChooseNewHaplotypes()
{
    if (m_haplotype_dlcells.size() == 0)
    {
        ConvertAllelesToDLCells();
    }
    if (m_haplotype_dlcells.size() == 1)
    {
        return m_haplotype_dlcells[0];
    }
    //Choose a new haplotype index
    long newindex = m_current_hapindex;
    while (newindex == m_current_hapindex)
    {
        newindex = registry.GetRandom().Long(m_haplotype_dlcells.size());
    }

    m_current_hapindex = newindex;
    return m_haplotype_dlcells[newindex];
}

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

vector<LocusCell> Haplotypes::ChooseRandomHaplotypes()
{
    if (m_haplotype_dlcells.size() == 0)
    {
        ConvertAllelesToDLCells();
    }
    //Choose a new haplotype index
    m_current_hapindex = registry.GetRandom().Long(m_haplotype_dlcells.size());
    return m_haplotype_dlcells[m_current_hapindex];
}

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

vector<LocusCell> Haplotypes::ChooseFirstHaplotypes()
{
    if (m_haplotype_dlcells.size() == 0)
    {
        ConvertAllelesToDLCells();
    }
    //Choose haplotype zero.
    m_current_hapindex = 0;
    return m_haplotype_dlcells[m_current_hapindex];
}

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

vector<LocusCell> Haplotypes::ChooseNextHaplotypes()
{
    assert(m_haplotype_dlcells.size() != 0);
    //Choose the next haplotype index
    if (static_cast<unsigned long>(m_current_hapindex) == m_haplotype_dlcells.size()-1)
    {
        //Already at the last one.
        vector<LocusCell> blankcells;
        return blankcells;
    }
    m_current_hapindex++;
    return m_haplotype_dlcells[m_current_hapindex];
}

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

StringVec1d Haplotypes::GetAlleles() const
{
    StringVec1d retvec;
    for (unsigned long res=0; res<m_haplotype_alleles.size(); res++)
    {
        for (unsigned long allele=0; allele<m_haplotype_alleles[res].size(); allele++)
        {
            retvec.push_back(m_haplotype_alleles[res][allele]);
        }
    }
    return retvec;
}

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

string Haplotypes::GetMarkerData() const
{
    assert(m_haplotype_alleles.size() == m_penetrances.size());
    string markerdata;
    for (unsigned long hapres=0; hapres<m_haplotype_alleles.size(); hapres++)
    {
        markerdata += ToString(m_haplotype_alleles[hapres]) + "  ";
        if (m_penetrances[hapres] < 1)
        {
            markerdata += "(" + ToString(m_penetrances[hapres]) + ") ";
        }
    }
    return markerdata;
}

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

StringVec1d Haplotypes::GetHaplotypesXML(long nspaces) const
{
    string spaces(nspaces, ' ');
    string spaces2(nspaces+2, ' ');

    StringVec1d retvec;
    for (unsigned long hap=0; hap<m_penetrances.size(); hap++)
    {
        retvec.push_back(spaces + MakeTag(xmlstr::XML_TAG_HAPLOTYPES));
        retvec.push_back(spaces2 + MakeTag(xmlstr::XML_TAG_PENETRANCE) + " "
                         + ToString(m_penetrances[hap]) + " "
                         + MakeCloseTag(xmlstr::XML_TAG_PENETRANCE));
        retvec.push_back(spaces2 + MakeTag(xmlstr::XML_TAG_ALLELES)
                         + ToString(m_haplotype_alleles[hap]) + " "
                         + MakeCloseTag(xmlstr::XML_TAG_ALLELES));
        retvec.push_back(spaces + MakeCloseTag(xmlstr::XML_TAG_HAPLOTYPES));
    }
    return retvec;
}

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

bool Haplotypes::MultipleHaplotypes()  const
{
    if (m_haplotype_dlcells.size() > 0)
    {
        return (m_haplotype_dlcells.size() > 1); //we're in phase 2
    }
    else
    {
        return (m_haplotype_alleles.size() > 1); //still in phase 1
    }
}

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

StringVec2d Haplotypes::SetToVecs(multiset<string> stringSet) const
{
    StringVec2d retvecs;
    for (multiset<string>::iterator newstring=stringSet.begin();
         newstring != stringSet.end(); newstring++)
    {
        multiset<string> partialSet = stringSet;
        partialSet.erase(partialSet.find(*newstring));
        if (partialSet.size() == 0)
        {
            StringVec1d strings;
            strings.push_back(*newstring);
            retvecs.push_back(strings);
            return retvecs;
        }
        StringVec2d partialVecs = SetToVecs(partialSet);
        for (StringVec2d::iterator partialVec = partialVecs.begin();
             partialVec != partialVecs.end(); partialVec++)
        {
            (*partialVec).push_back(*newstring);
            retvecs.push_back(*partialVec);
        }
    }
    return retvecs;
}

//------------------------------------------------------------------------------------
// Debugging function.

void Haplotypes::PrintCellsAndAlleles() const
{
    cerr << "Here are the original strings:" << endl;
    for (size_t set=0; set<m_haplotype_alleles.size(); set++)
    {
        cerr << ToString(m_haplotype_alleles[set]) << endl;
    }
    cerr << "And here are the corresponding DLCells.  They should match!" << endl;
    for (size_t set=0; set<m_haplotype_dlcells.size(); set++)
    {
        for (size_t allele=0; allele<m_haplotype_dlcells[set].size(); allele++)
        {
            cerr << registry.GetDataPack().GetRegion(m_regionnum).GetLocus(m_locusname).GetDataModel()->
                CellToData(m_haplotype_dlcells[set][allele][0], 0)
                 << " ";
        }
    }
    cerr << endl;
    cerr << "And finally, the DLCells as raw data." << endl;
    for (size_t set=0; set<m_haplotype_dlcells.size(); set++)
    {
        for (size_t allele=0; allele<m_haplotype_dlcells[set].size(); allele++)
        {
            cerr << m_haplotype_dlcells[set][allele][0]->DLsToString(0,0);
        }
    }
    cerr << endl;
}

//____________________________________________________________________________________
