// $Id: datatype.cpp,v 1.40 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: Peter Beerli, Mary Kuhner, Jon Yamato and Joseph Felsenstein */


#include <cassert>
#include <cstring>

#include "datatype.h"
#include "dlcalc.h"
#include "dlcell.h"
#include "dlmodel.h"
#include "datapack.h"
#include "locus.h"
#include "errhandling.h"
#include "stringx.h"

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

using std::vector;

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

string ToString(data_type dtype)
{
    switch(dtype)
    {
        case dtype_DNA:
            return "DNA";
        case dtype_SNP:
            return "SNP";
        case dtype_msat:
            return "Microsatellite";
        case dtype_kallele:
            return "K-Allele";
    }

    assert(false); //uncaught data type;
    return "unknown";
}

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

string ToString(data_source dsource)
{
    switch(dsource)
    {
        case dsource_study:
            return "Study";
            break;
        case dsource_panel:
            return "Panel";
            break;
    }

    assert(false); //uncaught data source;
    return "unknown";
}

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

long DataType::GetPairwiseMarkers(const Locus& locus) const
{
    return locus.GetNmarkers();
} // DataType::GetPairwiseMarkers

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

double DataType::DifferenceWithinPop(StringVec2d::const_iterator
                                     begin1, StringVec2d::const_iterator end) const
{
    double ndiffs = 0.0;

    // half triangular matrix
    StringVec2d::const_iterator group1 = begin1;
    StringVec2d::const_iterator group2;

    for ( ; group1 != end; ++group1)
    {
        group2 = group1;
        ++group2;  // workaround for compilers that don't have += for iterators

        for ( ; group2 != end; ++group2)
        {
            StringVec1d::const_iterator site1, site2;
            for (site1 = group1->begin(), site2 = group2->begin();
                 site1 != group1->end();
                 ++site1, ++site2)
            {
                if (!IsEquivalent(*site1, *site2)) ++ndiffs;
            }
        }
    }

    return ndiffs;
} // DifferenceWithinPop

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

double DataType::DifferenceBetweenPops(StringVec2d::const_iterator begin1,
                                       StringVec2d::const_iterator end1, StringVec2d::const_iterator begin2,
                                       StringVec2d::const_iterator end2) const
{
    double ndiffs = 0.0;
    StringVec2d::const_iterator group1 = begin1;
    StringVec2d::const_iterator group2 = begin2;

    for ( ; group1 != end1; ++group1)
    {
        for ( ; group2 != end2; ++group2)
        {
            StringVec1d::const_iterator site1, site2;
            for (site1 = group1->begin(), site2 = group2->begin();
                 site1 != group1->end();
                 ++site1, ++site2)
            {
                if (!IsEquivalent(*site1, *site2)) ++ndiffs;
            }
        }
    }

    return ndiffs;
} // DifferenceBetweenPops

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

DoubleVec1d DataType::CalcFw(const Locus& locus) const
{
    // NB This cannot be used with DIVMIG
    StringVec3d data(locus.GetPartitionGeneticData(force_MIG));
    DoubleVec1d fwithin(data.size(),0.0);
    DoubleVec1d::iterator fw;
    StringVec3d::const_iterator pop;

    for(pop = data.begin(), fw = fwithin.begin(); pop != data.end(); ++pop, ++fw)
    {
        double nn = GetPairwiseMarkers(locus) * (pop->size()*pop->size() - pop->size())/2.0;
        double within = DifferenceWithinPop(pop->begin(), pop->end());

        if (nn > 0.0) within /= nn;

        *fw = 1.0 - within;
    }

    return fwithin;
} // NucleotideType::CalcFw

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

DoubleVec1d DataType::CalcFb(const Locus& locus) const
{
    // NB This cannot be used with DIVMIG
    StringVec3d data(locus.GetPartitionGeneticData(force_MIG));
    long npops = data.size();
    DoubleVec1d fbetween(npops*npops,0.0);
    long pop1, pop2;

    for(pop1 = 0; pop1 < npops; ++pop1)
    {
        for(pop2 = pop1 + 1; pop2 < npops; ++pop2)
        {
            double nn = GetPairwiseMarkers(locus) * data[pop1].size() * data[pop2].size();
            double between = DifferenceBetweenPops(data[pop1].begin(),
                                                   data[pop1].end(), data[pop2].begin(),
                                                   data[pop2].end());


            if (nn > 0.0) between /= nn;

            fbetween[pop1 * npops + pop2] = 1.0 - between;
            fbetween[pop2 * npops + pop1] = 1.0 - between;
        }
    }

    return fbetween;
} // NucleotideType::CalcFb

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

