/*!
 * @file swp_surfvol_cubic_primary.cpp
 * @author Robert Patterson
 *
  Project:        sweepc (population balance solver)
  Sourceforge:    http://sourceforge.net/projects/mopssuite

  Copyright (C) 2013 Robert I A Patterson

  File purpose:
    Implementation of a particle the approximates cuboidal crystals by
    storing their surface area and volume
    @brief Implementation of cuboidal crystals described by volume and surface area.

  Licence:
    This file is part of "sweepc".

    sweepc is free software; you can redistribute it and/or
    modify it under the terms of the GNU General Public License
    as published by the Free Software Foundation; either version 2
    of the License, or (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.

  Contact:
    Prof Markus Kraft
    Dept of Chemical Engineering
    University of Cambridge
    New Museums Site
    Pembroke Street
    Cambridge
    CB2 3RA
    UK

    Email:       mk306@cam.ac.uk
    Website:     http://como.cheng.cam.ac.uk
*/

#include "swp_surfvol_cubic_primary.h"

#include "swp_aggmodel_type.h"
#include "swp_model_factory.h"

#include <stdexcept>
#include <boost/random/poisson_distribution.hpp>

using namespace Sweep;
using namespace Sweep::AggModels;

// CONSTRUCTORS AND DESTRUCTORS.

// Default constructor (protected).
SurfVolCubicPrimary::SurfVolCubicPrimary(void)
{
}

// Initialising constructor.
SurfVolCubicPrimary::SurfVolCubicPrimary(double time, const Sweep::ParticleModel &model)
: Primary(time, model)
{
}

// Copy constructor.
SurfVolCubicPrimary::SurfVolCubicPrimary(const SurfVolCubicPrimary &copy)
{
    *this = copy;
}

// Stream-reading constructor.
SurfVolCubicPrimary::SurfVolCubicPrimary(std::istream &in, const Sweep::ParticleModel &model)
{
    Deserialize(in, model);
}

// Default destructor.
SurfVolCubicPrimary::~SurfVolCubicPrimary()
{
    releaseMem();
}

// OPERATOR OVERLOADS.

// Assignment operator (Primary RHS).
SurfVolCubicPrimary &SurfVolCubicPrimary::operator=(const Primary &rhs)
{
    operator=(dynamic_cast<const SurfVolCubicPrimary&>(rhs));
    return *this;
}

// AGGREGATION MODEL.

// Returns the aggregation model which this primary describes.
AggModels::AggModelType SurfVolCubicPrimary::AggID(void) const {return AggModels::SurfVolCubic_ID;}

// BASIC DERIVED PARTICLE PROPERTIES.

// Calculates the derived properties from the unique properties.  This
// function is broadly similar to the version in the spherical Primary
// class except that it uses the surface-volume model.  Therefore the
// surface area is not altered and the collision diameter is calculated
// using the arithmetic mean function.
//
// Mobility diameter:
// This calculation is based on the work of Rogak et al., 1993 Aer. Sci.
// Tech. 18:25-47, who give the calculation of dmob in the FM, SF and
// transition regime.
void SurfVolCubicPrimary::UpdateCache(void)
{
    // Store the old surface area, this value may be 0 if the particle is in
    // the process of initialisation
    double s = m_surf;

    // Pretend that the primary is spherical and set the cache
    // accordingly.
    Primary::UpdateCache();

    m_surf = std::max(s, m_surf);

    // Calculate diameters.
    m_dcol = (m_diam + sqrt(m_surf / PI)) * 0.5;

    // Calculate mobility diameter
    m_dmob = PP_Diameter();
    if (false) {
        // SF regime mobility diameter
        m_dmob *= 0.9 * sqrt(m_pmodel->GetFractDim() / (m_pmodel->GetFractDim() + 2));
        m_dmob *= pow(PP_Count(), (1.0/m_pmodel->GetFractDim()));
    } else {
        // FM regime mobility diameter
        m_dmob *= sqrt(0.802*(PP_Count()-1) + 1);
    }

    if (m_dmob < m_diam) m_dmob = m_diam;
}

// Returns the number of primary particles if the aggregate is assumed
// to consist of mono-sized primaries.
unsigned int SurfVolCubicPrimary::PP_Count(void) const
{
    // Note the minimum number of primary particles must be 1.
    return std::max(1u, (unsigned int)((m_surf * m_surf * m_surf) /
                                       (36.0 * PI * m_vol * m_vol)));
}

// Returns the primary particle diameter if the aggregate is assumed
// to consist of mono-sized primaries.
double SurfVolCubicPrimary::PP_Diameter(void) const
{
    // This should always be <= equiv. sphere diameter.
    return 6.0 * m_vol / m_surf;
}

// OPERATIONS.

// Adjusts the primary with the given composition and
// tracker values changes n times.  If the particle cannot be adjust
// n times, then this function returns the number of times
// it was adjusted.
unsigned int SurfVolCubicPrimary::Adjust(const fvector &dcomp, const fvector &dvalues, rng_type &rng,
                                    unsigned int n)
{
    // Store initial surface and volume
    double surfOld = m_surf;
    double volOld  = m_vol;

    // Adjust the particle assuming that it is spherical.
    n = Primary::Adjust(dcomp, dvalues, rng, n);

    // Calculate change in volume.
    double dvol = 0.0;
    for (unsigned int i=0; i!=dcomp.size(); ++i) {
        dvol += dcomp[i] * m_pmodel->Components(i)->MolWt() /
                m_pmodel->Components(i)->Density();
    }
    dvol *= (double)n / NA;

    // Calculate change in surface area.
    double invRadius = 0.0;
    if (dvol > 0.0) {
        // Inverse growth radius.
        invRadius = sqrt(4.0 * PI / surfOld);
    } else {
        // Inverse oxidation radius.
        invRadius = surfOld / (3.0 * volOld);
    }

    // Save new surface area.
    double s = surfOld + (2.0 * dvol * invRadius);

    // Set correct surface area, which was incorrectly set by
    // Primary::Adjust.
    m_surf    = s;

    // This has a knock-on affect of changing the collision diameter.
    // Note, we can avoid recalling UpdateCache() here, because only
    // a couple of cached values will have changed.
    m_dcol = (m_diam + sqrt(m_surf / PI)) * 0.5;
    m_dmob = m_dcol;

    return n;
}

/*!
 * Combines this primary with another.
 *
 * \param[in]       rhs         Particle to add to current instance
 * \param[in,out]   rng         Random number generator
 *
 * \return      Reference to the current instance after rhs has been added
 */
SurfVolCubicPrimary &SurfVolCubicPrimary::Coagulate(const Primary &rhs, rng_type &rng)

{
    // Store the resultant surface area.
    double s = m_surf + rhs.SurfaceArea();

    // Perform the coagulation.
    Primary::Coagulate(rhs, rng);

    // One sixth of the surface area is assumed to be lost as a result of two faces, one
    // from each incoming particle covering each other.  More sophisticated formulae
    // could be implemented here.
    m_surf = s * 0.8333333333333333333;

    // This has a knock-on affect of changing the collision diameter.
    // Note, we can avoid recalling UpdateCache() here, because only
    // a couple of cached values will have changed.
    m_dcol = (m_diam + sqrt(m_surf / PI)) * 0.5;
    m_dmob = m_dcol;

    return *this;
}

/*!
 * Combines this primary with another.
 *
 * \param[in]       rhs         Particle to add to current instance
 * \param[in,out]   rng         Random number generator
 *
 * \return      Reference to the current instance after rhs has been added
 */
SurfVolCubicPrimary &SurfVolCubicPrimary::Fragment(const Primary &rhs, rng_type &rng)

{
    // Store the resultant surface area.
    double s = m_surf + rhs.SurfaceArea();

    // Perform the coagulation.
    Primary::Fragment(rhs, rng);

    // One sixth of the surface area is assumed to be lost as a result of two faces, one
    // from each incoming particle covering each other.  More sophisticated formulae
    // could be implemented here.
    m_surf = s * 0.8333333333333333333;

    // This has a knock-on affect of changing the collision diameter.
    // Note, we can avoid recalling UpdateCache() here, because only
    // a couple of cached values will have changed.
    m_dcol = (m_diam + sqrt(m_surf / PI)) * 0.5;
    m_dmob = m_dcol;

    return *this;
}

// READ/WRITE/COPY.

//! Returns a copy of the model data.
SurfVolCubicPrimary *const SurfVolCubicPrimary::Clone(void) const
{
    return new SurfVolCubicPrimary(*this);
}

// Writes the object to a binary stream.
void SurfVolCubicPrimary::Serialize(std::ostream &out) const
{
    if (out.good()) {
        // Output base class.
        Primary::Serialize(out);
    } else {
        throw std::invalid_argument("Output stream not ready "
                                    "(Sweep, SurfVolCubicPrimary::Serialize).");
    }
}

// Reads the object from a binary stream.
void SurfVolCubicPrimary::Deserialize(std::istream &in, const Sweep::ParticleModel &model)
{
    if (in.good()) {
        // Read base class.
        Primary::Deserialize(in, model);
    } else {
        throw std::invalid_argument("Input stream not ready "
                                    "(Sweep, SurfVolCubicPrimary::Deserialize).");
    }
}


#include <sstream>
#include <limits>
#include <cstdlib> /* srand, rand */
#include <cmath>

#include <boost/math/distributions/negative_binomial.hpp>
#include <boost/log/trivial.hpp>

#include "kmergraphwithcoverage.h"
#include "localPRG.h"

using namespace prg;

void KmerGraphWithCoverage::set_exp_depth_covg(const uint32_t edp)
{
    const bool exp_depth_covg_parameter_is_valid = edp > 0;
    if (!exp_depth_covg_parameter_is_valid) {
        fatal_error(
            "Error setting exp_depth_covg: exp_depth_covg is invalid, must be > 0, is ",
            edp);
    }
    exp_depth_covg = edp;
}

void KmerGraphWithCoverage::set_binomial_parameter_p(const float e_rate)
{
    BOOST_LOG_TRIVIAL(debug) << "Set p in kmergraph";

    const bool valid_parameters_to_set_p
        = (kmer_prg->k != 0) && (0 < e_rate and e_rate < 1);
    if (!valid_parameters_to_set_p) {
        fatal_error("Error setting binomial parameter p, invalid parameters: ",
            "kmer_prg->k = ", kmer_prg->k, ", e_rate = ", e_rate);
    }

    binomial_parameter_p = 1 / exp(e_rate * kmer_prg->k);
}

void KmerGraphWithCoverage::increment_covg(
    uint32_t node_id, pandora::Strand strand, uint32_t sample_id)
{
    const bool sample_is_valid
        = this->node_index_to_sample_coverage[node_id].size() > sample_id;
    if (!sample_is_valid) {
        fatal_error(
            "Error incrementing coverage: sample_id is invalid (", sample_id, ")");
    }

    // get a pointer to the value we want to increment
    uint16_t* coverage_ptr = nullptr;
    if (strand == pandora::Strand::Forward) {
        coverage_ptr = &(this->node_index_to_sample_coverage[node_id][sample_id].first);
    } else {
        coverage_ptr
            = &(this->node_index_to_sample_coverage[node_id][sample_id].second);
    }

    const bool safe_to_increase_covg { (*coverage_ptr) < UINT16_MAX };
    if (safe_to_increase_covg) {
        ++(*coverage_ptr);
    }
}

uint32_t KmerGraphWithCoverage::get_covg(
    uint32_t node_id, pandora::Strand strand, uint32_t sample_id) const
{

    if (this->node_index_to_sample_coverage[node_id].size() <= sample_id)
        return 0;

    if (strand == pandora::Strand::Forward) {
        return (uint32_t)(
            this->node_index_to_sample_coverage[node_id][sample_id].first);
    } else {
        return (uint32_t)(
            this->node_index_to_sample_coverage[node_id][sample_id].second);
    }
}

void KmerGraphWithCoverage::set_covg(
    uint32_t node_id, uint16_t value, pandora::Strand strand, uint32_t sample_id)
{
    const bool sample_is_valid
        = this->node_index_to_sample_coverage[node_id].size() > sample_id;
    if (!sample_is_valid) {
        fatal_error("Error setting coverage: sample_id is invalid (", sample_id, ")");
    }

    if (strand == pandora::Strand::Forward) {
        this->node_index_to_sample_coverage[node_id][sample_id].first = value;
    } else {
        this->node_index_to_sample_coverage[node_id][sample_id].second = value;
    }
}

