// $Id: dlcell.cpp,v 1.46 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 */


// This file contains the implementation code for the data-likelihood
// storage object.

#include <cassert>
#include <cmath>
#include <cstring>
#include <iostream>                     //debug

#include "dlcell.h"
#include "dlmodel.h"
#include "errhandling.h"
#include "definitions.h"
#include "registry.h"
#include "cellmanager.h"
#include "stringx.h"                    // for ToString() use in debug function DLsToString()

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

// this turns on the data likelihood details for testing
//#define DATA_LIKELIHOOD_DETAILS

using namespace std;

//------------------------------------------------------------------------------------
// DLCell
//------------------------------------------------------------------------------------

DLCell::DLCell(long markers, long cats, long bins)
    : Cell(),
      m_nmarkers(markers),
      m_ncats(cats),
      m_nbins(bins),
      m_norms(markers, 0.0)
{
    // obtain an internal array from the free store
    m_identifier.first = m_nmarkers;
    m_identifier.second = m_ncats;
    m_identifier.third = m_nbins;

    m_DLs = registry.GetCellManager().GetArray(m_identifier, *this);
} // DLCell constructor

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

DLCell::~DLCell()
{
    // return the internal array to the free store
    registry.GetCellManager().FreeArray(m_identifier, m_DLs);
    m_DLs = NULL;
} // DLCell destructor

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

Cell* DLCell::Copy() const
{
    Cell* pCell = Clone();
    pCell->CopyInitialize(*this);
    return pCell;

} // Copy

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

void DLCell::EmptyCell()
{
    for (long posn = 0; posn < m_nmarkers; ++posn)
    {
        for (long cat = 0; cat < m_ncats; ++cat)
        {
            for (long base = 0; base < m_nbins; ++base)
            {
                m_DLs[posn][cat][base] = 0;
            }
        }
    }
}

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

void DLCell::CopyInitialize(const Cell& src)
{
    const DLCell& srcell = dynamic_cast<const DLCell&>(src);

    assert(m_nmarkers == srcell.m_nmarkers);
    assert(m_ncats == srcell.m_ncats);
    assert(m_nbins == srcell.m_nbins);

    m_norms = srcell.m_norms;
    // m_nmarkers+1 because there is an extra cell at the end, to
    // allow an STL-like interface where one past the end is a valid
    // address

    long arraySize = (m_nmarkers+1) * m_ncats * m_nbins;
    //  cerr << "arraysize=" << arraySize << endl;
    // memcpy for speed--be careful!
    memcpy(m_DLs[0][0], srcell.m_DLs[0][0], arraySize*sizeof(double));

} // CopyInitialize

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

bool DLCell::IsValidPosition(long pos) const
{
    // for some purposes, pos == m_nmarkers is in fact valid; this
    // allows "one past the end" logic emulating the STL

    if (0 <= pos && pos <= m_nmarkers) return true;
    return false;
} // IsValidPosition

//------------------------------------------------------------------------------------
// Make a datalikelihood array appropriately sized for this Cell

cellarray DLCell::MakeArray()
{
    long len = m_nmarkers + 1;
    // provide a STL like interface into the m_DLs array.
    // (i.e. provides a legal reference one past the normal end
    // of the array)

    // set up a 3-dimensional array in which all positions are
    // contiguous, so that memcpy can be used on it (a speed
    // optimization).

    long site, category;
    cellarray dl;

    dl = new double**[len];

    dl[0] = new double*[len * m_ncats];
    for (site = 0; site < len; ++site)
    {
        dl[site] = dl[0] + site * m_ncats;
    }

    dl[0][0] = new double[len * m_ncats * m_nbins];
    for (site = 0; site < len; ++site)
    {
        for (category = 0; category < m_ncats; ++category)
        {
            dl[site][category] = dl[0][0] + site * m_ncats * m_nbins + category * m_nbins;
        }
    }

    return dl;

} // MakeArray

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