DoubleVec1d DataType::CalcXFw(const Locus& locus,const StringVec3d& data) const
{
    DoubleVec1d fwithin;

    StringVec3d::const_iterator part;
    for(part = data.begin(); part != data.end(); ++part)
    {
        double nn = GetPairwiseMarkers(locus) * (part->size()*part->size() - part->size())/2.0;
        double within = DifferenceWithinPop(part->begin(), part->end());

        if (nn > 0.0)
        {
            within /= nn;
        }

        fwithin.push_back(1.0 - within);
    }

    return fwithin;
} // NucleotideType::CalcXFw

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

// argument "data" is locus.GetAllCrossPartitionGeneticData()
DoubleVec1d DataType::CalcXFb(const Locus& locus, const StringVec3d& data) const
{
    long nparts = data.size();
    DoubleVec1d fbetween(nparts*nparts,0.0);
    long part1, part2;

    for(part1 = 0; part1 < nparts; ++part1)
    {
        for(part2 = part1 + 1; part2 < nparts; ++part2)
        {
            double nn = GetPairwiseMarkers(locus) * data[part1].size() * data[part2].size();
            double between = DifferenceBetweenPops(data[part1].begin(),
                                                   data[part1].end(), data[part2].begin(),
                                                   data[part2].end());

            if (nn > 0.0) between /= nn;

            fbetween[part1 * nparts + part2] = 1.0 - between;
            fbetween[part2 * nparts + part1] = 1.0 - between;
        }
    }

    return fbetween;
} // NucleotideType::CalcXFb

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

model_type NucleotideType::DefaultModelType()
{
    return F84;
} // DefaultModelType

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

long DataType::CalcNVarMarkers(const StringVec2d& data) const
{
    long nvarmarkers = 0;
    long marker, nmarkers = data[0].size();

    for (marker = 0; marker < nmarkers; ++ marker)
    {
        bool varying = false;
        long indiv, nindiv = data.size();

        for(indiv = 1; indiv < nindiv && !varying; ++indiv)
            if (!IsEquivalent(data[0][marker],data[indiv][marker]))
                varying = true;

        if (varying) ++nvarmarkers;
    }

    return nvarmarkers;
} // CalcNVarMarkers

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

wakestats DataType::CalcNPairDiffs(const StringVec2d& data) const
{
    wakestats sums(0L,0L);
    StringVec2d::const_iterator indiv1;
    for(indiv1 = data.begin(); indiv1 != data.end(); ++indiv1)
    {
        StringVec2d::const_iterator indiv2;
        for(indiv2 = indiv1+1; indiv2 != data.end(); ++indiv2)
        {
            assert(indiv1->size() == indiv2->size());
            long pos, npos = indiv1->size(), count = 0;
            for(pos = 0; pos < npos; ++pos)
                count += ((IsEquivalent((*indiv1)[pos],(*indiv2)[pos])) ?
                          0L : 1L);
            sums.first += count;
            sums.second += count * count;
        }
    }

    return sums;
} // CalcNPairDiffs

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

bool DataType::IsEquivalent(const string& data1, const string& data2) const
{
    return (data1 == data2);
} // IsEquivalent


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

bool NucleotideType::Proofread(const string &raw, StringVec1d &clean,
                               string& baddata) const
{
    long rawsize, position;
    char ch;

    clean.clear();
    rawsize = raw.size();

    for (position = 0; position < rawsize; ++position)
    {
        ch = toupper(raw[position]);
        if (isspace(ch))
            continue;

        if (!strchr("ACGTUMRWSYKVHDBNOX?-", ch))
        {
            baddata = ToString(ch);
            return false;
        }

        string cleanstr;
        cleanstr += ch;
        clean.push_back(cleanstr);
    }

    return true;
}

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

bool NucleotideType::IsEquivalent(const string& data1, const string& data2) const
{
    if (CaselessStrCmp(data1, data2)) return true;

    // otherwise, they might still be equivalent due to code ambiguities

    string unknowns = "NOXnox?-";  // all of these mean 'unknown'

    bool isunknown1 = (data1.find_first_of(unknowns) != string::npos);
    bool isunknown2 = (data2.find_first_of(unknowns) != string::npos);

    if (isunknown1 && isunknown2) return true;

    string ts = "TUtu";            // all of these mean T

    bool ist1 = (data1.find_first_of(ts) != string::npos);
    bool ist2 = (data2.find_first_of(ts) != string::npos);

    if (ist1 && ist2) return true;

    return false;
} // NucleotideType::IsEquivalent

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