void KmerGraphWithCoverage::set_negative_binomial_parameters(
    const float& nbin_prob, const float& nb_fail)
{
    if (nbin_prob == 0 and nb_fail == 0)
        return;

    const bool negative_binomial_parameters_were_previously_set
        = (negative_binomial_parameter_p > 0 and negative_binomial_parameter_p < 1)
        && (negative_binomial_parameter_r > 0);
    if (!(negative_binomial_parameters_were_previously_set)) {
        fatal_error("Error setting negative_binomial_parameters: "
                    "negative_binomial_parameter_p (",
            negative_binomial_parameter_p, ")", " or negative_binomial_parameter_r (",
            negative_binomial_parameter_r, ") ", "were not correctly set");
    }

    negative_binomial_parameter_p += nbin_prob;
    negative_binomial_parameter_r += nb_fail;
}

float KmerGraphWithCoverage::nbin_prob(uint32_t node_id, const uint32_t& sample_id)
{
    auto k = this->get_forward_covg(node_id, sample_id)
        + this->get_reverse_covg(node_id, sample_id);
    float return_prob
        = log(pdf(boost::math::negative_binomial(
                      negative_binomial_parameter_r, negative_binomial_parameter_p),
            k));
    return_prob = std::max(return_prob, std::numeric_limits<float>::lowest() / 1000);
    return return_prob;
}

float KmerGraphWithCoverage::lin_prob(uint32_t node_id, const uint32_t& sample_id)
{
    const bool reads_were_mapped_to_this_kmer_graph = num_reads != 0;
    if (!reads_were_mapped_to_this_kmer_graph) {
        fatal_error(
            "Impossible to compute lin_prob, no reads were mapped to this kmer graph");
    }
    auto k = this->get_forward_covg(node_id, sample_id)
        + this->get_reverse_covg(node_id, sample_id);
    return log(float(k) / num_reads);
}

float KmerGraphWithCoverage::bin_prob(uint32_t node_id, const uint32_t& sample_id)
{
    const bool reads_were_mapped_to_this_kmer_graph = num_reads != 0;
    if (!reads_were_mapped_to_this_kmer_graph) {
        fatal_error(
            "Impossible to compute bin_prob, no reads were mapped to this kmer graph");
    }
    return bin_prob(node_id, num_reads, sample_id);
}

float KmerGraphWithCoverage::bin_prob(
    const uint32_t& node_id, const uint32_t& num, const uint32_t& sample_id)
{
    const bool binomial_parameter_p_is_set_correctly = binomial_parameter_p != 1;
    if (!binomial_parameter_p_is_set_correctly) {
        fatal_error("Error when computing bin_prob: binomial_parameter_p (",
            binomial_parameter_p, ") is not correctly set");
    }

    const bool node_exists = node_id < kmer_prg->nodes.size();
    if (!node_exists) {
        fatal_error("Error when computing bin_prob: attempt to access inexistent node ",
            node_id);
    }

    uint32_t sum_coverages = this->get_forward_covg(node_id, sample_id)
        + this->get_reverse_covg(node_id, sample_id);

    float return_prob;
    if (node_id == (*(kmer_prg->sorted_nodes.begin()))->id
        or node_id == (*(kmer_prg->sorted_nodes.rbegin()))->id) {
        return_prob = 0; // is really undefined
    } else if (sum_coverages > num) {
        // under model assumptions this can't happen, but it inevitably will, so bodge
        return_prob
            = lognchoosek2(sum_coverages, this->get_forward_covg(node_id, sample_id),
                  this->get_reverse_covg(node_id, sample_id))
            + sum_coverages * log(binomial_parameter_p / 2);
        // note this may give disadvantage to repeat kmers
    } else {
        return_prob = lognchoosek2(num, this->get_forward_covg(node_id, sample_id),
                          this->get_reverse_covg(node_id, sample_id))
            + sum_coverages * log(binomial_parameter_p / 2)
            + (num - sum_coverages) * log(1 - binomial_parameter_p);
    }
    return return_prob;
}

float KmerGraphWithCoverage::get_prob(
    const std::string& prob_model, const uint32_t& node_id, const uint32_t& sample_id)
{
    if (prob_model == "nbin") {
        // is there no parameter check here?
        return nbin_prob(node_id, sample_id);
    } else if (prob_model == "bin") {
        const bool binomial_parameters_are_ok
            = (binomial_parameter_p < 1) && (num_reads > 0);
        if (!binomial_parameters_are_ok) {
            fatal_error("Error when computing kmer prob: binomial parameters are not "
                        "ok (binomial_parameter_p = ",
                binomial_parameter_p, ", ", "num_reads = ", num_reads);
        }
        return bin_prob(node_id, sample_id);
    } else if (prob_model == "lin") {
        // is there no parameter check here?
        return lin_prob(node_id, sample_id);
    } else {
        fatal_error("Invalid probability model for kmer coverage distribution: ",
            prob_model, ". Should be nbin, bin or lin");
    }
}

bool KmerGraphWithCoverage::coverage_is_zeroes(const uint32_t& sample_id)
{
    bool all_zero = true;
    for (const auto& node_ptr : kmer_prg->nodes) {
        const auto covg { this->get_forward_covg(node_ptr->id, sample_id)
            + this->get_reverse_covg(node_ptr->id, sample_id) };
        if (covg > 0) {
            BOOST_LOG_TRIVIAL(debug) << "Found non-zero coverage in kmer graph";
            all_zero = false;
            break;
        }
    }
    if (all_zero) {
        BOOST_LOG_TRIVIAL(debug) << "ALL ZEROES in kmer graph coverages";
    }
    return all_zero;
}

float KmerGraphWithCoverage::find_max_path(std::vector<KmerNodePtr>& maxpath,
    const std::string& prob_model, const uint32_t& max_num_kmers_to_average,
    const uint32_t& sample_id)
{
    // TODO: FIX THIS INNEFICIENCY I INTRODUCED
    const std::vector<KmerNodePtr> sorted_nodes(
        this->kmer_prg->sorted_nodes.begin(), this->kmer_prg->sorted_nodes.end());
    this->kmer_prg->check();

    // also check not all 0 covgs
    auto coverages_all_zero = coverage_is_zeroes(sample_id);
    if (coverages_all_zero)
        return std::numeric_limits<float>::lowest();

    // create vectors to hold the intermediate values
    std::vector<float> max_sum_of_log_probs_from_node(sorted_nodes.size(), 0);
    std::vector<uint32_t> length_of_maxpath_from_node(sorted_nodes.size(), 0);
    std::vector<uint32_t> prev_node_along_maxpath(
        sorted_nodes.size(), sorted_nodes.size() - 1);
    float max_mean;
    int max_length;
    const float tolerance = 0.000001;

    for (uint32_t j = sorted_nodes.size() - 1; j != 0; --j) {
        max_mean = std::numeric_limits<float>::lowest();
        max_length = 0; // tie break with longest kmer path
        const auto& current_node = sorted_nodes[j - 1];
        for (uint32_t i = 0; i != current_node->out_nodes.size(); ++i) {
            const auto& considered_outnode = current_node->out_nodes[i].lock();
            const bool is_terminus_and_most_likely
                = considered_outnode->id == sorted_nodes.back()->id
                and thresh > max_mean + tolerance;
            const bool avg_log_likelihood_is_most_likely
                = max_sum_of_log_probs_from_node[considered_outnode->id]
                    / length_of_maxpath_from_node[considered_outnode->id]
                > max_mean + tolerance;
            const bool avg_log_likelihood_is_close_to_most_likely = max_mean
                    - max_sum_of_log_probs_from_node[considered_outnode->id]
                        / length_of_maxpath_from_node[considered_outnode->id]
                <= tolerance;
            const bool is_longer_path
                = length_of_maxpath_from_node[considered_outnode->id]
                > (uint)max_length;

            if (is_terminus_and_most_likely or avg_log_likelihood_is_most_likely
                or (avg_log_likelihood_is_close_to_most_likely and is_longer_path)) {
                max_sum_of_log_probs_from_node[current_node->id]
                    = get_prob(prob_model, current_node->id, sample_id)
                    + max_sum_of_log_probs_from_node[considered_outnode->id];
                length_of_maxpath_from_node[current_node->id]
                    = 1 + length_of_maxpath_from_node[considered_outnode->id];
                prev_node_along_maxpath[current_node->id] = considered_outnode->id;

                if (length_of_maxpath_from_node[current_node->id]
                    > max_num_kmers_to_average) {
                    uint32_t prev_node = prev_node_along_maxpath[current_node->id];
                    for (uint step = 0; step < max_num_kmers_to_average; step++) {
                        prev_node = prev_node_along_maxpath[prev_node];
                    }
                    max_sum_of_log_probs_from_node[current_node->id]
                        -= get_prob(prob_model, sorted_nodes[prev_node]->id, sample_id);
                    length_of_maxpath_from_node[current_node->id] -= 1;

                    // this remains as an assert, as it is a code check
                    // Note: I think we might even be able to remove this
                    assert(length_of_maxpath_from_node[current_node->id]
                        == max_num_kmers_to_average);
                }

                if (considered_outnode->id != sorted_nodes.back()->id) {
                    max_mean = max_sum_of_log_probs_from_node[considered_outnode->id]
                        / length_of_maxpath_from_node[considered_outnode->id];
                    max_length = length_of_maxpath_from_node[considered_outnode->id];
                } else {
                    max_mean = thresh;
                }
            }
        }
    }

    // extract path
    uint32_t prev_node = prev_node_along_maxpath[sorted_nodes[0]->id];
    while (prev_node < sorted_nodes.size() - 1) {
        maxpath.push_back(this->kmer_prg->nodes[prev_node]);
        prev_node = prev_node_along_maxpath[prev_node];

        if (maxpath.size() > 1000000) {
            fatal_error("I think I've found an infinite loop - is "
                        "something wrong with this kmergraph?");
        }
    }

    const bool path_was_found_through_the_kmer_PRG = length_of_maxpath_from_node[0] > 0;
    if (!path_was_found_through_the_kmer_PRG) {
        fatal_error("Error when finding max path: found no path through kmer prg");
    }

    return prob_path(maxpath, sample_id, prob_model);
}

std::vector<std::vector<KmerNodePtr>> KmerGraphWithCoverage::get_random_paths(
    uint32_t num_paths)
{
    // find a random path through kmergraph picking ~uniformly from the outnodes at each
    // point
    std::vector<std::vector<KmerNodePtr>> rpaths;
    std::vector<KmerNodePtr> rpath;
    uint32_t i;

    time_t now;
    now = time(nullptr);
    srand((unsigned int)now);

    if (!kmer_prg->nodes.empty()) {
        for (uint32_t j = 0; j != num_paths; ++j) {
            i = rand() % kmer_prg->nodes[0]->out_nodes.size();
            rpath.push_back(kmer_prg->nodes[0]->out_nodes[i].lock());
            while (rpath.back() != kmer_prg->nodes[kmer_prg->nodes.size() - 1]) {
                if (rpath.back()->out_nodes.size() == 1) {
                    rpath.push_back(rpath.back()->out_nodes[0].lock());
                } else {
                    i = rand() % rpath.back()->out_nodes.size();
                    rpath.push_back(rpath.back()->out_nodes[i].lock());
                }
            }
            rpath.pop_back();
            rpaths.push_back(rpath);
            rpath.clear();
        }
    }
    return rpaths;
}

float KmerGraphWithCoverage::prob_path(const std::vector<KmerNodePtr>& kpath,
    const uint32_t& sample_id, const std::string& prob_model)
{
    float return_prob_path = 0;
    for (uint32_t i = 0; i != kpath.size(); ++i) {
        return_prob_path += get_prob(prob_model, kpath[i]->id, sample_id);
    }
    uint32_t len = kpath.size();
    if (kpath[0]->path.length() == 0) {
        len -= 1;
    }
    if (kpath.back()->path.length() == 0) {
        len -= 1;
    }
    if (len == 0) {
        len = 1;
    }
    return return_prob_path / len;
}

void KmerGraphWithCoverage::save_covg_dist(const std::string& filepath)
{
    std::ofstream handle;
    handle.open(filepath);

    for (const auto& kmer_node_ptr : kmer_prg->nodes) {
        const KmerNode& kmer_node = *kmer_node_ptr;

        uint32_t sample_id = 0;
        for (const auto& sample_coverage :
            node_index_to_sample_coverage[kmer_node.id]) {
            handle << kmer_node.id << " " << sample_id << " " << sample_coverage.first
                   << " " << sample_coverage.second;

            sample_id++;
        }
    }
    handle.close();
}