void DLCell::SwapDLs(Cell_ptr othercell, long pos)
{
    // This creates a 2D array (categories x bins) and uses it to
    // swap the information for a single marker between two DLCells.
    // All bets are off if they are not the same type!

    // WARNING -- Not exception safe due to raw new.

    // take pointers to the two things to be swapped
    double** otherDLs = othercell->GetSiteDLs(pos);
    double** myDLs = GetSiteDLs(pos);

    // create temporary storage for the swap
    double** newDLs = new double* [m_ncats];
    newDLs[0] = new double [m_ncats * m_nbins];
    long cat;
    for(cat = 1; cat < m_ncats; ++cat)
        newDLs[cat] = newDLs[0] + cat * m_nbins;

    // memcpy for speed
    memcpy(newDLs[0], myDLs[0], m_ncats * m_nbins * sizeof(double));

    // swap
    SetSiteDLs(pos, otherDLs);
    othercell->SetSiteDLs(pos, newDLs);

    // release temporary storage
    delete [] newDLs[0];
    delete [] newDLs;

} // SwapDLs

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

void DLCell::SumMarkers(long startpos, long endpos, bool normalize)
{
    long cat, bin, pos;

    if (normalize)
    {

        for (cat = 0; cat < m_ncats; ++cat)
        {
            for (bin = 0; bin < m_nbins; ++bin)
            {
                double result = 0.0;
                for (pos = startpos; pos != endpos; ++pos)
                {
                    double markerval = log(m_DLs[pos][cat][bin]) + GetNorms(pos);
                    if (markerval > EXPMIN) result += exp(markerval);
                }
#ifdef DATA_LIKELIHOOD_DETAILS
                cerr << "SumMarkers-normalized cat:" << cat
                     << " bin:" << bin
                     << " result:" << result
                     << endl;
#endif
                if (result == 0.0)   // never found a good value?
                    m_DLs[endpos][cat][bin] = exp(EXPMIN);
                else
                    m_DLs[endpos][cat][bin] = result;
            }
        }

        // renormalize and re-set norms
        double newnorm = Normalize(m_DLs[endpos]);
        SetNorms(newnorm, endpos);

    }
    else                                // no normalization
    {

        for (cat = 0; cat < m_ncats; ++cat)
        {
            for (bin = 0; bin < m_nbins; ++bin)
            {
                double result = 0.0;
                for (pos = startpos; pos != endpos; ++pos)
                {
                    result += m_DLs[pos][cat][bin];
                }
#ifdef DATA_LIKELIHOOD_DETAILS
                cerr << "SumMarkers-not normalized cat:" << cat
                     << " bin:" << bin
                     << " result:" << result
                     << endl;
#endif

                m_DLs[endpos][cat][bin] = result;
            }
        }
    }

} // SumMarkers

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

bool DLCell::IsSameAs(const Cell_ptr othercell, long pos) const
{
    double** otherDLs = othercell->GetSiteDLs(pos);
    double** myDLs = GetSiteDLs(pos);
    long cat, bin;
    for(cat = 0; cat < m_ncats; ++cat)
        for (bin = 0; bin < m_nbins; ++bin)
            if (myDLs[cat][bin] != otherDLs[cat][bin]) return false;

    return true;
} // DLCell::IsSameAs

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

long DLCell::DiffersFrom(Cell_ptr othercell) const
{
    long marker;
    for(marker = 0; marker < m_nmarkers; ++marker)
        if (!IsSameAs(othercell,marker)) return marker;

    return FLAGLONG;

} // DLCell::DiffersFrom

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

void DLCell::SetAllCategoriesTo(DoubleVec1d& state, long posn)
{
    for (long cat = 0; cat < m_ncats; cat++)
    {
        for (unsigned long nstate=0; nstate < state.size(); nstate++)
        {
            m_DLs[posn][cat][nstate] = state[nstate];
        }
    }
}

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

DoubleVec1d DLCell::GetStateFor(long posn, long cat) const
{
    DoubleVec1d state;
    for (long nstate = 0; nstate < m_nbins; nstate++)
    {
        state.push_back(m_DLs[posn][cat][nstate]);
    }
    return state;
}


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

void DLCell::SetStateTo (long posn, long cat, DoubleVec1d state)
{
    for (long nstate = 0; nstate < m_nbins; nstate++)
    {
        m_DLs[posn][cat][nstate] = state[nstate];
    }
}

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

void DLCell::AddTo(const Cell_ptr othercell)
{
    for (long pos = 0; pos < m_nmarkers; pos++)
    {
        double** otherDLs = othercell->GetSiteDLs(pos);
        double** myDLs = GetSiteDLs(pos);
        long cat, bin;
        for(cat = 0; cat < m_ncats; ++cat)
        {
            for (bin = 0; bin < m_nbins; ++bin)
            {
                myDLs[cat][bin] += otherDLs[cat][bin];
            }
        }
    }
}

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