DNASeqType::DNASeqType()
{
    m_name   = "DNA Sequence";
    m_xmltag = lamarcstrings::DNA;
}

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

LocusCell DNASeqType::CreateDLCell(const Locus& locus) const
{
    long len = locus.GetNmarkers();
    long cat = locus.GetDataModel()->GetNcategories();
    vector<Cell_ptr> newcells;
    Cell_ptr cell(new DNACell(len, cat));

    newcells.push_back(cell);

    return LocusCell(newcells);
} // CreateDLCell

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

LocusCell DNASeqType::CreateInitializedDLCell(const Locus& locus,
                                              const StringVec1d& tdata) const
{
    long len = locus.GetNmarkers();
    long cat = locus.GetDataModel()->GetNcategories();
    vector<Cell_ptr> newcells;
    Cell_ptr cell(new DNACell(len, cat));

    cell->Initialize(tdata, locus.GetDataModel());
    newcells.push_back(cell);

    return LocusCell(newcells);
} // CreateInitializedDLCell

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

DLCalc_ptr DNASeqType::CreateDLCalculator(const Locus& locus) const
{
    return DLCalc_ptr(new DNACalculator(locus));
}

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

SNPDataType::SNPDataType()
{
    m_name   = "SNP data";
    m_xmltag = lamarcstrings::SNP;
}

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

long SNPDataType::GetPairwiseMarkers(const Locus& locus) const
{
    return locus.GetNsites();
} // SNPDataType::GetPairwiseMarkers

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

LocusCell SNPDataType::CreateDLCell(const Locus& locus) const
{
    long len = locus.GetNmarkers();
    int  cat = locus.GetDataModel()->GetNcategories();
    vector<Cell_ptr> newcells;

    // A SNP data type actually has *two* DLCells, one for variable and one
    // for non-variable sites
    Cell_ptr variant(new DNACell(len, cat));
    Cell_ptr invariant(new SNPCell(5, cat));

    newcells.push_back(variant);
    newcells.push_back(invariant);
    return LocusCell(newcells);
}

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

LocusCell SNPDataType::CreateInitializedDLCell(const Locus& locus,
                                               const StringVec1d& tdata) const
{
    long len = locus.GetNmarkers();
    int  cat = locus.GetDataModel()->GetNcategories();
    vector<Cell_ptr> newcells;

    // A SNP data type actually has *two* DLCells, one for variable and one
    // for non-variable sites

    Cell_ptr variant(new DNACell(len, cat));
    variant->Initialize(tdata, locus.GetDataModel());

    Cell_ptr invariant(new SNPCell(5, cat));
    invariant->Initialize(tdata, locus.GetDataModel());

    newcells.push_back(variant);
    newcells.push_back(invariant);

    return LocusCell(newcells);
} // CreateInitializedDLCell

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

DLCalc_ptr SNPDataType::CreateDLCalculator(const Locus& locus) const
{
    return DLCalc_ptr(new SNPCalculator(locus));
}

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

LocusCell AlleleType::CreateDLCell(const Locus& locus) const
{
    long len = locus.GetNmarkers();
    long cat = locus.GetDataModel()->GetNcategories();
    long bins = locus.GetDataModel()->GetNbins();
    vector<Cell_ptr> newcells;
    Cell_ptr cell(new AlleleCell(len, cat, bins));

    newcells.push_back(cell);

    return LocusCell(newcells);
} // CreateDLCell

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

LocusCell AlleleType::CreateInitializedDLCell(const Locus& locus,
                                              const StringVec1d& tdata) const
{
    long len = locus.GetNmarkers();
    long cat = locus.GetDataModel()->GetNcategories();
    long bins = locus.GetDataModel()->GetNbins();
    vector<Cell_ptr> newcells;
    Cell_ptr cell(new AlleleCell(len, cat, bins));

    cell->Initialize(tdata, locus.GetDataModel());
    newcells.push_back(cell);

    return newcells;
} // CreateInitializedDLCell

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

MSType::MSType()
    : AlleleType()
{
    m_name   = "Microsatellite";
    m_xmltag = lamarcstrings::MICROSAT;
}

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