// save the KmerGraph as gfa
// TODO: THIS SHOULD BE RECODED, WE ARE DUPLICATING CODE HERE (SEE KmerGraph::save())!!!
void KmerGraphWithCoverage::save(
    const fs::path& filepath, const std::shared_ptr<LocalPRG> localprg) const
{
    uint32_t sample_id = 0;

    fs::ofstream handle(filepath);
    if (handle.is_open()) {
        handle << "H\tVN:Z:1.0\tbn:Z:--linear --singlearr" << std::endl;
        for (const auto& c : kmer_prg->nodes) {
            handle << "S\t" << c->id << "\t";

            if (localprg != nullptr) {
                handle << localprg->string_along_path(c->path);
            } else {
                handle << c->path;
            }

            handle << "\tFC:i:" << this->get_forward_covg(c->id, sample_id)
                   << "\tRC:i:" << this->get_reverse_covg(c->id, sample_id)
                   << std::endl;

            for (uint32_t j = 0; j < c->out_nodes.size(); ++j) {
                handle << "L\t" << c->id << "\t+\t" << c->out_nodes[j].lock()->id
                       << "\t+\t0M" << std::endl;
            }
        }
        handle.close();
    } else {
        fatal_error("Unable to open kmergraph file ", filepath);
    }
}

// TODO: THIS SHOULD BE RECODED, WE ARE DUPLICATING CODE HERE (SEE KmerGraph::load())!!!
// TODO: remove this method?
void KmerGraphWithCoverage::load(const std::string& filepath)
{
    // TODO: this might be dangerous, recode this?
    auto kmer_prg = const_cast<KmerGraph*>(this->kmer_prg);
    kmer_prg->clear();
    uint32_t sample_id = 0;

    std::string line;
    std::vector<std::string> split_line;
    std::stringstream ss;
    uint32_t id = 0, from, to;
    uint16_t covg;
    prg::Path p;
    uint32_t num_nodes = 0;

    std::ifstream myfile(filepath);
    if (myfile.is_open()) {

        while (getline(myfile, line).good()) {
            if (line[0] == 'S') {
                split_line = split(line, "\t");

                const bool line_is_consistent = split_line.size() >= 4;
                if (!line_is_consistent) {
                    fatal_error("Error reading GFA. Offending line: ", line);
                }

                id = std::stoi(split_line[1]);
                num_nodes = std::max(num_nodes, id);
            }
        }
        myfile.clear();
        myfile.seekg(0, myfile.beg);
        kmer_prg->nodes.reserve(num_nodes);
        std::vector<uint16_t> outnode_counts(num_nodes + 1, 0),
            innode_counts(num_nodes + 1, 0);

        while (getline(myfile, line).good()) {
            if (line[0] == 'S') {
                split_line = split(line, "\t");

                const bool line_is_consistent = split_line.size() >= 4;
                if (!line_is_consistent) {
                    fatal_error("Error reading GFA. Offending line: ", line);
                }

                id = stoi(split_line[1]);
                ss << split_line[2];
                char c = ss.peek();

                if (!isdigit(c)) {
                    fatal_error("Error reading GFA: cannot read in this sort of "
                                "kmergraph GFA as it ",
                        "does not label nodes with their PRG path. ",
                        "Offending line: ", line);
                }

                ss >> p;
                ss.clear();
                // add_node(p);
                KmerNodePtr n = std::make_shared<KmerNode>(id, p);

                const bool id_is_consistent = (id == kmer_prg->nodes.size()
                    or num_nodes - id == kmer_prg->nodes.size());
                if (!id_is_consistent) {
                    fatal_error("Error reading GFA: node ID is inconsistent.",
                        "id = ", id, ", ", "nodes.size() = ", kmer_prg->nodes.size(),
                        ", ", "num_nodes = ", num_nodes);
                }

                kmer_prg->nodes.push_back(n);
                kmer_prg->sorted_nodes.insert(n);
                if (kmer_prg->k == 0 and p.length() > 0) {
                    kmer_prg->k = p.length();
                }
                covg = (uint16_t)(stoul(split(split_line[3], "FC:i:")[0]));
                this->set_forward_covg(n->id, covg, sample_id);
                covg = (uint16_t)(stoul(split(split_line[4], "RC:i:")[0]));
                this->set_reverse_covg(n->id, covg, sample_id);
                if (split_line.size() >= 6) {
                    n->num_AT = std::stoi(split_line[5]);
                }
            } else if (line[0] == 'L') {
                split_line = split(line, "\t");

                const bool line_is_consistent = split_line.size() >= 5;
                if (!line_is_consistent) {
                    fatal_error("Error reading GFA. Offending line: ", line);
                }

                const int from_node = stoi(split_line[1]);
                const int to_node = stoi(split_line[3]);
                const bool from_node_in_range = from_node < (int)outnode_counts.size();
                const bool to_node_in_range = to_node < (int)innode_counts.size();
                if (!from_node_in_range) {
                    fatal_error(
                        "Error reading GFA: from_node out of range: ", from_node,
                        ">=", outnode_counts.size(), ". Offending line: ", line);
                }
                if (!to_node_in_range) {
                    fatal_error("Error reading GFA: to_node out of range: ", to_node,
                        ">=", innode_counts.size(), ". Offending line: ", line);
                }

                outnode_counts[stoi(split_line[1])] += 1;
                innode_counts[stoi(split_line[3])] += 1;
            }
        }

        if (id == 0) {
            reverse(kmer_prg->nodes.begin(), kmer_prg->nodes.end());
        }

        id = 0;
        for (const auto& n : kmer_prg->nodes) {
            const bool id_is_consistent = (kmer_prg->nodes[id]->id == id)
                && (n->id < outnode_counts.size()) && (n->id < innode_counts.size());
            if (!id_is_consistent) {
                fatal_error("Error reading GFA: node: ", n,
                    " has inconsistent id, should be ", id);
            }
            id++;
            n->out_nodes.reserve(outnode_counts[n->id]);
            n->in_nodes.reserve(innode_counts[n->id]);
        }

        myfile.clear();
        myfile.seekg(0, myfile.beg);

        while (getline(myfile, line).good()) {
            if (line[0] == 'L') {
                split_line = split(line, "\t");

                const bool line_is_consistent = split_line.size() >= 5;
                if (!line_is_consistent) {
                    fatal_error("Error reading GFA. Offending line: ", line);
                }

                if (split_line[2] == split_line[4]) {
                    from = std::stoi(split_line[1]);
                    to = std::stoi(split_line[3]);
                } else {
                    // never happens
                    from = std::stoi(split_line[3]);
                    to = std::stoi(split_line[1]);
                }
                kmer_prg->add_edge(kmer_prg->nodes[from], kmer_prg->nodes[to]);
            }
        }
    } else {
        fatal_error("Error reading GFA: unable to open kmergraph file: ", filepath);
    }
}


///////////////////////////////////////////////////////////////////////////////
// tail_variate_means.hpp
//
//  Copyright 2006 Daniel Egloff, Olivier Gygi. Distributed under the Boost
//  Software License, Version 1.0. (See accompanying file
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)

#ifndef BOOST_ACCUMULATORS_STATISTICS_TAIL_VARIATE_MEANS_HPP_DE_01_01_2006
#define BOOST_ACCUMULATORS_STATISTICS_TAIL_VARIATE_MEANS_HPP_DE_01_01_2006

#include <numeric>
#include <vector>
#include <limits>
#include <functional>
#include <sstream>
#include <stdexcept>
#include <boost/throw_exception.hpp>
#include <boost/parameter/keyword.hpp>
#include <boost/mpl/placeholders.hpp>
#include <boost/type_traits/is_same.hpp>
#include <boost/accumulators/framework/accumulator_base.hpp>
#include <boost/accumulators/framework/extractor.hpp>
#include <boost/accumulators/numeric/functional.hpp>
#include <boost/accumulators/framework/parameters/sample.hpp>
#include <boost/accumulators/statistics_fwd.hpp>
#include <boost/accumulators/statistics/tail.hpp>
#include <boost/accumulators/statistics/tail_variate.hpp>
#include <boost/accumulators/statistics/tail_mean.hpp>
#include <boost/accumulators/statistics/parameters/quantile_probability.hpp>
#include <boost/serialization/vector.hpp>

#ifdef _MSC_VER
# pragma warning(push)
# pragma warning(disable: 4127) // conditional expression is constant
#endif

namespace boost { namespace accumulators
{

namespace impl
{
    /**
        @brief Estimation of the absolute and relative tail variate means (for both left and right tails)

        For all \f$j\f$-th variates associated to the \f$\lceil n(1-\alpha)\rceil\f$ largest samples (or the
        \f$\lceil n(1-\alpha)\rceil\f$ smallest samples in case of the left tail), the absolute tail means
        \f$\widehat{ATM}_{n,\alpha}(X, j)\f$ are computed and returned as an iterator range. Alternatively,
        the relative tail means \f$\widehat{RTM}_{n,\alpha}(X, j)\f$ are returned, which are the absolute
        tail means normalized with the (non-coherent) sample tail mean \f$\widehat{NCTM}_{n,\alpha}(X)\f$.

        \f[
            \widehat{ATM}_{n,\alpha}^{\mathrm{right}}(X, j) =
                \frac{1}{\lceil n(1-\alpha) \rceil}
                \sum_{i=\lceil \alpha n \rceil}^n \xi_{j,i}
        \f]

        \f[
            \widehat{ATM}_{n,\alpha}^{\mathrm{left}}(X, j) =
                \frac{1}{\lceil n\alpha \rceil}
                \sum_{i=1}^{\lceil n\alpha \rceil} \xi_{j,i}
        \f]

        \f[
            \widehat{RTM}_{n,\alpha}^{\mathrm{right}}(X, j) =
                \frac{\sum_{i=\lceil n\alpha \rceil}^n \xi_{j,i}}
            {\lceil n(1-\alpha)\rceil\widehat{NCTM}_{n,\alpha}^{\mathrm{right}}(X)}
        \f]

        \f[
            \widehat{RTM}_{n,\alpha}^{\mathrm{left}}(X, j) =
                \frac{\sum_{i=1}^{\lceil n\alpha \rceil} \xi_{j,i}}
            {\lceil n\alpha\rceil\widehat{NCTM}_{n,\alpha}^{\mathrm{left}}(X)}
        \f]
    */

    ///////////////////////////////////////////////////////////////////////////////
    // tail_variate_means_impl
    //  by default: absolute tail_variate_means
    template<typename Sample, typename Impl, typename LeftRight, typename VariateTag>
    struct tail_variate_means_impl
      : accumulator_base
    {
        typedef typename numeric::functional::fdiv<Sample, std::size_t>::result_type float_type;
        typedef std::vector<float_type> array_type;
        // for boost::result_of
        typedef iterator_range<typename array_type::iterator> result_type;

        tail_variate_means_impl(dont_care) {}

        template<typename Args>
        result_type result(Args const &args) const
        {
            std::size_t cnt = count(args);

            std::size_t n = static_cast<std::size_t>(
                std::ceil(
                    cnt * ( ( is_same<LeftRight, left>::value ) ? args[quantile_probability] : 1. - args[quantile_probability] )
                )
            );

            std::size_t num_variates = tail_variate(args).begin()->size();

            this->tail_means_.clear();
            this->tail_means_.resize(num_variates, Sample(0));

            // If n is in a valid range, return result, otherwise return NaN or throw exception
            if (n < static_cast<std::size_t>(tail(args).size()))
            {
                this->tail_means_ = std::accumulate(
                    tail_variate(args).begin()
                  , tail_variate(args).begin() + n
                  , this->tail_means_
                  , numeric::plus
                );

                float_type factor = n * ( (is_same<Impl, relative>::value) ? non_coherent_tail_mean(args) : 1. );

                std::transform(
                    this->tail_means_.begin()
                  , this->tail_means_.end()
                  , this->tail_means_.begin()
#ifdef BOOST_NO_CXX98_BINDERS
                  , std::bind(std::divides<float_type>(), std::placeholders::_1, factor)
#else
                  , std::bind2nd(std::divides<float_type>(), factor)
#endif
                );
            }
            else
            {
                if (std::numeric_limits<float_type>::has_quiet_NaN)
                {
                    std::fill(
                        this->tail_means_.begin()
                      , this->tail_means_.end()
                      , std::numeric_limits<float_type>::quiet_NaN()
                    );
                }
                else
                {
                    std::ostringstream msg;
                    msg << "index n = " << n << " is not in valid range [0, " << tail(args).size() << ")";
                    boost::throw_exception(std::runtime_error(msg.str()));
                }
            }
            return make_iterator_range(this->tail_means_);
        }

        // make this accumulator serializeable
        template<class Archive>
        void serialize(Archive & ar, const unsigned int file_version)
        {
            ar & tail_means_;
        }

    private:

        mutable array_type tail_means_;

    };

} // namespace impl

///////////////////////////////////////////////////////////////////////////////
// tag::absolute_tail_variate_means
// tag::relative_tail_variate_means
//
namespace tag
{
    template<typename LeftRight, typename VariateType, typename VariateTag>
    struct absolute_tail_variate_means
      : depends_on<count, non_coherent_tail_mean<LeftRight>, tail_variate<VariateType, VariateTag, LeftRight> >
    {
        typedef accumulators::impl::tail_variate_means_impl<mpl::_1, absolute, LeftRight, VariateTag> impl;
    };
    template<typename LeftRight, typename VariateType, typename VariateTag>
    struct relative_tail_variate_means
      : depends_on<count, non_coherent_tail_mean<LeftRight>, tail_variate<VariateType, VariateTag, LeftRight> >
    {
        typedef accumulators::impl::tail_variate_means_impl<mpl::_1, relative, LeftRight, VariateTag> impl;
    };
    struct abstract_absolute_tail_variate_means
      : depends_on<>
    {
    };
    struct abstract_relative_tail_variate_means
      : depends_on<>
    {
    };
}

///////////////////////////////////////////////////////////////////////////////
// extract::tail_variate_means
// extract::relative_tail_variate_means
//
namespace extract
{
    extractor<tag::abstract_absolute_tail_variate_means> const tail_variate_means = {};
    extractor<tag::abstract_relative_tail_variate_means> const relative_tail_variate_means = {};