void DLCell::SubtractFrom(const Cell_ptr othercell)
{
    for (long pos = 0; pos < m_nmarkers; pos++)
    {
        double** otherDLs = othercell->GetSiteDLs(pos);
        double** myDLs = GetSiteDLs(pos);
        long cat, bin;
        for(cat = 0; cat < m_ncats; ++cat)
        {
            for (bin = 0; bin < m_nbins; ++bin)
            {
                myDLs[cat][bin] -= otherDLs[cat][bin];
            }
        }
    }
}

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

void DLCell::MultiplyBy(double mult)
{
    for (long pos = 0; pos < m_nmarkers; pos++)
    {
        double** myDLs = GetSiteDLs(pos);
        long cat, bin;
        for(cat = 0; cat < m_ncats; ++cat)
        {
            for (bin = 0; bin < m_nbins; ++bin)
            {
                myDLs[cat][bin] *= mult;
            }
        }
    }
}

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

void DLCell::MultiplyBy(const Cell_ptr othercell)
{
    for (long pos = 0; pos < m_nmarkers; pos++)
    {
        double** otherDLs = othercell->GetSiteDLs(pos);
        double** myDLs = GetSiteDLs(pos);
        long cat, bin;
        for(cat = 0; cat < m_ncats; ++cat)
        {
            for (bin = 0; bin < m_nbins; ++bin)
            {
                myDLs[cat][bin] *= otherDLs[cat][bin];
            }
        }
    }
}

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

LongVec1d DLCell::GetOnes(long marker) const
{
    LongVec1d ones;
    double** dls = GetSiteDLs(marker);
    //Just mess with the first category
    for (long bin = 0; bin < m_nbins; ++bin)
    {
        if (dls[0][bin] == 1.0)
        {
            ones.push_back(bin);
        }
    }
    return ones;
}

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

string DLCell::DLsToString(long start, long end) const
{
    string lines;

    long line;
    for(line = start; line <= end; ++line)
    {
        lines += "marker " + ToString(line) + ": ";
        double** dls = GetSiteDLs(line);
        long cat;
        for(cat = 0; cat < m_ncats; ++cat)
        {
            lines += "{";
            long bin;
            for(bin = 0; bin < m_nbins; ++bin)
            {
                lines += ToString(dls[cat][bin]);
                if (bin != m_nbins - 1) lines += ",";
            }
            lines += "}";
            if (cat != m_ncats - 1) lines += ", ";
        }
        lines += "\n";
    }

    return lines;

} // DLCell::DLsToString

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

double DLCell::Normalize(double** siteDLs)
{
    double biggest = NEG_MAX;
    long cat, bin;
    for(cat = 0; cat < m_ncats; ++cat)
    {
        for(bin = 0; bin < m_nbins; ++bin)
        {
            if (siteDLs[cat][bin] > biggest) biggest = siteDLs[cat][bin];
        }
    }

    for(cat = 0; cat < m_ncats; ++cat)
    {
        for(bin = 0; bin < m_nbins; ++bin)
        {
            siteDLs[cat][bin] /= biggest;
        }
    }

    return log(biggest);

} // Normalize

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

void DLCell::SetSiteDLs(long posn, double **siteDLs)
{
    assert(IsValidPosition(posn));
    long siteSize = m_ncats * m_nbins;
    memcpy(m_DLs[posn][0], siteDLs[0], siteSize * sizeof(double));
}

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

void DLCell::AddToSiteDLs(long posn, double **siteDLs)
{
    assert(IsValidPosition(posn));
    long cat, bin;
    for(cat = 0; cat < m_ncats; ++cat)
    {
        for (bin = 0; bin < m_nbins; ++bin)
        {
            m_DLs[posn][cat][bin] += siteDLs[cat][bin];
        }
    }
}

//------------------------------------------------------------------------------------
// NullCell
//------------------------------------------------------------------------------------

NullCell::NullCell()
    : Cell()
{
    // deliberately blank
} // NullCell constructor

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

void NullCell::Initialize (const StringVec1d&, const DataModel_ptr)
{
    assert (false); // should never call this!

} // Initialize

