// $Id: phenotypes.cpp,v 1.6 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: Lucian Smith, Mary Kuhner, Jon Yamato and Joseph Felsenstein */


#include <cassert>

#include "errhandling.h"
#include "locuscell.h" //Have to include this to use Haplotypes
#include "mathx.h"
#include "phenotypes.h"
#include "registry.h"
#include "stringx.h"
#include "vectorx.h"
#include "xml_strings.h"

using namespace std;

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

typedef map<multiset<string>, pair<StringVec1d, DoubleVec1d> >::iterator PhenMapIter;
typedef map<multiset<string>, pair<StringVec1d, DoubleVec1d> >::const_iterator PhenMapConstIter;

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

Phenotypes::Phenotypes(long regnum, string lname)
    : m_regionnum(regnum),
      m_locusname(lname),
      m_phenomap(),
      m_hapsmap()
{
}

//Needed a blank copy in the locus--we replace it if we have a real one.
//
//Boy howdy does this all need to be refactored.  Whew, it stinks.
Phenotypes::Phenotypes(string lname)
    : m_regionnum(FLAGLONG),
      m_locusname(lname),
      m_phenomap(),
      m_hapsmap()
{
}

void Phenotypes::AddPhenotype(const StringVec1d& alleles, string name, double penetrance)
{
    assert (m_regionnum != FLAGLONG);
    multiset<string> alleleSet = VecToSet(alleles);
    PhenMapIter genotype = m_phenomap.find(alleleSet);
    if (genotype == m_phenomap.end())
    {
        StringVec1d names;
        DoubleVec1d penetrances;
        m_phenomap.insert(make_pair(alleleSet, make_pair(names, penetrances)));
        genotype = m_phenomap.find(alleleSet);
        assert(genotype != m_phenomap.end());
    }
    (*genotype).second.first.push_back(name);
    (*genotype).second.second.push_back(penetrance);

}

Haplotypes Phenotypes::ChooseHaplotypes(const StringVec1d& alleles)
{
    multiset<string> alleleSet = VecToSet(alleles);
    PhenMapIter genotype = m_phenomap.find(alleleSet);
    if (genotype == m_phenomap.end())
    {
        string msg = "Unable to find any phenotypes for the set of alleles \"";
        for (size_t allele = 0; allele<alleles.size(); allele++)
        {
            msg += " " + alleles[allele];
        }
        msg += " \".";
        throw data_error(msg);
    }
    DoubleVec1d penetrances = (*genotype).second.second; //for convenience
    if (penetrances.size() == 1)
    {
        return GetHaplotypes((*genotype).second.first[0]);
    }
    double rand = registry.GetRandom().Float();
    double sum = 0;
    for (size_t which=0; which < penetrances.size(); which++)
    {
        sum += penetrances[which];
        if (rand < sum)
        {
            return GetHaplotypes((*genotype).second.first[which]);
        }
    }
    //Just in case
    return GetHaplotypes(*(*genotype).second.first.rbegin());
}

Haplotypes Phenotypes::GetHaplotypes(string name)
{
    if (m_hapsmap.size() == 0)
    {
        MakeHaplotypes();
    }
    map<string, Haplotypes>::iterator hap = m_hapsmap.find(name);
    if (hap == m_hapsmap.end())
    {
        throw data_error("Unable to find a set of haplotypes for the phenotype named " + name + ".");
    }
    return (*hap).second;
}

StringVec1d Phenotypes::GetPhenotypesXML(long nspaces) const
{
    StringVec1d xmllines;
    if (m_phenomap.size() == 0)
    {
        return xmllines;
    }
    string line = MakeIndent(MakeTag(xmlstr::XML_TAG_PHENOTYPES), nspaces);
    xmllines.push_back(line);
    nspaces += INDENT_DEPTH;

    for (PhenMapConstIter genotype=m_phenomap.begin(); genotype != m_phenomap.end(); genotype++)
    {
        line = MakeIndent(MakeTag(xmlstr::XML_TAG_GENOTYPE), nspaces);
        xmllines.push_back(line);
        nspaces += INDENT_DEPTH;

        line = MakeTag(xmlstr::XML_TAG_ALLELES) + " "
            + ToString(MultisetElemToString((*genotype).first))
            + MakeCloseTag(xmlstr::XML_TAG_ALLELES);
        xmllines.push_back(MakeIndent(line, nspaces));

        for (size_t phenotype=0; phenotype<(*genotype).second.second.size(); phenotype++)
        {
            line = MakeIndent(MakeTag(xmlstr::XML_TAG_PHENOTYPE), nspaces);
            xmllines.push_back(line);
            nspaces += INDENT_DEPTH;

            line = MakeTag(xmlstr::XML_TAG_PHENOTYPE_NAME) + " "
                + (*genotype).second.first[phenotype] + " "
                + MakeCloseTag(xmlstr::XML_TAG_PHENOTYPE_NAME);
            xmllines.push_back(MakeIndent(line, nspaces));

            line = MakeTag(xmlstr::XML_TAG_PENETRANCE) + " "
                + ToString((*genotype).second.second[phenotype]) + " "
                + MakeCloseTag(xmlstr::XML_TAG_PENETRANCE);
            xmllines.push_back(MakeIndent(line, nspaces));

            nspaces -= INDENT_DEPTH;
            line = MakeIndent(MakeCloseTag(xmlstr::XML_TAG_PHENOTYPE), nspaces);
            xmllines.push_back(line);
        }

        nspaces -= INDENT_DEPTH;
        line = MakeIndent(MakeCloseTag(xmlstr::XML_TAG_GENOTYPE), nspaces);
        xmllines.push_back(line);
    }
    nspaces -= INDENT_DEPTH;
    line = MakeIndent(MakeCloseTag(xmlstr::XML_TAG_PHENOTYPES), nspaces);
    xmllines.push_back(line);
    return xmllines;
}

void Phenotypes::MakeHaplotypes()
{
    for(PhenMapIter genotype=m_phenomap.begin(); genotype != m_phenomap.end(); genotype++)
    {
        ScaleToSumToOne((*genotype).second.second);
        assert((*genotype).second.first.size() == (*genotype).second.second.size());
        for (size_t phenNum=0; phenNum<(*genotype).second.second.size(); phenNum++)
        {
            multiset<string> alleles = (*genotype).first;
            string phenName = (*genotype).second.first[phenNum];
            double penetrance = (*genotype).second.second[phenNum];

            map<string, Haplotypes>::iterator hap = m_hapsmap.find(phenName);
            if (hap == m_hapsmap.end())
            {
                Haplotypes newhaps(m_regionnum, m_locusname);
                m_hapsmap.insert(make_pair(phenName, newhaps));
                hap = m_hapsmap.find(phenName);
                assert(hap != m_hapsmap.end());
            }
            (*hap).second.AddHaplotype(alleles, penetrance);
        }
    }
}

multiset<string> Phenotypes::VecToSet(StringVec1d vec)
{
    multiset<string> retset;
    for (size_t str=0; str<vec.size(); str++)
    {
        retset.insert(vec[str]);
    }
    return retset;
}

//____________________________________________________________________________________