    BOOST_ACCUMULATORS_IGNORE_GLOBAL(tail_variate_means)
    BOOST_ACCUMULATORS_IGNORE_GLOBAL(relative_tail_variate_means)
}

using extract::tail_variate_means;
using extract::relative_tail_variate_means;

// tail_variate_means<LeftRight, VariateType, VariateTag>(absolute) -> absolute_tail_variate_means<LeftRight, VariateType, VariateTag>
template<typename LeftRight, typename VariateType, typename VariateTag>
struct as_feature<tag::tail_variate_means<LeftRight, VariateType, VariateTag>(absolute)>
{
    typedef tag::absolute_tail_variate_means<LeftRight, VariateType, VariateTag> type;
};

// tail_variate_means<LeftRight, VariateType, VariateTag>(relative) ->relative_tail_variate_means<LeftRight, VariateType, VariateTag>
template<typename LeftRight, typename VariateType, typename VariateTag>
struct as_feature<tag::tail_variate_means<LeftRight, VariateType, VariateTag>(relative)>
{
    typedef tag::relative_tail_variate_means<LeftRight, VariateType, VariateTag> type;
};

// Provides non-templatized extractor
template<typename LeftRight, typename VariateType, typename VariateTag>
struct feature_of<tag::absolute_tail_variate_means<LeftRight, VariateType, VariateTag> >
  : feature_of<tag::abstract_absolute_tail_variate_means>
{
};

// Provides non-templatized extractor
template<typename LeftRight, typename VariateType, typename VariateTag>
struct feature_of<tag::relative_tail_variate_means<LeftRight, VariateType, VariateTag> >
  : feature_of<tag::abstract_relative_tail_variate_means>
{
};

// So that absolute_tail_means can be automatically substituted
// with absolute_weighted_tail_means when the weight parameter is non-void.
template<typename LeftRight, typename VariateType, typename VariateTag>
struct as_weighted_feature<tag::absolute_tail_variate_means<LeftRight, VariateType, VariateTag> >
{
    typedef tag::absolute_weighted_tail_variate_means<LeftRight, VariateType, VariateTag> type;
};

template<typename LeftRight, typename VariateType, typename VariateTag>
struct feature_of<tag::absolute_weighted_tail_variate_means<LeftRight, VariateType, VariateTag> >
  : feature_of<tag::absolute_tail_variate_means<LeftRight, VariateType, VariateTag> >
{
};

// So that relative_tail_means can be automatically substituted
// with relative_weighted_tail_means when the weight parameter is non-void.
template<typename LeftRight, typename VariateType, typename VariateTag>
struct as_weighted_feature<tag::relative_tail_variate_means<LeftRight, VariateType, VariateTag> >
{
    typedef tag::relative_weighted_tail_variate_means<LeftRight, VariateType, VariateTag> type;
};

template<typename LeftRight, typename VariateType, typename VariateTag>
struct feature_of<tag::relative_weighted_tail_variate_means<LeftRight, VariateType, VariateTag> >
  : feature_of<tag::relative_tail_variate_means<LeftRight, VariateType, VariateTag> >
{
};

}} // namespace boost::accumulators

#ifdef _MSC_VER
# pragma warning(pop)
#endif

#endif


/*******************************************************************************

  Copyright (c) 2017, Honda Research Institute Europe GmbH

  Redistribution and use in source and binary forms, with or without
  modification, are permitted provided that the following conditions are
  met:

  1. Redistributions of source code must retain the above copyright notice,
   this list of conditions and the following disclaimer.

  2. Redistributions in binary form must reproduce the above copyright
   notice, this list of conditions and the following disclaimer in the
   documentation and/or other materials provided with the distribution.

  3. Neither the name of the copyright holder nor the names of its
   contributors may be used to endorse or promote products derived from
   this software without specific prior written permission.

  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
  IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
  THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
  PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
  CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
  PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
  PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
  LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
  NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

*******************************************************************************/

/******************************************************************************

  http://eigen.tuxfamily.org/dox/TopicPreprocessorDirectives.html:

  EIGEN_NO_MALLOC - if defined, any request from inside the Eigen to allocate
                    memory from the heap results in an assertion failure. This
                    is useful to check that some routine does not allocate
                    memory dynamically. Not defined by default.

  EIGEN_RUNTIME_NO_MALLOC - if defined, a new switch is introduced which can be
                            turned on and off by calling set_is_malloc_allowed
                            (bool). If malloc is not allowed and Eigen tries to
                            allocate memory dynamically anyway, an assertion
                            failure results. Not defined by default.

  Mapping C-Arrays:
   http://eigen.tuxfamily.org/dox/group__TutorialMapClass.html#
          TutorialMapPlacementNew

  Malloc checking:
   Eigen::internal::set_is_malloc_allowed(true, false) will create a fatal
   error if set to true, and disable checking for false. However, this is
   not thread-safe, and therefore should only be used for testing.

******************************************************************************/

#include <Rcs_eigen.h>
#include <Rcs_macros.h>

#define EIGEN_RUNTIME_NO_MALLOC

#include <Eigen/Core>
#include <Eigen/Dense>

#include <Eigen/QR>
#include <Eigen/SVD>
#include <Eigen/Cholesky>
#include <Eigen/LU>

#include <iostream>
#include <cstdio>

using namespace Eigen;

#if defined (MATND_MAX_STACK_VECTOR_SIZE)
#undef MATND_MAX_STACK_VECTOR_SIZE
#endif

#if defined (MATND_MAX_STACK_MATRIX_SIZE)
#undef MATND_MAX_STACK_MATRIX_SIZE
#endif

#define MATND_MAX_STACK_VECTOR_SIZE (25)
#define MATND_MAX_STACK_MATRIX_SIZE (25*25)

// These types have an unknown size at compile time, and therefore are
// allocated on the heap. The storage order of the matrices is set to row
// major, so that the Eigen's Map class can be used to cast them over a MatNd
// type without any additional memory allocation.
typedef Matrix<double, Dynamic, Dynamic, RowMajor> HeapMat;
typedef Matrix<double, Dynamic, 1, ColMajor> HeapVec;

// These types have a known size at compile time, and therefore are allocated
// on the stack. The storage order of the matrices is set to row major, so that
// the Eigen's Map class can be used to cast them over a MatNd type without any
// additional memory allocation.
typedef Matrix<double, Dynamic, Dynamic, RowMajor,
        MATND_MAX_STACK_VECTOR_SIZE,
        MATND_MAX_STACK_VECTOR_SIZE> StackMat;
typedef Matrix<double, Dynamic, 1, ColMajor,
        MATND_MAX_STACK_VECTOR_SIZE, 1> StackVec;

template <typename T>
static inline void MatNd_fromEigen3(MatNd* dst, const T& src,
                                    bool reshape=false)
{
  if (reshape==false)
  {
    RCHECK_EQ((int)dst->m, src.rows());
    RCHECK_EQ((int)dst->n, src.cols());
  }
  else
  {
    MatNd_reshape(dst, src.rows(), src.cols());
  }

  for (size_t m=0; m<dst->m; m++)
  {
    for (size_t n=0; n<dst->n; n++)
    {
      MatNd_set2(dst, m, n, src(m,n));
    }
  }
}

template <typename T>
static inline double CholDet(const T& L)
{
  double det = 1.0;

  for (int i=0; i<L.rows(); i++)
  {
    det *= L(i, i)*L(i, i);
  }

  return det;
}

template <typename MatT>
static inline double CholDC(MatNd* L_, const MatNd* A_)
{
  Map<MatT> A(A_->ele, A_->m, A_->n);
  LLT<MatT> llt;
  llt.compute(A);

  if (llt.info() != Eigen::Success)
  {
    RLOG(1, "Error on LLT!");
    MatNd_setZero(L_);
    return 0.0;
  }

  MatT L(llt.matrixL());
  double det = CholDet<MatT>(L);
  MatNd_fromEigen3<MatT>(L_, L);

  return det;
}

template <typename MatT>
static inline double CholSolve(MatNd* x_, const MatNd* A_, const MatNd* b_)
{
  Map<const MatT> A(A_->ele, A_->m, A_->n);
  LLT<MatT> llt;
  llt.compute(A);

  if (llt.info() != Eigen::Success)
  {
    RLOG(1, "Error on LLT!");
    return 0.0;
  }

  Map<MatT> x(x_->ele, x_->m, x_->n);
  Map<const MatT> b(b_->ele, b_->m, b_->n);
  x = llt.solve(b);

  // Compute determinant
  MatT L(llt.matrixL());

  return CholDet<MatT>(L);
}

template <typename MatT>
static inline double CholInv(MatNd* A_inv_, const MatNd* A_)
{
  size_t n = A_->m;
  Map<const MatT> A(A_->ele, n, n);

  LLT<MatT> llt;
  llt.compute(A);

  if (llt.info() != Eigen::Success)
  {
    RLOG(1, "Error on LLT!");
    MatNd_setZero(A_inv_);
    return 0.0;
  }

  Map<MatT> invA(A_inv_->ele, n, n);
  invA = llt.solve(MatT::Identity(n,n));

  MatT L(llt.matrixL());

  return CholDet<MatT>(L);
}

template <typename MatT>
static inline int SvdDC(MatNd* U, MatNd* S, MatNd* V, const MatNd* J,
                        double eps)
{
  Map<MatT> mapJ(J->ele, J->m, J->n);

  // We need to compute the full SVD, since otherwise there will dynamic
  // memory allocation for the StackMat types due to the QR preconditioning
  // (on Linux)
  JacobiSVD<MatT> svd(mapJ, ComputeFullU | ComputeFullV);

  // Cut off singular values smaller than eps
  int rank = svd.nonzeroSingularValues();

  for (int i=0; i<svd.nonzeroSingularValues(); ++i)
  {
    if (svd.singularValues()(i) < eps)
    {
      rank = i;
      break;
    }
  }

  if (S != NULL)
  {
    MatNd_fromEigen3<MatT>(S, svd.singularValues(), true);
    S->m = rank;
  }

  if (U != NULL)
  {
    MatT svdU = svd.matrixU();
    MatNd_reshape(U, svdU.rows(), rank);

    for (unsigned int m=0; m<U->m; m++)
    {
      for (unsigned int n=0; n<U->n; n++)
      {
        MatNd_set2(U, m, n, svdU(m,n));
      }
    }
  }

  if (V != NULL)
  {
    MatT svdV = svd.matrixV();
    MatNd_reshape(V, svdV.rows(), rank);

    for (unsigned int m=0; m<V->m; m++)
    {
      for (unsigned int n=0; n<V->n; n++)
      {
        MatNd_set2(V, m, n, svdV(m,n));
      }
    }
  }

  return rank;
}