DoubleVec1d NullCell::GetStateFor(long, long) const
{
    throw implementation_error("No state for this data likelihood.");
}

//------------------------------------------------------------------------------------
// NucCell
//------------------------------------------------------------------------------------

NucCell::NucCell(long markers, long cats)
    : DLCell(markers, cats, BASES)
{
    // deliberately left blank
} // NucCell constructor

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

void NucCell::Initialize(const StringVec1d &sequence, const DataModel_ptr trans)
{
    long posn, base, cat;

    string postring;
    vector<double> likes;

    // could be OPTIMIZED

    for (posn = 0; posn < m_nmarkers; ++posn)
    {
        postring = sequence[posn];
        likes = trans->DataToLikes(postring);
        for (cat = 0; cat < m_ncats; ++cat)
        {
            for (base = 0; base < m_nbins; ++base)
            {
                m_DLs[posn][cat][base] = likes[base];
            }
        }
    }
} // Initialize

//------------------------------------------------------------------------------------
// DNACell
//------------------------------------------------------------------------------------

DNACell::DNACell(long markers, long cats)
    : NucCell(markers, cats)
{
    // intentionally blank
} // DNACell constructor

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

Cell *DNACell::Clone() const
{
    Cell *pDLCell = new DNACell(m_nmarkers, m_ncats);
    return pDLCell;
}

//------------------------------------------------------------------------------------
// SNPCell
//------------------------------------------------------------------------------------

SNPCell::SNPCell(long markers, long cats)
    : NucCell(markers, cats)
{
    // deliberately blank
} // SNPCell constructor

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

SNPCell::~SNPCell()
{
    // deliberately blank
}

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

Cell *SNPCell::Clone() const
{
    Cell *pDLCell  = new SNPCell(m_nmarkers, m_ncats);
    return pDLCell;
}

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

void SNPCell::Initialize(const StringVec1d&, const DataModel_ptr trans)
{
    // We do not use the sequence input -- SNPCell::Initialize should
    // only be used to initialize invariant cells. We keep the first
    // argument for conformity with the base class interface.

    long categ, invar, base;

    // In the m_DLs array, the basic idea is to set all entries to 0.0
    // except for the entries where the invariant equals the base,
    // where they are set to 1.0.
    // HOWEVER,
    // since we've now got a data uncertainty model, each of the
    // invariant bases needs to go through the error correction
    // transformation first.

    for (invar = 0; invar < INVARIANTS; ++invar)
    {
        vector<double> likes = trans->DataToLikes(SINGLEBASES[invar]);
        for (categ = 0; categ < m_ncats; ++categ)
        {
            for (base = 0; base < m_nbins; ++base)
            {
                m_DLs[invar][categ][base] = likes[base];
            }
        }
    }
}

//------------------------------------------------------------------------------------
// AlleleCell
//------------------------------------------------------------------------------------

AlleleCell::AlleleCell(long markers, long cats, long bins)
    : DLCell(markers, cats, bins)
{
    // deliberately blank
} // AlleleCell constructor

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

Cell *AlleleCell::Clone() const
{
    Cell *pDLCell  = new AlleleCell(m_nmarkers, m_ncats, m_nbins);
    return pDLCell;
} // Clone

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

void AlleleCell::Initialize(const StringVec1d &sequence, const DataModel_ptr trans)
{
    long posn, bin, cat;
    vector<double> likes;

    // could be OPTIMIZED

    for (posn = 0; posn < m_nmarkers; ++posn)
    {
        likes = trans->DataToLikes(sequence[posn], posn);
        for (cat = 0; cat < m_ncats; ++cat)
        {
            for (bin = 0; bin < m_nbins; ++bin)
            {
                m_DLs[posn][cat][bin] = likes[bin];
            }
        }
    }
} // Initialize

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

triplet::triplet()
    : first(0),
      second(0),
      third(0)
{
    // intentionally blank
} // triplet default constructor

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

triplet::triplet(long f, long s, long t)
    : first(f),
      second(s),
      third(t)
{
    // intentionally blank
} // triplet constructor

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

bool triplet::operator<(const triplet& rhs) const
{
    if (first != rhs.first) return (first < rhs.first);
    if (second != rhs.second) return (second < rhs.second);
    return (third < rhs.third);
} // triplet operator<

//____________________________________________________________________________________