bool MSType::Proofread(const string &raw, StringVec1d &clean,
                       string& baddata) const
{
    unsigned long index;
    char c;
    string allele;

    for (index = 0; index < raw.size(); ++index)
    {
        c = raw[index];
        if (!isspace(c)) allele += c;
        else
        {
            if (!allele.empty()) clean.push_back(allele);
            allele.erase();
        }
    }

    // needed in case the last allele is not followed by spaces....
    if (!allele.empty()) clean.push_back(allele);

    // error checking:  microsat data must be either an integer or ?
    for (index = 0; index < clean.size(); ++index)
    {
        if (clean[index] == "?") continue;  // means unknown data
        if (!IsInteger(clean[index]))
        {
            baddata = clean[index];
            return false;  // illegal data found
        }
    }

    return true;
}

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

DLCalc_ptr MSType::CreateDLCalculator(const Locus& locus) const
{
    return DLCalc_ptr(new AlleleCalculator(locus));
}

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

model_type MSType::DefaultModelType()
{
    return Brownian;
} // DefaultModelType

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

KAlleleType::KAlleleType()
    : AlleleType()
{
    m_name   = "K-Allele";
    m_xmltag = lamarcstrings::KALLELE;
} // KAlleleType constructor

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

bool KAlleleType::Proofread(const string& raw, StringVec1d& clean,
                            string&) const //string 'baddata' not used.
{
    unsigned long index;
    char c;
    string allele;

    for (index = 0; index < raw.size(); ++index)
    {
        c = raw[index];
        if (!isspace(c)) allele += c;
        else
        {
            if (!allele.empty()) clean.push_back(allele);
            allele.erase();
        }
    }

    // needed in case the last allele is not followed by spaces....
    if (!allele.empty()) clean.push_back(allele);

    // no error checking possible, as any character but
    // whitespace is allowable

    return true;
} // Proofread

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

DLCalc_ptr KAlleleType::CreateDLCalculator(const Locus& locus) const
{
    return DLCalc_ptr(new AlleleCalculator(locus));
} // CreateDLCalculator

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

model_type KAlleleType::DefaultModelType()
{
    return KAllele;
} // DefaultModelType

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


// Free helper function

DataType* CreateDataType(const string tag)
{
    if (CaselessStrCmp(tag, lamarcstrings::DNA)) return new DNASeqType;
    if (CaselessStrCmp(tag, lamarcstrings::SNP)) return new SNPDataType;
    if (CaselessStrCmp(tag, lamarcstrings::MICROSAT)) return new MSType;
    if (CaselessStrCmp(tag, lamarcstrings::KALLELE)) return new KAlleleType;
    data_error e("Unknown data type encountered: " + tag);
    throw e;
} // CreateDataType


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

bool ModelTypeAcceptsDataType(model_type modelType, data_type dtype)
{
    switch(modelType)
    {
        case F84:
        case GTR:
            // Nucleotide models can accept only DNA and SNP data types.
            switch(dtype)
            {
                case dtype_DNA:
                case dtype_SNP:
                    return true;
                case dtype_msat:
                case dtype_kallele:
                    return false;
            }
            assert(false); //uncaught data type.
            return false;
            break;
        case Brownian:
        case Stepwise:
        case MixedKS:
            // Brownian model is appropriate only for microsat data because it
            //  relies on ordering of alleles.
            // The stepwise model can accept only microsatellite data.
            // The MixedKS model includes a Stepwise aspect, and therefore has
            //  the same restrictions.
            switch(dtype)
            {
                case dtype_msat:
                    return true;
                case dtype_DNA:
                case dtype_SNP:
                case dtype_kallele:
                    return false;
            }
            assert(false); //uncaught data type.
            return false;
            break;
        case KAllele:
            // KAllele model can be used for elecrophoretic, phenotypic,
            // or possibly microsatellite types.
            switch(dtype)
            {
                case dtype_msat:
                case dtype_kallele:
                    return true;
                case dtype_DNA:
                case dtype_SNP:
                    return false;
            }
            assert(false); //uncaught data type.
            return false;
            break;
    };

    assert(false); //uncaught data model.
    return false;
}

model_type DefaultModelForDataType(data_type dtype)
{
    switch(dtype)
    {
        case dtype_DNA:
        case dtype_SNP:
            return F84;
        case dtype_msat:
            return Brownian;
        case dtype_kallele:
            return KAllele;
    }

    assert(false); //uncaught data type.
    return F84;
}

//____________________________________________________________________________________