template <typename MatT>
static inline double SvdSolve(MatNd* x_, const MatNd* A_, const MatNd* b_)
{
  Map<MatT> x(x_->ele, x_->m, x_->n);
  Map<MatT> A(A_->ele, A_->m, A_->n);
  Map<MatT> b(b_->ele, b_->m, b_->n);

  // We need to compute the full SVD, since otherwise there will dynamic
  // memory allocation for the StackMat types due to the QR preconditioning
  // (on Linux)
  JacobiSVD<MatT> svd(A, ComputeFullU | ComputeFullV);
  x = svd.solve(b);

  // Compute determinant
  double det = 1.0;
  MatT S = svd.singularValues();

  // Cut off singular values smaller than eps
  for (int i=0; i<S.rows(); ++i)
  {
    det *= S(i);
  }

  return det;
}

template <typename MatT>
static inline size_t SvdInverse(MatNd* A_, double eps=1.0e-12)
{
  Map<MatT> A(A_->ele, A_->m, A_->n);

  // We need to compute the full SVD, since otherwise there will dynamic
  // memory allocation for the StackMat types due to the QR preconditioning
  // (on Linux)
  JacobiSVD<MatT> svd;
  //  svd.setThreshold(eps);
  svd.compute(A, ComputeThinU | ComputeThinV);

  MatT invS = svd.singularValues();

  // Cut off singular values smaller than eps
  size_t rank = 0;
  for (int i=0; i<invS.rows(); ++i)
  {
    if (invS(i) > eps)
    {
      invS(i) = 1.0/invS(i);
      rank++;
    }
  }

  // inv = V diag(S^-1) U^T
  A.noalias() = svd.matrixV()*invS.asDiagonal()*svd.matrixU().transpose();

  return rank;
}

template <typename MatT>
static inline void QRDC(MatNd* Q_, MatNd* R_, const MatNd* A_)
{
  Map<const MatT> A(A_->ele, A_->m, A_->n);
  Map<MatT> Q(Q_->ele, Q_->m, Q_->n);
  Map<MatT> R(R_->ele, R_->m, R_->n);

  HouseholderQR<MatT> qr(A);

  Q = qr.householderQ();
  R = qr.matrixQR().template triangularView<Eigen::Upper>();
}

template <typename MatT>
static inline bool EigenVectors(MatNd* V_, double* d_, const MatNd* A_)
{
  Map<const MatT> A(A_->ele, A_->m, A_->n);

  Eigen::SelfAdjointEigenSolver<MatT> eigensolver(A);

  if (eigensolver.info() != Eigen::Success)
  {
    return false;
  }

  Map<MatT> V(V_->ele, V_->m, V_->n);
  Map<MatT> d(d_, A_->m, 1);
  V = eigensolver.eigenvectors();
  d = eigensolver.eigenvalues();

  return true;
}

template <typename MatT>
static inline void ColPivHhQR(MatNd* X_, const MatNd* A_, const MatNd* B_)
{
  Map<const MatT> A(A_->ele, A_->m, A_->n);
  Map<const MatT> B(B_->ele, B_->m, B_->n);
  Map<MatT> X(X_->ele, X_->m, X_->n);

  // A X = B
  X = A.colPivHouseholderQr().solve(B);
}

#ifdef __cplusplus
extern "C" {
#endif

double MatNd_choleskyDecomposition_E3(MatNd* L, const MatNd* A)
{
  RCHECK_EQ(A->m, A->n);
  RCHECK_EQ(L->m, L->n);
  RCHECK_EQ(A->m, L->m);

  if (A->m > MATND_MAX_STACK_VECTOR_SIZE)
  {
    return CholDC<HeapMat>(L, A);
  }
  else
  {
#ifdef EIGEN_RUNTIME_NO_MALLOC
    Eigen::internal::set_is_malloc_allowed(false);
#endif
    double det = CholDC<StackMat>(L, A);
#ifdef EIGEN_RUNTIME_NO_MALLOC
    Eigen::internal::set_is_malloc_allowed(true);
#endif
    return det;
  }
}

double MatNd_choleskySolve_E3(MatNd* x, const MatNd* A, const MatNd* b)
{
  RCHECK_EQ(A->m, A->n);
  RCHECK_EQ(x->m, A->m);
  RCHECK_EQ(b->m, x->m);
  RCHECK_EQ(b->n, x->n);

  if (A->m > MATND_MAX_STACK_VECTOR_SIZE)
  {
    return CholSolve<HeapMat>(x, A, b);
  }
  else
  {
#ifdef EIGEN_RUNTIME_NO_MALLOC
    Eigen::internal::set_is_malloc_allowed(false);
#endif
    double det = CholSolve<StackMat>(x, A, b);
#ifdef EIGEN_RUNTIME_NO_MALLOC
    Eigen::internal::set_is_malloc_allowed(true);
#endif
    return det;
  }
}

double MatNd_choleskyInverse_E3(MatNd* invA, const MatNd* A)
{
  RCHECK_EQ(A->m, A->n);
  RCHECK_EQ(invA->m, invA->n);
  RCHECK_EQ(A->m, invA->m);

  if (A->m > MATND_MAX_STACK_VECTOR_SIZE)
  {
    return CholInv<HeapMat>(invA,A);
  }
  else
  {
#ifdef EIGEN_RUNTIME_NO_MALLOC
    Eigen::internal::set_is_malloc_allowed(false);
#endif
    double det = CholInv<StackMat>(invA,A);
#ifdef EIGEN_RUNTIME_NO_MALLOC
    Eigen::internal::set_is_malloc_allowed(true);
#endif
    return det;
  }
}

void MatNd_mul_E3(MatNd* C_, const MatNd* A_, const MatNd* B_)
{
  if ((A_->m > MATND_MAX_STACK_VECTOR_SIZE) ||
      (B_->n > MATND_MAX_STACK_VECTOR_SIZE))
  {
    Map<HeapMat> C(C_->ele, C_->m, C_->n);
    Map<HeapMat> A(A_->ele, A_->m, A_->n);
    Map<HeapMat> B(B_->ele, B_->m, B_->n);

    // The noalias function inhibits the creation of a temporary object which
    // would lead to dynamic memory alllocation.
    C.noalias() = A*B;
  }
  else
  {
#ifdef EIGEN_RUNTIME_NO_MALLOC
    Eigen::internal::set_is_malloc_allowed(false);
#endif

    Map<StackMat> C(C_->ele, C_->m, C_->n);
    Map<StackMat> A(A_->ele, A_->m, A_->n);
    Map<StackMat> B(B_->ele, B_->m, B_->n);

    C.noalias() = A*B;

#ifdef EIGEN_RUNTIME_NO_MALLOC
    Eigen::internal::set_is_malloc_allowed(true);
#endif
  }
}

unsigned int MatNd_SVD(MatNd* U, MatNd* S, MatNd* V, const MatNd* J,
                       double eps)
{
  if ((J->m > MATND_MAX_STACK_VECTOR_SIZE) ||
      (J->n > MATND_MAX_STACK_VECTOR_SIZE))
  {
    return SvdDC<HeapMat>(U, S, V, J, eps);
  }
  else
  {
#ifdef EIGEN_RUNTIME_NO_MALLOC
    Eigen::internal::set_is_malloc_allowed(false);
#endif
    int rank = SvdDC<StackMat>(U, S, V, J, eps);
#ifdef EIGEN_RUNTIME_NO_MALLOC
    Eigen::internal::set_is_malloc_allowed(true);
#endif
    return rank;
  }
}

double MatNd_SVDSolve(MatNd* x, const MatNd* A, const MatNd* b)
{
  if ((A->m > MATND_MAX_STACK_VECTOR_SIZE) ||
      (A->n > MATND_MAX_STACK_VECTOR_SIZE))
  {
    return SvdSolve<HeapMat>(x, A, b);
  }
  else
  {
#ifdef EIGEN_RUNTIME_NO_MALLOC
    Eigen::internal::set_is_malloc_allowed(false);
#endif
    double det = SvdSolve<StackMat>(x, A, b);
#ifdef EIGEN_RUNTIME_NO_MALLOC
    Eigen::internal::set_is_malloc_allowed(true);
#endif
    return det;
  }
}

unsigned int MatNd_SVDInverse(MatNd* inv, const MatNd* src)
{
  MatNd_copy(inv, src);
  return MatNd_SVDInverseSelf(inv);
}

unsigned int MatNd_SVDInverseSelf(MatNd* A)
{
  if ((A->m > MATND_MAX_STACK_VECTOR_SIZE) ||
      (A->n > MATND_MAX_STACK_VECTOR_SIZE))
  {
    return SvdInverse<HeapMat>(A);
  }
  else
  {
#ifdef EIGEN_RUNTIME_NO_MALLOC
    Eigen::internal::set_is_malloc_allowed(false);
#endif
    int rank = SvdInverse<StackMat>(A);
#ifdef EIGEN_RUNTIME_NO_MALLOC
    Eigen::internal::set_is_malloc_allowed(true);
#endif
    return rank;
  }
}

unsigned int MatNd_rank(const MatNd* J, double eps)
{
  return MatNd_SVD(NULL, NULL, NULL, J, eps);
}

void MatNd_QRDecomposition(MatNd* Q, MatNd* R, const MatNd* A)
{
  if ((A->m > MATND_MAX_STACK_VECTOR_SIZE) ||
      (A->n > MATND_MAX_STACK_VECTOR_SIZE) ||
      (A->n > A->m))   // That's a bug in Eigen3
  {
    QRDC<HeapMat>(Q, R, A);
  }
  else
  {
#ifdef EIGEN_RUNTIME_NO_MALLOC
    Eigen::internal::set_is_malloc_allowed(false);
#endif
    QRDC<StackMat>(Q, R, A);
#ifdef EIGEN_RUNTIME_NO_MALLOC
    Eigen::internal::set_is_malloc_allowed(true);
#endif
  }
}

bool MatNd_getEigenVectors(MatNd* V, double* d, const MatNd* A)
{
  RCHECK_EQ(A->m, A->n);
  RCHECK_EQ(A->n, V->m);
  RCHECK_EQ(V->m, V->n);

  bool success;

  if (A->m > MATND_MAX_STACK_VECTOR_SIZE)
  {
    success = EigenVectors<HeapMat>(V, d, A);
  }
  else
  {
#ifdef EIGEN_RUNTIME_NO_MALLOC
    Eigen::internal::set_is_malloc_allowed(false);
#endif
    success = EigenVectors<StackMat>(V, d, A);
#ifdef EIGEN_RUNTIME_NO_MALLOC
    Eigen::internal::set_is_malloc_allowed(true);
#endif
  }

  return success;
}

void MatNd_HouseholderQR(MatNd* X, const MatNd* A, const MatNd* B)
{
  if ((A->m > MATND_MAX_STACK_VECTOR_SIZE) ||
      (A->n > MATND_MAX_STACK_VECTOR_SIZE))
  {
    return ColPivHhQR<HeapMat>(X, A, B);
  }
  else
  {
#ifdef EIGEN_RUNTIME_NO_MALLOC
    Eigen::internal::set_is_malloc_allowed(false);
#endif
    return ColPivHhQR<StackMat>(X, A, B);
#ifdef EIGEN_RUNTIME_NO_MALLOC
    Eigen::internal::set_is_malloc_allowed(true);
#endif
  }
}

#ifdef __cplusplus
}
#endif


/*
 * Copyright 2008-2016 Jan Gasthaus
 *
 * 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.
 */

#include "libplump/hpyp_model.h"

#include <cmath>
#include <boost/bind.hpp>
#include <boost/shared_ptr.hpp>

#include "libplump/utils.h"
#include "libplump/subseq.h"
#include "libplump/stirling.h"
#include "libplump/random.h"
#include "libplump/hpyp_restaurants.h"

namespace gatsby { namespace libplump {

HPYPModel::HPYPModel(seq_type& seq,
                     INodeManager& nodeManager,
                     const IAddRemoveRestaurant& restaurant,
                     IParameters& parameters,
                     int numTypes)
    : seq(seq),
      contextTree_(new ContextTree(nodeManager, seq)),
      contextTree(*contextTree_),
      restaurant(restaurant),
      parameters(parameters),
      numTypes(numTypes) {
    baseProb = 1./((double) numTypes);
  }

void HPYPModel::insertRoot(e_type obs) {
  WrappedNodeList root_path = this->contextTree.findLongestSuffix(0,0);
  d_vec discount_path = this->parameters.getDiscounts(root_path);
  d_vec concentration_path = this->parameters.getConcentrations(root_path,
                                                                discount_path);
  d_vec prob_path = this->computeProbabilityPath(root_path,
                                                 discount_path,
                                                 concentration_path,
                                                 obs);
  this->updatePath(root_path,
                   prob_path,
                   discount_path,
                   concentration_path,
                   obs);
}

/**
 * Insert a context into the tree and handle a potentially occuring
 * split.
 */
WrappedNodeList HPYPModel::insertContext(l_type start, l_type stop) {
  typedef ContextTree::InsertionResult InsertionResult;

  // insert context
  InsertionResult insertionResult = contextTree.insert(start, stop);

  // handle split if one occurred
  if (insertionResult.action != InsertionResult::INSERT_ACTION_NO_SPLIT) {
    WrappedNode nodeA, nodeC;
    WrappedNodeList::iterator i = insertionResult.path.end();

    // move iterator to point to the split node which is either the
    // last or second to last element on the path, depending on
    // whether the inserted node was a suffix
    switch (insertionResult.action) {
      case InsertionResult::INSERT_ACTION_SPLIT :
        // split where inserted node is not a suffix of a node in
        // the tree
        i--; i--; // split node is above the newly inserted node
        break;
      case InsertionResult::INSERT_ACTION_SPLIT_SUFFIX :
        // split where inserted node is a suffix -> inserted node
        // is the shorter node
        i--; // split node is the newly inserted node
        break;
      case InsertionResult::INSERT_ACTION_NO_SPLIT :
        break;
    }
    nodeC = *i;
    i--; // move iterator to parent
    nodeA = *i;
    this->handleSplit(nodeA,
                      insertionResult.splitChild,
                      nodeC);
  }
  return insertionResult.path;
}

d_vec HPYPModel::insertContextAndObservation(l_type start,
                                             l_type stop,
                                             e_type obs) {
  // insert context (and handle a potential split)
  WrappedNodeList path = insertContext(start, stop);

  // insert observation
  d_vec discountPath = this->parameters.getDiscounts(path);
  d_vec concentrationPath = this->parameters.getConcentrations(path,
                                                               discountPath);
  d_vec probabilityPath = this->computeProbabilityPath(path,
                                                       discountPath,
                                                       concentrationPath,
                                                       obs);
  this->parameters.accumulateParameterGradient(this->restaurant,
      path, probabilityPath, discountPath, concentrationPath, obs);

  this->updatePath(path, probabilityPath, discountPath,
                   concentrationPath, obs);

  // static int j = 0;
  //if (j == 1) {
    this->parameters.stepParameterGradient(10e-4);
  //  j = 0;
  //}

  return probabilityPath;
}

d_vec HPYPModel::insertObservation(l_type start, l_type stop, e_type obs,
    WrappedNodeList* cached_path) {
  tracer << "HPYPModel::insertObservation(" << start << ", " << stop
         << ", " << obs << ")" << std::endl;

  WrappedNodeList path;
  if (cached_path != NULL) {
    path = *cached_path;
  } else {
    path = this->contextTree.findLongestSuffix(start,stop);
  }

  tracer << "  HPYPModel::insertObservation: longest suffix path: "
         << std::endl << this->contextTree.pathToString(path) << std::endl;

  d_vec discount_path = this->parameters.getDiscounts(path);
  d_vec concentration_path = this->parameters.getConcentrations(path,
                                                                discount_path);
  d_vec prob_path = this->computeProbabilityPath(path,
                                                 discount_path,
                                                 concentration_path,
                                                 obs);
  this->updatePath(path, prob_path, discount_path,
                   concentration_path, obs);
  return prob_path;
}

void HPYPModel::removeObservation(
    l_type start, l_type stop, e_type obs,
    const HPYPModel::PayloadDataPath& payloadDataPath,
    WrappedNodeList* cached_path) {
  tracer << "HPYPModel::removeObservation(" << start << ", " << stop
         << ", " << obs << ")" << std::endl;

  WrappedNodeList path;
  if (cached_path != NULL) {
    path = *cached_path;
  } else {
    path = this->contextTree.findLongestSuffix(start,stop);
  }
  assert(path.back().end == (*cached_path).back().end);

  tracer << "  HPYPModel::removeObservation: longest suffix path: "
         << std::endl << contextTree.pathToString(path) << std::endl;

  d_vec discountPath = this->parameters.getDiscounts(path);

  this->removeObservationFromPath(path,
                                  discountPath,
                                  obs,
                                  payloadDataPath);
}

void HPYPModel::removeAddSweep(l_type start, l_type stop) {
  // start timer
  clock_t start_t,end_t;
  start_t = clock();
  for (int i = start; i < stop; ++i) {
    HPYPModel::PayloadDataPath payloadDataPath;
    WrappedNodeList path = this->contextTree.findNode(start, i);
    //WrappedNodeList p2 = this->contextTree.findLongestSuffix(start, i);
    //std::cout << path.back().end << ", " << p2.back().end
    //          << ", " << path.size() << ", " << p2.size() <<std::endl;
    //assert(path.back().end == p2.back().end);
    this->removeObservation(start, i, this->seq[i], payloadDataPath, &path);
    this->insertObservation(start, i, this->seq[i], &path);

    if (i%10000==0) {
      end_t = clock();
      std::cerr << makeProgressBarString(i/(double)stop) << " "
                << ((double)i*CLOCKS_PER_SEC)/(end_t-start_t) << " chars/sec"
                <<  "\r";
    }
  }
}

d_vec HPYPModel::computeLosses(l_type start, l_type stop) {
  d_vec losses;

  // deal with first symbol: add loss and insert customer
  losses.push_back(log2((double) this->numTypes));
  insertRoot(this->seq[start]);

  // start timer
  clock_t start_t,end_t;
  start_t = clock();

  for (l_type i=start+1; i < stop; i++) {
    d_vec prob_path = this->insertContextAndObservation(start,i,this->seq[i]);
    double prob = prob_path[prob_path.size()-2];
    losses.push_back(-log2(prob));

    if (i%10000==0) {
      end_t = clock();
      std::cerr << makeProgressBarString(i/(double)stop) << " "
                << ((double)i*CLOCKS_PER_SEC)/(end_t-start_t) << " chars/sec"
                <<  "\r";
    }

  }
  end_t = clock();

  std::cerr << makeProgressBarString(1) << " "
            << ((double)stop*CLOCKS_PER_SEC)/(end_t-start_t) << " chars/sec"
            <<  std::endl;

  return losses;
}

d_vec HPYPModel::computeLossesWithDeletion(l_type start, l_type stop, l_type lag) {
  d_vec losses;

  // deal with first symbol: add loss and insert customer
  losses.push_back(log2((double) this->numTypes));
  insertRoot(this->seq[start]);

  // start timer
  clock_t start_t,end_t;
  start_t = clock();

  for (l_type i=start+1; i < stop; i++) {
    d_vec prob_path = this->insertContextAndObservation(start,i,this->seq[i]);
    double prob = prob_path[prob_path.size()-2];
    losses.push_back(-log2(prob));
    if (i - lag >= start) {
      HPYPModel::PayloadDataPath payloadDataPath;
      WrappedNodeList path = this->contextTree.findNode(start, i - lag);
      this->removeObservation(start, i - lag, this->seq[i - lag], payloadDataPath, &path);
    }

    if (i%10000==0) {
      end_t = clock();
      std::cerr << makeProgressBarString(i/(double)stop) << " "
        << ((double)i*CLOCKS_PER_SEC)/(end_t-start_t) << " chars/sec"
        <<  "\r";
    }

  }
  end_t = clock();

  std::cerr << makeProgressBarString(1) << " "
    << ((double)stop*CLOCKS_PER_SEC)/(end_t-start_t) << " chars/sec"
    <<  std::endl;

  return losses;
}

d_vec HPYPModel::predictSequence(l_type start, l_type stop, PredictMode mode) {
  d_vec probs;
  for (l_type i = start; i < stop; i++) {
    switch(mode) {
      case ABOVE:
        probs.push_back(this->predict(start, i, this->seq[i]));
        break;
      case FRAGMENT:
        probs.push_back(this->predictWithFragmentation(start, i, this->seq[i]));
        break;
      case BELOW:
        probs.push_back(this->predictBelow(start, i, this->seq[i]));
        break;
    }
  }
  return probs;
}

void HPYPModel::buildTree(l_type stop) {
  this->insertRoot(this->seq[0]);
  for (l_type i=1; i < stop; ++i) {
    this->insertContextAndObservation(0, i, this->seq[i]);
  }
}

void HPYPModel::updateTree(l_type start, l_type stop) {
  for (l_type i = start; i < stop; ++i) {
    this->insertContextAndObservation(0, i, this->seq[i]);
  }
}

double HPYPModel::predict(l_type start, l_type stop, e_type obs) {
  WrappedNodeList path = this->contextTree.findLongestSuffix(start,stop);
  d_vec discount_path = this->parameters.getDiscounts(path);
  d_vec concentration_path = this->parameters.getConcentrations(path,
                                                                discount_path);

  d_vec prob_path = this->computeProbabilityPath(path,
                                                 discount_path,
                                                 concentration_path,
                                                 obs);
  return prob_path.back();
}

/**
 * Predict prob; in the case of required fragmentation, predict form
 * _below_ the split point!
 */
double HPYPModel::predictBelow(l_type start, l_type stop, e_type obs) {
  WrappedNodeList path =
      this->contextTree.findLongestSuffixVirtual(start,stop).second;
  d_vec discount_path = this->parameters.getDiscounts(path);
  d_vec concentration_path = this->parameters.getConcentrations(path,
                                                                discount_path);

  d_vec prob_path = this->computeProbabilityPath(path,
                                                 discount_path,
                                                 concentration_path,
                                                 obs);
  return prob_path.back();
}

double HPYPModel::predictWithFragmentation(l_type start,
                                           l_type stop,
                                           e_type obs) {
  std::pair<int, WrappedNodeList> path = contextTree.findLongestSuffixVirtual(
      start, stop);

  d_vec discountPath = parameters.getDiscounts(path.second);
  d_vec concentrationPath = parameters.getConcentrations(path.second,
                                                         discountPath);
  d_vec probabilityPath = this->computeProbabilityPath(path.second,
                                                       discountPath,
                                                       concentrationPath,
                                                       obs);

  double probability = 0;
  if (path.first != 0) {
    // create a new payload for the node we are predicting from
    // fragmentation -- last probability on the path needs to be recomputed
    // by creating a new, split node of length path.second.
    void* splitNode = this->restaurant.getFactory().make();
    WrappedNodeList::iterator it = path.second.end();
    it--; it--; // one before last; parent of node we need to split
    int parentLength = it->end - it->start;
    // path.first is length of parent after split
    double discountAfter = parameters.getDiscount(path.first, it->end - it->start);
    it++; // last node
    double discountFragmented = this->parameters.getDiscount(parentLength,
                                                             path.first);
    this->restaurant.updateAfterSplit(
        it->payload,
        splitNode,
        discountPath.back(),
        discountFragmented,
        true); // update splitNode only
    double concentrationFragmented = this->parameters.getConcentration(
        discountFragmented, parentLength, path.first);
    probability = this->restaurant.computeProbability(
        splitNode, obs, probabilityPath[probabilityPath.size()-2],
        discountFragmented, concentrationFragmented);
    this->restaurant.getFactory().recycle(splitNode);
  } else {
    probability = probabilityPath.back();
  }
  return probability;
}

d_vec HPYPModel::predictiveDistribution(l_type start, l_type stop) {
  d_vec predictive;
  predictive.reserve(this->numTypes);
  WrappedNodeList path = this->contextTree.findLongestSuffix(start,stop);
  d_vec discount_path = this->parameters.getDiscounts(path);
  d_vec concentration_path = this->parameters.getConcentrations(path,
                                                                discount_path);

  for(int i = 0; i < this->numTypes; ++i) {
    d_vec prob_path = this->computeProbabilityPath(path,
                                                   discount_path,
                                                   concentration_path,
                                                   i);
    predictive.push_back(prob_path.back());
  }

  return predictive;
}

d_vec HPYPModel::predictiveDistributionWithMixing(l_type start,
                                                  l_type stop,
                                                  d_vec& mixingWeights) {
  d_vec predictive;
  predictive.reserve(this->numTypes);
  const WrappedNodeList path = this->contextTree.findLongestSuffix(start, stop);
  const d_vec discount_path = this->parameters.getDiscounts(path);
  const d_vec concentration_path
      = this->parameters.getConcentrations(path, discount_path);

  for(int i = 0; i < this->numTypes; ++i) {
    d_vec prob_path = this->computeProbabilityPath(path,
                                                   discount_path,
                                                   concentration_path,
                                                   i);
    double prob = 0;
    double sum = 0;
    for (int j = 0;
         j < std::min<int>(mixingWeights.size(), prob_path.size());
         ++j) {
      prob += mixingWeights[j]*prob_path[j];
      sum += mixingWeights[j];
    }
    predictive.push_back(prob + (1 - sum)*prob_path.back());
  }

  return predictive;
}

/**
 * Given a path from the root to some node (not a leaf),
 * perform add/remove Gibbs sampling of the last node by repeatedly
 * removing and adding customers, cus times for each type s.
 */
void HPYPModel::addRemoveSamplePath(
    const WrappedNodeList& path,
    const d_vec& discountPath,
    const d_vec& concentrationPath,
    const HPYPModel::PayloadDataPath& payloadDataPath,
    double baseProb) {
  assert(path.size() > 0);
  assert(path.size() == discountPath.size());
  assert(path.size() == concentrationPath.size());

  bool useAdditionalData = payloadDataPath.size() == path.size();
  void* main = path.back().payload;
  const IAddRemoveRestaurant& r = this->restaurant; // shortcut

  IHPYPBaseRestaurant::TypeVector types = r.getTypeVector(main);
  for(IHPYPBaseRestaurant::TypeVectorIterator it = types.begin();
      it != types.end(); ++it) { // for each type of customer

    e_type type = *it;
    l_type cw = r.getC(main, type);

    if (cw == 1) {
      continue; // no point in reseating in a 1 customer restaurant
    }

    d_vec probabilityPath = this->computeProbabilityPath(path,
                                                         discountPath,
                                                         concentrationPath,
                                                         type);
    WrappedNodeList::const_iterator current;
    for (l_type i = 0; i < cw; ++i) { // for each customer of this type
      current = --path.end(); // set current to last restaurant in path

      // index into d/alpha vectors for current restaurant
      int j = discountPath.size() - 1;

      bool goUp = true;
      while(goUp && j != -1) {
        void* additionalData = NULL;
        if (useAdditionalData) {
          additionalData = payloadDataPath[j].get();
        }
        bool removed = r.removeCustomer(current->payload,
                                        type,
                                        discountPath[j],
                                        additionalData);
        if (removed) {
          --current;
          --j;
        } else {
          goUp = false;
        }
      }

      // recompute probabilities back down; need not recompute last probability
      while(j < (int)probabilityPath.size() - 1) {
        if (j == -1) {
          // can't recompute base distribution probabilities at prob_path[0]
          j = 0; //
          ++current;
        }
        probabilityPath[j+1] = r.computeProbability(current->payload,
                                                    type,
                                                    probabilityPath[j],
                                                    discountPath[j],
                                                    concentrationPath[j]);
        ++j;
        ++current;
      }

      // set current and j to the last restaurant on path
      current = --path.end();
      j = discountPath.size() - 1;
      goUp = true;
      while(goUp && j != -1) {
        void* additionalData = NULL;
        if (useAdditionalData) {
          additionalData = payloadDataPath[j].get();
        }
        bool inserted = r.addCustomer(current->payload,
                                      type,
                                      probabilityPath[j],
                                      discountPath[j],
                                      concentrationPath[j],
                                      additionalData);
        if (inserted) {
          current--;
          j--;
        } else {
          goUp = false;
        }
      }

    }

  }

}

/**
 * Given a path from the root to some node (not a leaf),
 * perform add/remove Gibbs sampling of the last node by repeatedly
 * removing and adding customers, cus times for each type s.
 */
void HPYPModel::directGibbsSamplePath(
    const WrappedNodeList& path,
    const d_vec& discountPath,
    const d_vec& concentrationPath,
    const HPYPModel::PayloadDataPath& payloadDataPath,
    double baseProb) {
  assert(path.size() > 0);
  assert(path.size() == discountPath.size());
  assert(path.size() == concentrationPath.size());

  bool useAdditionalData = payloadDataPath.size() == path.size();
  // XXX: HACK! We assume this is the type of addData for the used restaurant
  stirling_generator_full_log *stirlingGenCurrent = NULL;
  stirling_generator_full_log *stirlingGenParent = NULL;
  void* main = path.back().payload;
  const BaseCompactRestaurant& r = (BaseCompactRestaurant&)this->restaurant; // shortcut

  IHPYPBaseRestaurant::TypeVector types = r.getTypeVector(main);
  for(IHPYPBaseRestaurant::TypeVectorIterator it = types.begin();
      it != types.end(); ++it) { // for each type of customer

    e_type type = *it;
    l_type cw = r.getC(main, type);

    if (cw == 1) {
      continue; // no point in reseating in a 1 customer restaurant
    }

    WrappedNodeList::const_iterator current;
    current = --path.end(); // set current to last restaurant in path

    // index into d/alpha vectors for current restaurant
    int j = discountPath.size() - 1;

    bool goUp = true;
    while(goUp && j != -1) {
      goUp = false;
      stirlingGenCurrent = (stirling_generator_full_log*)payloadDataPath[j].get();
      void* currentPayload = (*current).payload;
      void* parentPayload = NULL; // initialized below
      int currentCw = r.getC(currentPayload, type);
      int currentTw = r.getT(currentPayload, type);
      int otherT = r.getT(currentPayload) - currentTw;
      std::vector<double> logProbs(currentCw, 0);
      std::vector<double> logProbs1(currentCw, 0);
      std::vector<double> logProbs2(currentCw, 0);
      std::vector<double> logProbs3(currentCw, 0);
      std::vector<double> logProbs4(currentCw, 0);
      if (j > 0) { // not at the top, so we have a CRP parent
        stirlingGenParent = (stirling_generator_full_log*)payloadDataPath[j-1].get();
        current--; // move to parent
        parentPayload = (*current).payload;
        current++;

        int parentTw = r.getT(parentPayload, type);
        int parentCw = r.getC(parentPayload, type);
        int parentOtherC = r.getC(parentPayload) - currentTw;
        for (int tw = 1; tw <= currentCw; ++tw) {
          int newParentCw = parentCw - currentTw + tw;
          if (newParentCw < parentTw) {
            logProbs4[tw-1] = -INFINITY;
          } else {
            logProbs1[tw-1] = logKramp(concentrationPath[j] + discountPath[j], discountPath[j], otherT + tw - 1);
            logProbs2[tw-1] = - logKramp(concentrationPath[j-1] + 1, 1, parentOtherC + tw - 1);
            logProbs3[tw-1] = stirlingGenCurrent->getLog(currentCw, tw);
            logProbs4[tw-1] = stirlingGenParent->getLog(newParentCw, parentTw);
            //std::cerr <<  logKramp(concentrationPath[j] + discountPath[j], discountPath[j], otherT + tw - 1) << std::endl;
            //std::cerr <<  logKramp(concentrationPath[j-1], 1, parentOtherC + tw - 1)<< std::endl;
            //std::cerr <<  stirlingGenCurrent->getLog(cw, tw)<< std::endl;
            //std::cerr << cw << ", " << tw << std::endl;
            //std::cerr <<  stirlingGenParent->getLog(newParentCw, parentTw)<< std::endl;
            //std::cerr << newParentCw << ", " << parentTw << std::endl;
          }
        }
        //std::cerr << parentCw << ", " << parentTw << ", " << parentOtherC << ", " << currentTw << ", " << otherT << std::endl;
      } else { // at the root, take base prob into account
        for (int tw = 1; tw <= currentCw; ++tw) {
          logProbs1[tw-1] = logKramp(concentrationPath[j] + discountPath[j], discountPath[j], otherT + tw - 1);
          logProbs2[tw-1] = stirlingGenCurrent->getLog(currentCw, tw);
          logProbs3[tw-1] = tw * log(baseProb);
        }
      }
      // subtract max for stability
      //std::cerr << iterableToString(logProbs1) << std::endl;
      //std::cerr << iterableToString(logProbs2) << std::endl;
      //std::cerr << iterableToString(logProbs3) << std::endl;
      //std::cerr << iterableToString(logProbs4) << std::endl;
      subMax_vec(logProbs1);
      subMax_vec(logProbs2);
      subMax_vec(logProbs3);
      subMax_vec(logProbs4);
      add_vec(logProbs, logProbs1);
      add_vec(logProbs, logProbs2);
      add_vec(logProbs, logProbs3);
      add_vec(logProbs, logProbs4);
      subMax_vec(logProbs);

      //std::cerr << iterableToString(logProbs1) << std::endl;
      //std::cerr << iterableToString(logProbs2) << std::endl;
      //std::cerr << iterableToString(logProbs3) << std::endl;
      //std::cerr << iterableToString(logProbs4) << std::endl;
      //std::cerr << iterableToString(logProbs) << std::endl;
      exp_vec(logProbs);
      //std::cerr << iterableToString(logProbs) << std::endl;
      // if (max == -INFINITY) { // if all choices are improbable, choose uniformly
      //     // we can do better by normalizing the component individually
      //   for (int ii = 0; ii < logProbs.size(); ++ii)
      //     logProbs[ii] = 1.0;
      // }
      //std::cerr << iterableToString(logProbs) << std::endl;
      int sampledTw = sample_unnormalized_pdf(logProbs, 0) + 1;
      //std::cerr << "Old tw: " << currentTw << ", newTw: " << sampledTw << std::endl;

      r.setT(currentPayload, type, sampledTw);
      if (j > 0) {
        int newCw =  r.getC(parentPayload, type) - currentTw + sampledTw;
        assert(newCw >= r.getT(parentPayload, type));
        r.setC(parentPayload, type, newCw);
      }

      if (sampledTw != currentTw) {
        --current;
        --j;
      } else {
        goUp = false;
      }
    }

  }

}

boost::shared_ptr<void> HPYPModel::makeAdditionalDataPtr(void* payload,
                                                         double discount,
                                                         double concentration) const {
  return boost::shared_ptr<void>(
      this->restaurant.createAdditionalData(payload,
        discount,
        concentration),
      boost::bind(&IAddRemoveRestaurant::freeAdditionalData,
        &(this->restaurant),
        _1));
}

void HPYPModel::runGibbsSampler(bool directGibbs) {
  ContextTree::DFSPathIterator pathIterator = contextTree.getDFSPathIterator();
  d_vec discountPath = parameters.getDiscounts(*pathIterator);
  d_vec concentrationPath = parameters.getConcentrations(*pathIterator,
                                                         discountPath);

  // initialize payloadDataPath; by using shared_ptr with the proper
  // destruction function, all clean-up should be automatic.
  HPYPModel::PayloadDataPath payloadDataPath;
  int j = 0;
  for (WrappedNodeList::const_iterator it = (*pathIterator).begin();
       it != (*pathIterator).end(); ++it) {
      payloadDataPath.push_back(this->makeAdditionalDataPtr(
            it->payload, discountPath[j], concentrationPath[j]));
    j++;
  }

  if (directGibbs) {
    this->directGibbsSamplePath(*pathIterator, discountPath, concentrationPath,
                                payloadDataPath, baseProb);
  } else {
    this->addRemoveSamplePath(*pathIterator, discountPath, concentrationPath,
                              payloadDataPath, baseProb);
  }

  size_t pathLength = (*pathIterator).size();

  while(pathIterator.hasMore()) { // loop over all paths in the tree
    ++pathIterator;
    if ((*pathIterator).size() == 0) {
      break;
    }

    if ((*pathIterator).size() == pathLength) {
      // sibling
      discountPath.pop_back();
      parameters.extendDiscounts(*pathIterator, discountPath);
      concentrationPath.pop_back();
      parameters.extendConcentrations(*pathIterator,
                                      discountPath,
                                      concentrationPath);
      payloadDataPath.pop_back();
      payloadDataPath.push_back(
          this->makeAdditionalDataPtr((*pathIterator).back().payload,
                                      discountPath.back(),
                                      concentrationPath.back()));

    } else {
      if ((*pathIterator).size() == pathLength - 1) {
        // we went up -- just drop the last term
        discountPath.pop_back();
        concentrationPath.pop_back();
        payloadDataPath.pop_back();
      } else {
        // we went up one and then down some number of levels -- recompute
        discountPath.pop_back();
        concentrationPath.pop_back();
        parameters.extendDiscounts(*pathIterator, discountPath);
        parameters.extendConcentrations(*pathIterator,
                                        discountPath,
                                        concentrationPath);
        payloadDataPath.pop_back();
        WrappedNodeList::const_iterator it = (*pathIterator).begin();
        // move it to the first item not covered by the payloadDataPath
        for (size_t i = 0; i < payloadDataPath.size(); ++i) {
          ++it;
        }
        for (size_t i = payloadDataPath.size(); i < discountPath.size(); ++i) {
          payloadDataPath.push_back(
              this->makeAdditionalDataPtr(it->payload,
                                          discountPath[i],
                                          concentrationPath[i]));

          ++it;
        }
        assert(it == (*pathIterator).end());
      }
    }
    pathLength = (*pathIterator).size();

    tracer << (*pathIterator).size() << " " << discountPath.size() << " "
           << concentrationPath.size() << " " <<  payloadDataPath.size()
           << std::endl;

    if (directGibbs) {
      this->directGibbsSamplePath(*pathIterator, discountPath, concentrationPath,
                                  payloadDataPath, baseProb);
    } else {
      this->addRemoveSamplePath(*pathIterator, discountPath, concentrationPath,
                                payloadDataPath, baseProb);
    }
  }
}

d_vec HPYPModel::computeProbabilityPath(const WrappedNodeList& path,
                                        const d_vec& discount_path,
                                        const d_vec& concentration_path,
                                        e_type obs) {
  d_vec out;
  out.reserve(path.size() + 1);

  double prob = this->baseProb; // base distribution
  out.push_back(prob);

  int j = 0;
  for(WrappedNodeList::const_iterator it = path.begin();
      it!= path.end();
      ++it) {
    prob = this->restaurant.computeProbability(it->payload,
                                               obs,
                                               prob,
                                               discount_path[j],
                                               concentration_path[j]);
    out.push_back(prob);
    j++;
  }
  return out;
}

void HPYPModel::updatePath(const WrappedNodeList& path,
                           const d_vec& prob_path,
                           const d_vec& discount_path,
                           const d_vec& concentration_path,
                           e_type obs) {

  unsigned int j=path.size()-1;
  double newTable = 1;
  for(WrappedNodeList::const_reverse_iterator it = path.rbegin();
      it != path.rend();
      ++it) {
    newTable = this->restaurant.addCustomer(it->payload,
                                            obs,
                                            prob_path[j],
                                            discount_path[j],
                                            concentration_path[j],
                                            NULL,
                                            newTable);
    if (newTable==0) {
      break;
    }
    j--;
  }
}

void HPYPModel::removeObservationFromPath(
    const WrappedNodeList& path,
    const d_vec& discountPath,
    e_type obs,
    const HPYPModel::PayloadDataPath& payloadDataPath) {
  int j = path.size()-1;

  double frac_t = 1;

  for(WrappedNodeList::const_reverse_iterator it = path.rbegin();
    it != path.rend(); it++) {

    void* payloadData = NULL;
    if (payloadDataPath.size() == path.size()) {
      payloadData = payloadDataPath[j].get();
    }

    frac_t = this->restaurant.removeCustomer(it->payload,
                                             obs,
                                             discountPath[j],
                                             payloadData, frac_t);
    if (frac_t == 0.)
      break;
    j--;
  }
}

void HPYPModel::handleSplit(const WrappedNode& nodeA,
                            const WrappedNode& nodeB,
                            const WrappedNode& nodeC) {
    int lengthA = nodeA.end - nodeA.start;
    int lengthB = nodeB.end - nodeB.start;
    int lengthC = nodeC.end - nodeC.start;

    // parent context should be shorter than both its children
    assert(lengthA < lengthB && lengthA < lengthC);
    // length of node that required splitting should be longer than the result
    assert(lengthC < lengthB);

    double discBBeforeSplit = this->parameters.getDiscount(lengthA, lengthB);
    double discBAfterSplit  = this->parameters.getDiscount(lengthC, lengthB);
    this->restaurant.updateAfterSplit(nodeB.payload,
                                      nodeC.payload,
                                      discBBeforeSplit,
                                      discBAfterSplit);
}

bool HPYPModel::checkConsistency(const WrappedNode& node,
                      const std::list<WrappedNode>& children) const {
  bool consistent = this->restaurant.checkConsistency(node.payload);
  std::map<e_type, int> table_counts;
  for(std::list<WrappedNode>::const_iterator it = children.begin();
      it != children.end(); ++it) {
    IHPYPBaseRestaurant::TypeVector keys =
        this->restaurant.getTypeVector(it->payload);
    for(IHPYPBaseRestaurant::TypeVectorIterator key_it = keys.begin();
        key_it != keys.end(); ++key_it) {
      table_counts[*key_it] += this->restaurant.getT(it->payload, *key_it);
    }
  }

  for(std::map<e_type, int>::iterator it = table_counts.begin();
      it != table_counts.end(); ++it) {
    consistent = (this->restaurant.getC(node.payload, (*it).first)
                  >= (*it).second) && consistent;
    if (!consistent) {
      std::cerr << "Child table sum is: " << (*it).second
                << ", parent customers is: "
                << this->restaurant.getC(node.payload, (*it).first)
                << std::endl;
    }
  }
  return consistent;
}

void HPYPModel::prunePath(WrappedNodeList& path) {

}

bool HPYPModel::checkConsistency() const {
  CheckConsistencyVisitor v(*this);
  this->contextTree.visitDFSWithChildren(v);
  return v.consistent;
}

double HPYPModel::computeLogRestaurantProb(
    const WrappedNodeList& path,
    const d_vec& discountPath,
    const d_vec& concentrationPath,
    const HPYPModel::PayloadDataPath& payloadDataPath,
    double baseProb) const {
  assert(path.size() > 0);
  assert(path.size() == discountPath.size());
  assert(path.size() == concentrationPath.size());

  void* payload = path.back().payload;
  const BaseCompactRestaurant& r = (BaseCompactRestaurant&)this->restaurant; // shortcut
  double logProb = 0;
  int j = discountPath.size() - 1;
  l_type c = r.getC(payload);
  if (c == 1) {
      // deterministic restaurant
      return 0;
  }
  l_type t = r.getT(payload);
  //std::cerr << "c: " << c << ", t: " << t << std::endl;
  //std::cerr << "a: " << concentrationPath[j] << ", d: " << discountPath[j] << std::endl;
  logProb += logKramp(concentrationPath[j] + discountPath[j], discountPath[j], t - 1);
  //std::cerr << logProb << std::endl;
  logProb -= logKramp(concentrationPath[j] + 1, 1, c - 1);
  //std::cerr << logProb << std::endl;

  stirling_generator_full_log *stirlingGen = (stirling_generator_full_log*)payloadDataPath[j].get();
  IHPYPBaseRestaurant::TypeVector types = r.getTypeVector(payload);
  for(IHPYPBaseRestaurant::TypeVectorIterator it = types.begin();
      it != types.end(); ++it) { // for each type of customer

    e_type type = *it;
    l_type cw = r.getC(payload, type);
    l_type tw = r.getT(payload, type);

    logProb += stirlingGen->getLog(cw, tw);

    if (j == 0) { // at the root, take base prob into account
      logProb += tw * log(baseProb);
    }
  }
  //std::cerr << logProb << std::endl;
  return logProb;
}

double HPYPModel::computeLogJoint() const {
  double logJoint = 0;
  ContextTree::DFSPathIterator pathIterator = contextTree.getDFSPathIterator();
  d_vec discountPath = parameters.getDiscounts(*pathIterator);
  d_vec concentrationPath = parameters.getConcentrations(*pathIterator,
                                                         discountPath);

  // initialize payloadDataPath; by using shared_ptr with the proper
  // destruction function, all clean-up should be automatic.
  HPYPModel::PayloadDataPath payloadDataPath;
  int j = 0;
  for (WrappedNodeList::const_iterator it = (*pathIterator).begin();
       it != (*pathIterator).end(); ++it) {
      payloadDataPath.push_back(this->makeAdditionalDataPtr(
            it->payload, discountPath[j], concentrationPath[j]));
    j++;
  }

  logJoint += computeLogRestaurantProb(*pathIterator, discountPath, concentrationPath, payloadDataPath, baseProb);

  size_t pathLength = (*pathIterator).size();

  while(pathIterator.hasMore()) { // loop over all paths in the tree
    ++pathIterator;
    if ((*pathIterator).size() == 0) {
      break;
    }

    if ((*pathIterator).size() == pathLength) {
      // sibling
      discountPath.pop_back();
      parameters.extendDiscounts(*pathIterator, discountPath);
      concentrationPath.pop_back();
      parameters.extendConcentrations(*pathIterator,
                                      discountPath,
                                      concentrationPath);
      payloadDataPath.pop_back();
      payloadDataPath.push_back(
          this->makeAdditionalDataPtr((*pathIterator).back().payload,
                                      discountPath.back(),
                                      concentrationPath.back()));

    } else {
      if ((*pathIterator).size() == pathLength - 1) {
        // we went up -- just drop the last term
        discountPath.pop_back();
        concentrationPath.pop_back();
        payloadDataPath.pop_back();
      } else {
        // we went up one and then down some number of levels -- recompute
        discountPath.pop_back();
        concentrationPath.pop_back();
        parameters.extendDiscounts(*pathIterator, discountPath);
        parameters.extendConcentrations(*pathIterator,
                                        discountPath,
                                        concentrationPath);
        payloadDataPath.pop_back();
        WrappedNodeList::const_iterator it = (*pathIterator).begin();
        // move it to the first item not covered by the payloadDataPath
        for (size_t i = 0; i < payloadDataPath.size(); ++i) {
          ++it;
        }
        for (size_t i = payloadDataPath.size(); i < discountPath.size(); ++i) {
          payloadDataPath.push_back(
              this->makeAdditionalDataPtr(it->payload,
                                          discountPath[i],
                                          concentrationPath[i]));

          ++it;
        }
        assert(it == (*pathIterator).end());
      }
    }
    pathLength = (*pathIterator).size();

    logJoint += computeLogRestaurantProb(*pathIterator, discountPath, concentrationPath, payloadDataPath, baseProb);

  }
  return logJoint;
}

HPYPModel::ToStringVisitor::ToStringVisitor(seq_type& seq,
    const IHPYPBaseRestaurant& restaurant)
    : outstream(), seq(seq), restaurant(restaurant) {}

void HPYPModel::ToStringVisitor::operator()(const WrappedNode& n) {
  for (l_type i = 0; i < n.depth; ++i) {
    this->outstream << " ";
  }
  this->outstream << SubSeq::toString(n.start, n.end, seq);
  this->outstream << " " << this->restaurant.toString(n.payload);
  this->outstream << std::endl;
}

HPYPModel::CheckConsistencyVisitor::CheckConsistencyVisitor(
    const HPYPModel& model) : consistent(true), model(model) {}

void HPYPModel::CheckConsistencyVisitor::operator()(
    WrappedNode& n, std::list<WrappedNode>& children) {
  bool nodeConsistent = this->model.checkConsistency(n, children);
  if (!nodeConsistent) {
    std::cerr << "Node " << n.toString() << " not consistent!" << std::endl;
  }
  consistent = consistent && nodeConsistent;
}

HPYPModel::LogJointVisitor::LogJointVisitor(
    const HPYPModel& model) : logJoint(0), model(model) {}

void HPYPModel::LogJointVisitor::operator()(
    // FIXME
    WrappedNode& n, std::list<WrappedNode>& children) {
    int c = model.restaurant.getC(n.payload);
    int t = model.restaurant.getT(n.payload);
    IHPYPBaseRestaurant::TypeVector keys =
        model.restaurant.getTypeVector(n.payload);
    for(IHPYPBaseRestaurant::TypeVectorIterator key_it = keys.begin();
        key_it != keys.end(); ++key_it) {
      int cw = model.restaurant.getC(n.payload, *key_it);
      int tw = model.restaurant.getT(n.payload, *key_it);
    }
}

std::string HPYPModel::toString() {
  HPYPModel::ToStringVisitor visitor(this->seq, this->restaurant);
  this->contextTree.visitDFS(visitor);
  return visitor.outstream.str();
}

}} // namespace gatsby::libplump