/*

Copyright (c) 2005-2021, University of Oxford.
All rights reserved.

University of Oxford means the Chancellor, Masters and Scholars of the
University of Oxford, having an administrative office at Wellington
Square, Oxford OX1 2JD, UK.

This file is part of Chaste.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
 * Redistributions of source code must retain the above copyright notice,
   this list of conditions and the following disclaimer.
 * 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.
 * Neither the name of the University of Oxford 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.

*/

#ifndef ABSTRACTGENERALIZEDRUSHLARSENCARDIACCELL_HPP_
#define ABSTRACTGENERALIZEDRUSHLARSENCARDIACCELL_HPP_

#include "ChasteSerialization.hpp"
#include <boost/serialization/base_object.hpp>
#include "ClassIsAbstract.hpp"
#include "AbstractCardiacCell.hpp"
#include "PetscTools.hpp"

/*
Megan E. Marsh, Raymond J. Spiteri
Numerical Simulation Laboratory
University of Saskatchewan
December 2011
Partial support provided by research grants from the National
Science and Engineering Research Council (NSERC) of Canada
and the MITACS/Mprime Canadian Network of Centres of Excellence.
*/

/**
 * This is the base class for cardiac cells solved using the GRL methods (GRL1 and GRL2).
 * Modified from AbstractRushLarsenCardiacCell.hpp
 */
class AbstractGeneralizedRushLarsenCardiacCell : public AbstractCardiacCell
{
private:
    /** Needed for serialization. */
    friend class boost::serialization::access;
    /**
     * Archive the member variables.
     *
     * @param archive
     * @param version
     */
    template<class Archive>
    void serialize(Archive & archive, const unsigned int version)
    {
        // This calls serialize on the base class.
        archive & boost::serialization::base_object<AbstractCardiacCell>(*this);
    }

public:
    /**
     * Standard constructor for a cell.
     *
     * @param numberOfStateVariables  the size of the ODE system
     * @param voltageIndex  the index of the variable representing the transmembrane
     *     potential within the state variable vector
     * @param pIntracellularStimulus  the intracellular stimulus function
     *
     * Some notes for future reference:
     *  \li It's a pity that inheriting from AbstractCardiacCell forces us to store a
     *      null pointer (for the unused ODE solver) in every instance.  We may want
     *      to revisit this design decision at a later date.
     */
    AbstractGeneralizedRushLarsenCardiacCell(
            unsigned numberOfStateVariables,
            unsigned voltageIndex,
            boost::shared_ptr<AbstractStimulusFunction> pIntracellularStimulus);

    /** Virtual destructor */
    virtual ~AbstractGeneralizedRushLarsenCardiacCell();

    /**
     * Simulates this cell's behaviour between the time interval [tStart, tEnd],
     * with timestep #mDt.
     *
     * The length of the time interval must be a multiple of the timestep.
     *
     * @param tStart  beginning of the time interval to simulate
     * @param tEnd  end of the time interval to simulate
     * @param tSamp  sampling interval for returned results (defaults to #mDt)
     * @return  the values of each state variable, at intervals of tSamp.
     */
    OdeSolution Compute(double tStart, double tEnd, double tSamp=0.0);

    /**
     * Simulates this cell's behaviour between the time interval [tStart, tEnd],
     * with timestep #mDt.  The transmembrane potential is kept fixed throughout,
     * but the other state variables are updated.
     *
     * The length of the time interval must be a multiple of the timestep.
     *
     * @param tStart  beginning of the time interval to simulate
     * @param tEnd  end of the time interval to simulate
     */
    void ComputeExceptVoltage(double tStart, double tEnd);

    /**
     * Simulate this cell's behaviour between the time interval [tStart, tEnd],
     * with timestemp #mDt, updating the internal state variable values.
     *
     * @param tStart  beginning of the time interval to simulate
     * @param tEnd  end of the time interval to simulate
     */
    void SolveAndUpdateState(double tStart, double tEnd);

    /**
     * @return whether the ODE system has an analytic Jacobian (#mHasAnalyticJacobian).
     */
    bool HasAnalyticJacobian() const;

    /**
     * Force the use of a numerical Jacobian, even if an analytic form is provided.
     * This is needed for a handful of troublesome models.
     *
     * @param useNumericalJacobian  Whether to use a numerical instead of the analytic Jacobian.
     */
    void ForceUseOfNumericalJacobian(bool useNumericalJacobian = true);

private:
// LCOV_EXCL_START
    /**
     * This function should never be called - the cell class incorporates its own solver.
     *
     * @param time
     * @param rY
     * @param rDY
     */
    void EvaluateYDerivatives(double time, const std::vector<double>& rY, std::vector<double>& rDY)
    {
        NEVER_REACHED;
    }
// LCOV_EXCL_STOP

protected:
    /**
     * Update the values of all variables except the transmembrane potential using a GRL method.
     *
     * @param time  the current simulation time
     */
    virtual void ComputeOneStepExceptVoltage(double time)=0;

    /**
     * Perform a forward Euler step to update the transmembrane potential.
     *
     * @param time  the current simulation time
     */
    virtual void UpdateTransmembranePotential(double time)=0;

    /** The diagonal of the Jacobian, working memory for use by subclasses. */
    std::vector<double> mPartialF;

    /** The derivatives, working memory for use by subclasses. */
    std::vector<double> mEvalF;

    /** The state at the beginning of the current step, working memory for use by subclasses. */
    std::vector<double> mYInit;

    /** Whether we have an analytic Jacobian. */
    bool mHasAnalyticJacobian;
};

CLASS_IS_ABSTRACT(AbstractGeneralizedRushLarsenCardiacCell)

#endif // ABSTRACTGENERALIZEDRUSHLARSENCARDIACCELL_HPP_


/*
    Copyright (c) 2014, Philipp Krähenbühl
    All rights reserved.

    Redistribution and use in source and binary forms, with or without
    modification, are permitted provided that the following conditions are met:
        * Redistributions of source code must retain the above copyright
        notice, this list of conditions and the following disclaimer.
        * 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.
        * Neither the name of the Stanford University 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 Philipp Krähenbühl ''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 Philipp Krähenbühl 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.
*/
#include "color.h"
#include "util/util.h"
#include "util/sse_defs.h"
#include <vector>
#include <Eigen/Core>
using namespace Eigen;
#include <emmintrin.h>

template<typename F>
std::vector<float> initTable( F f, int N, float s ) {
	std::vector<float> r( N );
	for( int i=0; i<N; i++ )
		r[i] = f(s*i);
	return r;
}
static const std::vector<float> ltable = initTable( [](float y){ return (y > ((6.0/29)*(6.0/29)*(6.0/29)) ? 116*pow((double)y,1.0/3.0)-16 : y*((29.0/3)*(29.0/3)*(29.0/3)))/270.; }, 1025, 1.0 / 1024. );
static const std::vector<float> stable = initTable( [](float v){ return (v < 0.4045f ? v / 12.92f : pow( (v+0.055f)/1.055f, 2.4f ) ); }, 1025, 1.0 / 1024. );

inline float mapTable( float v, const float * t, int N ) {
	if( v < 0 )  v = 0;
	if( v > N-1 ) v = N-1;
	return t[(int)v];
}
template<typename T>
T min( T a, T b ) {
	return std::min(a,b);
}
template<> __m128 min<__m128>( __m128 a, __m128 b ) {
	return _mm_min_ps(a,b);
}
template<typename T>
T max( T a, T b ) {
	return std::max(a,b);
}
template<> __m128 max<__m128>( __m128 a, __m128 b ) {
	return _mm_max_ps(a,b);
}
template<typename T>
static T c( float a ) {
	return a;
}
template<> __m128 c<__m128>( float a ) {
	return _mm_set1_ps( a );
}
inline __m128 mapTable( __m128 v, const float * t, int N ) {
	float * vv = (float*)&v;
	return _mm_set_ps( mapTable( vv[3], t, N ), mapTable( vv[2], t, N ), mapTable( vv[1], t, N ), mapTable( vv[0], t, N ) );
}
template<typename T>
void rgb2xyz( T & x, T & y, T & z, const T & r, const T & g, const T & b ) {
	x = c<T>(0.430574f)*r + c<T>(0.341550f)*g + c<T>(0.178325f)*b;
	y = c<T>(0.222015f)*r + c<T>(0.706655f)*g + c<T>(0.071330f)*b;
	z = c<T>(0.020183f)*r + c<T>(0.129553f)*g + c<T>(0.939180f)*b;
}
template<typename T>
T gcorr( const T & v ) {
	return mapTable( c<T>(1024)*v, stable.data(), stable.size() );
}
template<typename T>
void srgb2xyz( T & x, T & y, T & z, T r, T g, T b ) {
	r = gcorr( r );
	g = gcorr( g );
	b = gcorr( b );
	x = c<T>(0.4124564f)*r + c<T>(0.3575761f)*g + c<T>(0.1804375f)*b;
	y = c<T>(0.2126729f)*r + c<T>(0.7151522f)*g + c<T>(0.0721750f)*b;
	z = c<T>(0.0193339f)*r + c<T>(0.1191920f)*g + c<T>(0.9503041f)*b;
}
template<typename T> void xyz2luv( T& l, T& u, T& v, T x, T y, T z ) {
	const T un = c<T>(0.197833f), vn = c<T>(0.468331f);
	l = mapTable( c<T>(1024.f)*y, ltable.data(), ltable.size() );
	z = c<T>(1.f) / (x + c<T>(15.f)*y + c<T>(3.f)*z + c<T>(1e-35f) );
	u = l * (c<T>(13.f)*(c<T>(4.f)*x*z - un)) + c<T>(88.f/270.f);
	v = l * (c<T>(13.f)*(c<T>(9.f)*y*z - vn)) + c<T>(134.f/270.f);
}
template<typename T> void xyz2lab( T& l, T& a, T& b, T x, T y, T z ) {
	x *= c<T>(1.f/0.950456f);
	z *= c<T>(1.f/1.088754f);
	x = mapTable( c<T>(1024)*x, ltable.data(), ltable.size() );
	y = mapTable( c<T>(1024)*y, ltable.data(), ltable.size() );
	z = mapTable( c<T>(1024)*z, ltable.data(), ltable.size() );
	l = y;
	a = c<T>(500.f/116.f)*(x-y);
	b = c<T>(200.f/116.f)*(y-z);
}
enum ColorType{
	LUV,
	LAB
};
template<bool SRGB,ColorType type,typename T>
static void convertRGB( T& c1, T& c2, T& c3, const T & R, const T & G, const T & B ) {
	T x,y,z;
	if (SRGB)
		srgb2xyz( x,y,z, R,G,B );
	else
		rgb2xyz( x,y,z, R,G,B );
	if (type==LUV)
		xyz2luv( c1,c2,c3, x,y,z );
	else
		xyz2lab( c1,c2,c3, x,y,z );
}
template<bool SRGB,ColorType type>
static void convertRGB( Image & luv, const Image & rgb ) {
	if( rgb.C()!= 3 )
		throw std::invalid_argument( "RGB image required!" );
	const int W = rgb.W(), H = rgb.H();
	luv.create( W, H, 3 );
	int i;
	for( i=0; i+3<W*H; i+=4 ) {
		__m128 r = _mm_set_ps( rgb[3*i+0], rgb[3*(i+1)+0], rgb[3*(i+2)+0], rgb[3*(i+3)+0] );
		__m128 g = _mm_set_ps( rgb[3*i+1], rgb[3*(i+1)+1], rgb[3*(i+2)+1], rgb[3*(i+3)+1] );
		__m128 b = _mm_set_ps( rgb[3*i+2], rgb[3*(i+1)+2], rgb[3*(i+2)+2], rgb[3*(i+3)+2] );
		__m128 l,u,v;
		convertRGB<SRGB,type>( l, u, v, r, g, b );
		float * ll = (float*)&l, * uu = (float*)&u, * vv = (float*)&v;
		for( int k=0; k<4; k++ ){
			luv[3*(i+k)+0] = ll[k];
			luv[3*(i+k)+1] = uu[k];
			luv[3*(i+k)+2] = vv[k];
		}
	}
	for( i=0; i<W*H; i++ )
		convertRGB<SRGB,type>( luv[3*i+0],luv[3*i+1],luv[3*i+2], rgb[3*i+0],rgb[3*i+1],rgb[3*i+2] );
}
void rgb2luv( Image & luv, const Image & rgb ) {
	convertRGB<false,LUV>( luv, rgb );
}
void srgb2luv( Image & luv, const Image & rgb ) {
	convertRGB<true,LUV>( luv, rgb );
}
void rgb2lab( Image & lab, const Image & rgb ) {
	convertRGB<false,LAB>( lab, rgb );
}
void srgb2lab( Image & lab, const Image & rgb ) {
	convertRGB<true,LAB>( lab, rgb );
}
void rgb2hsv( Image & hsv, const Image & rgb ) {
	if( rgb.C()!= 3 )
		throw std::invalid_argument( "RGB image required!" );
	const int W = rgb.W(), H = rgb.H();
	hsv.create( W, H, 3 );
	const int N = W*H;
	for( int i=0; i<N; i++ ) {
		float s,v;
		hsv[3*i+2] = v = std::max(rgb[3*i+0],std::max(rgb[3*i+1],rgb[3*i+2]));
		hsv[3*i+1] = s = v-std::min(rgb[3*i+0],std::min(rgb[3*i+1],rgb[3*i+2]));
		if( v == rgb[3*i+0] )
			hsv[3*i+0] = (rgb[3*i+1] - rgb[3*i+2]) / (6.*s+1e-10);
		else if( v == rgb[3*i+1] )
			hsv[3*i+0] = (2+rgb[3*i+2] - rgb[3*i+0]) / (6.*s+1e-10);
		else if( v == rgb[3*i+2] )
			hsv[3*i+0] = (4+rgb[3*i+0] - rgb[3*i+1]) / (6.*s+1e-10);
		if( hsv[3*i+0] < 0 )
			hsv[3*i+0] += 1;
	}
}


// This file is part of MLDB. Copyright 2015 mldb.ai inc. All rights reserved.

/* decision_tree_generator.cc
   Jeremy Barnes, 15 March 2006
   Copyright (c) 2006 Jeremy Barnes  All rights reserved.
   $Source$

   Generator for decision trees.
*/

#include "decision_tree_generator.h"
#include "mldb/plugins/jml/jml/registry.h"
#include "training_index.h"
#include "weighted_training.h"
#include "stump_training_core.h"
#include "stump_training.h"
#include "stump_training_bin.h"
#include "stump_regress.h"
#include "binary_symmetric.h"
#include "mldb/utils/smart_ptr_utils.h"
#include "mldb/base/thread_pool.h"
#include <boost/timer/timer.hpp>

#include <random>
#include <mutex>

using namespace std;

namespace ML {

// Tuning parameters for the compacting code.  This controls how frequently
// we compact.  Compacting makes operations faster, but takes time and
// memory.

/// If there are less than 1 in COMPACT_SPARSENESS_RATIO non-zero entries,
/// then we compact the dataset.
static constexpr int COMPACT_SPARSENESS_RATIO=8;

/// On entry, if we have a density less than COMPACT_INITIAL_DENSITY,
/// we compact the dataset.
static constexpr float COMPACT_INITIAL_DENSITY=0.5;

/// Glue code to allow the templated version of compactDataset to find
/// the first weight, no matter what types are passed in.
float getFirstWeight(const float * f)
{
    return *f;
}

float getFirstWeight(float f)
{
    return f;
}

template<class Weights>
void compact_dataset(const Training_Data & data,
                     const vector<float> & in_class,
                     const Weights & weights,
                     const distribution<float> & binary_weights,
                     int num_non_zero,
                     Training_Data & new_data,
                     vector<float> & new_in_class,
                     Weights & new_weights,
                     distribution<float> & new_binary_weights,
                     const vector<Feature> & features,
                     const Feature & predicted)
{
    /* Compact the dataset, since we only need a fraction of it.  We
       simply recreate all of the strucutres, but with the entries that
       would have been zero not there anymore.

       Note that this can lead to high memory usage, and is not necessary.
       If there are problems with memory, we could disable this code and
       the algorithm would continue to function.
    */

    new_in_class.reserve(num_non_zero);

    if (binary_weights.empty())
        new_weights.reserve(num_non_zero);
    else new_binary_weights.reserve(num_non_zero);

    for (unsigned i = 0;  i < in_class.size();  ++i) {
        // NOTE: here we assume that either all class weights are zero, or
        // none are zero.  That way we only need to look at the first weight.
        if (binary_weights.empty()) {
            if (in_class[i] <= 0.0f || getFirstWeight(weights[i]) == 0.0f) continue;
            new_in_class.push_back(in_class[i]);
            new_weights.push_back(weights[i]);
        }
        else {
            if (in_class[i] <= 0.0f || binary_weights[i] == 0.0f) continue;
            new_in_class.push_back(in_class[i]);
            new_binary_weights.push_back(binary_weights[i]);
        }
    }

    // Create a filtered version of the dataset
    new_data.initFiltered(data, in_class, predicted, features);
}

/*****************************************************************************/
/* DECISION_TREE_GENERATOR                                                   */
/*****************************************************************************/

Decision_Tree_Generator::
Decision_Tree_Generator()
{
    defaults();
}

Decision_Tree_Generator::~Decision_Tree_Generator()
{
}

void
Decision_Tree_Generator::
configure(const Configuration & config, vector<string> & unparsedKeys)
{
    Classifier_Generator::configure(config, unparsedKeys);

    config.findAndRemove(trace, "trace", unparsedKeys);
    config.findAndRemove(max_depth, "max_depth", unparsedKeys);
    config.findAndRemove(update_alg, "update_alg", unparsedKeys);
    config.findAndRemove(random_feature_propn, "random_feature_propn", unparsedKeys);
    config.findAndRemove(verbosity, "verbosity", unparsedKeys);
}

void
Decision_Tree_Generator::
defaults()
{
    Classifier_Generator::defaults();
    trace = 0;
    max_depth = -1;
    update_alg = Stump::PROB;
    random_feature_propn = 1.0;
}

Config_Options
Decision_Tree_Generator::
options() const
{
    Config_Options result = Classifier_Generator::options();
    result
        .add("trace", trace, "0-",
             "trace execution of training in a very fine-grained fashion")
        .add("max_depth", max_depth, "0- or -1",
             "give maximum tree depth.  -1 means go until data separated")
        .add("update_alg", update_alg,
             "select the type of output that the tree gives")
        .add("random_feature_propn", random_feature_propn, "0.0-1.0",
             "proportion of the features to enable (for random forests)");

    return result;
}

void
Decision_Tree_Generator::
init(std::shared_ptr<const Feature_Space> fs, Feature predicted)
{
    Classifier_Generator::init(fs, predicted);
    model = Decision_Tree(fs, predicted);
}

std::shared_ptr<Classifier_Impl>
Decision_Tree_Generator::
generate(Thread_Context & context,
         const Training_Data & training_set,
         const Training_Data & validation_set,
         const distribution<float> & training_ex_weights,
         const distribution<float> & validate_ex_weights,
         const std::vector<Feature> & features, int) const
{
    boost::timer::cpu_timer timer;

    Feature predicted = model.predicted();

    boost::multi_array<float, 2> weights
        = expand_weights(training_set, training_ex_weights, predicted);

    Decision_Tree current
        = train_weighted(context, training_set, weights, features, max_depth);

    if (verbosity > 2)
        cerr << current.print() << endl;

    return std::make_shared<Decision_Tree>(std::move(current));
}

std::shared_ptr<Classifier_Impl>
Decision_Tree_Generator::
generate(Thread_Context & context,
         const Training_Data & training_set,
         const boost::multi_array<float, 2> & weights,
         const std::vector<Feature> & features,
         float & Z,
         int recursion) const
{
    //boost::timer::cpu_timer timer;

    //Feature predicted = model.predicted();

    Decision_Tree current
        = train_weighted(context, training_set, weights, features, max_depth);

    if (verbosity > 2) cerr << current.print() << endl;

    return make_sp(current.make_copy());
}

Decision_Tree
Decision_Tree_Generator::
train_weighted(Thread_Context & context,
               const Training_Data & data,
               const boost::multi_array<float, 2> & weights,
               const std::vector<Feature> & features,
               int max_depth) const
{
    Decision_Tree result = model;

    Feature predicted = model.predicted();

    /* Record which examples are in our class and with what weight they
       are there. */
    distribution<float> in_class(data.example_count(), 1.0);

    int nlw = weights.shape()[1];

    for (unsigned i = 0;  i < data.example_count();  ++i) {
        bool nonZero = false;
        for (unsigned j = 0;  j < nlw && !nonZero;  ++j) {
            nonZero = weights[i][j] != 0;
        }

        in_class[i] = nonZero;
    }

    bool regression_problem
        = result.feature_space()->info(predicted).type() == REAL;

    if (random_feature_propn < 0.0 || random_feature_propn > 1.0)
        throw Exception("random_feature_propn is not between 0.0 and 1.0");

    vector<Feature> filtered_features;
    if (random_feature_propn < 1.0) {

        int iter = 0;
        while (filtered_features.empty() && iter < 50) {
            typedef mt19937 engine_type;
            engine_type engine(context.random());
            std::uniform_real_distribution<> rng(0, 1);

            for (unsigned i = 0;  i < features.size();  ++i) {
                if (rng(engine) < random_feature_propn)
                    filtered_features.push_back(features[i]);
            }
        }

        if (filtered_features.empty())
            throw Exception("random_feature_propn is too low");
    }
    else filtered_features.insert(filtered_features.end(),
                                  features.begin(), features.end());

    if (max_depth == -1)
        max_depth = 50;

    if (regression_problem) {
        vector<float> weights_vec(data.example_count());
        for (unsigned x = 0;  x < weights_vec.size();  ++x)
            weights_vec[x] = weights[x][0];

        result.tree.root = train_recursive_regression
            (context, data, weights_vec, filtered_features, in_class,
             0, max_depth, result.tree);
    }
    else {

        convert_bin_sym(const_cast<boost::multi_array<float, 2> &>(weights),
                        data, predicted, features);

        int advance = get_advance(weights);
        vector<const float *> weights_vec(data.example_count());

        for (unsigned x = 0;  x < weights_vec.size();  ++x)
            weights_vec[x] = &weights[x][0];

        int numNonZero = 0;

        for (unsigned i = 0;  i < weights_vec.size();  ++i) {
            if (weights_vec[i][0] != 0.0)
                ++numNonZero;
        }

        //cerr << "numNonZero = " << numNonZero << endl;

        if (numNonZero < COMPACT_INITIAL_DENSITY * weights_vec.size()) {
            Training_Data new_data(data.feature_space());
            distribution<float> new_in_class;
            vector<const float *> new_weights;
            distribution<float> new_binary_weights;

            compact_dataset(data, in_class, weights_vec, {} /* binary weights */, numNonZero,
                            new_data, new_in_class, new_weights, new_binary_weights,
                            features, model.predicted());

            result.tree.root = train_recursive
                (context, new_data, new_weights, advance, filtered_features,
                 new_in_class, 0, max_depth, result.tree);
        }
        else {
            result.tree.root = train_recursive
                (context, data, weights_vec, advance, filtered_features, in_class,
                 0, max_depth, result.tree);
        }

        result.encoding = Stump::update_to_encoding(update_alg);

        /* Validate that the examples in the training set are indeed split
           the way it says */
        // TODO
    }

    return result;
}

namespace {

/** Structure in which we hold the results of the line search over potential
    split points. */
template<class W, class Z, class Tracer = No_Trace>
struct Tree_Accum {

    Tree_Accum(const Feature_Space & fs, int nl,
               const Tracer & tracer = Tracer())
        : tracer(tracer), best_w(nl),
          best_arg(numeric_limits<float>::quiet_NaN()),
          best_z(Z::worst),
          best_feature(MISSING_FEATURE), fs(fs)
    {
    }

    Tracer tracer;

    Z calc_z;

    W best_w;
    float best_arg;
    std::atomic<float> best_z;  // read without the lock and strictly decrementing
    Feature best_feature;

    bool has_result() const { return best_feature != MISSING_FEATURE; }

    std::mutex lock;

    const Feature_Space & fs;

    /** Method that gets called when we start a new feature.  We use it to
        pre-cache part of the work from the Z calculation, as we are
        assured that the MISSING buckets of W will never change after this
        method is called.

        Return value is used to allow an early exit from the training process,
        due to it being impossible for this feature to have a high enough
        value to be included.
    */
    bool start(const Feature & feature, const W & w, double & missing)
    {
        bool optional = fs.info(feature).optional();
        missing = calc_z.missing(w, optional);
        bool keep_going = calc_z.can_beat(w, missing, best_z);

        return keep_going;
    }

    /** Method that gets called when we have found a potential split point. */
    float add_z(const Feature & feature, const W & w, float arg, float z)
    {
        if (false) {
        // Check that the dataset was split evenly enough, ie that at least
        // 10% of the data is in one bucket

            double w_true = w(0,true,0) + w(0,true,1);
            double w_false = w(0,false,0) + w(0,false,1);
            double w_missing = w(0,MISSING,0) + w(0,MISSING,1);
            double w_total = (w_true + w_false + w_missing);
            w_true /= w_total;  w_false /= w_total;  w_missing /= w_total;
            double threshold = 0.2;

            int n = (w_true > threshold)
                + (w_false > threshold)
                + (w_missing > threshold);

            if (n < 2) { z += (1 - z) * 0.9; };
        }

        bool print_feat = false;
        //print_feat = fs.print(feature) == "language_cosine";
        if (tracer || print_feat)
            tracer("tree accum", 3)
                << "  accum: feature " << feature << " arg " << arg
                << " (" << fs.print(feature, arg)
                << "; 0x" << format("%08x", reinterpret_as_int(arg))
                << ") z " << z << "  " << fs.print(feature)
                << (z < best_z ? " ****" : "")
                << endl;

        if (tracer || print_feat)
            tracer("tree accum", 4) << w.print() << endl;

        auto isBetter = [&] ()
        {
            // Break ties in a consistent manner: first compare z, then feature, then argument
            // so that the output is always deterministic
            return less_all(z, best_z.load(), feature, best_feature, arg, best_arg);
        };

        if (z <= best_z) {  // first check only checks z to avoid taking the lock
            std::unique_lock<std::mutex> guard(lock);

            if (isBetter()) {

#if 0 // MLDB-784... sparse features can have a non-finite split
                if (!isfinite(arg)) {  // will never happen
                    static std::mutex mutex;
                    std::unique_lock<std::mutex> guard(mutex);

                    cerr << "Best arg had non-finite split" << endl;
                    cerr << "feature = " << fs.print(feature) << endl;
                    cerr << "info = " << fs.info(feature) << endl;
                    cerr << "z = " << z << endl;
                    cerr << "arg = " << arg << endl;

                    cerr << "Beating previous best" << endl;
                    cerr << "feature = " << fs.print(best_feature) << endl;
                    cerr << "info = " << fs.info(best_feature) << endl;
                    cerr << "z = " << best_z << endl;
                    cerr << "arg = " << best_arg << endl;
                }
#endif // MLDB-784

                if (tracer || print_feat)
                    tracer("tree accum", 4) << w.print() << endl;
                // A better one.  This replaces whatever we had accumulated so
                // far.
                best_z = z;
                best_w = w;
                best_arg = arg;
                best_feature = feature;
            }
        }

        return z;
    }

    float add(const Feature & feature, const W & w, float arg, double missing)
    {
        // If the decision tree generator is having a really tough time
        // separating the classes, and it's a bucketed feature, than it
        // may send back a -INF for arg (which means split on missing or
        // not missing), even if there is no missing feature, due to
        // numerical issues.  Since the decision tree can't handle a split
        // point of -INIFINITY, we return that we don't want this split
        // point so that it will continue looking for something better.

        //if (!isfinite(arg)) {
        //    return Z::none;
        //}

        float z = calc_z.non_missing(w, missing);
        return add_z(feature, w, arg, z);
    }

    float add_presence(const Feature & feature, const W & w, float arg,
                       double missing)
    {
        float z = calc_z.non_missing_presence(w, missing);
        return add_z(feature, w, arg, z);
    }

    void finish(const Feature & feature)
    {
        // nothing to do here, at the moment
    }

    Split split()
    {
        if (!has_result()) return Split();
        Split result(best_feature, best_arg, fs);
        return result;
    }

    double z() const
    {
        return best_z;
    }
};

template<class W>
void
get_probs(distribution<float> & probs, const W & w_,
          Stump::Update update, float epsilon = 0.0)
{
#if 1
    W w = w_;
    for (unsigned j = 1;  j < 3;  ++j) {
        for (unsigned l = 0;  l < w.nl();  ++l) {
            w(l, 0, true) += w(l, j, true);
            w(l, 0, false) += w(l, j, false);
        }
    }

    probs.resize(w.nl());

    //cerr << "get_probs(): w_ = " << endl << w_.print() << endl << "w = " << endl
    //     << w.print() << endl << " probs = " << probs << endl;

    C_any c(update);
    return c(probs.data(), 0, w, epsilon, false);

#else
    distribution<float> result(w.nl());

    for (unsigned j = 0;  j < 3;  ++j)
        for (unsigned l = 0;  l < w.nl();  ++l)
            result[l] += w(l, j, true);

    result.normalize();
    return result;
#endif
}

void
fillin_leaf(Tree::Leaf & leaf,
            const Training_Data & data,
            const Feature & predicted,
            const vector<const float *> & weights,
            int advance,
            const distribution<float> & in_class,
            Stump::Update update_alg,
            float examples = -1.0)
{
    /* Use the stump trainer to accumulate for us. */

    if (examples == -1.0) examples = in_class.total();

    //cerr << "new_leaf: advance = " << advance << endl;

    if (advance < 0)
        throw Exception("invalid advance");

    leaf.examples = examples;

    if (advance != 0) {
        typedef W_normal W;
        typedef Z_normal Z;
        typedef Stump_Trainer<W, Z> Trainer;
        Trainer trainer;

        W w = trainer.calc_default_w(data, predicted, in_class, weights, advance);

        double epsilon = xdiv<double>(1.0, examples);
        get_probs(leaf.pred, w, update_alg, epsilon);

        //cerr << "new_leaf norm: dist = " << dist << " W = " << endl
        //     << w.print() << endl;
    }
    else {
        typedef W_binsym W;
        typedef Z_binsym Z;
        typedef Stump_Trainer<W, Z> Trainer;
        Trainer trainer;

        //cerr << "getting default W" << endl;

        W w = trainer.calc_default_w(data, predicted, in_class, weights, advance);

        double epsilon = xdiv<double>(1.0, examples);
        get_probs(leaf.pred, w, update_alg, epsilon);

        //cerr << "new_leaf binsym: dist = " << dist << " W = " << endl
        //     << w.print() << endl;
    }
}

void
fillin_leaf_regression(Tree::Leaf & leaf,
                       const Training_Data & data,
                       const Feature & predicted,
                       const vector<float> & weights,
                       const distribution<float> & in_class,
                       float examples = -1.0)
{
    /* Calculate the weighted mean over the examples in this class. */
    int nx = data.example_count();

    double total_weight = 0.0;
    double total_val = 0.0;

    const vector<Label> & labels = data.index().labels(predicted);

    for (unsigned x = 0;  x < nx;  ++x) {
        float w = weights[x] * in_class[x];
        float val = labels[x].value();
        total_weight += w;
        total_val += w * val;
    }

    distribution<float> dist(1, total_val / total_weight);

    if (examples == -1.0) examples = in_class.total();

    leaf.pred = dist;
    leaf.examples = examples;
}

void split_dataset(const Training_Data & data,
                   const Split & split,
                   const distribution<float> & in_class,
                   distribution<float> & class_true,
                   distribution<float> & class_false,
                   distribution<float> & class_missing,
                   double & total_true,
                   double & total_false,
                   double & total_missing,
                   bool validate)
{
    // TODO: we could use the index for the feature instead?

    /* Split these examples based upon what the split said. */
    int nx = data.example_count();

    class_true = distribution<float>(nx);
    class_false = distribution<float>(nx);
    class_missing = distribution<float>(nx);

    total_true = 0.0;
    total_false = 0.0;
    total_missing = 0.0;

    Joint_Index index = data.index().dist(split.feature(), BY_EXAMPLE,
                                          IC_VALUE | IC_DIVISOR | IC_EXAMPLE);

    int last_example = -1;
    for (unsigned i = 0;  i < index.size();  ++i) {
        int example = index[i].example();

        // Any which we skipped over are missing
        for (++last_example; last_example < example; ++last_example) {
            float w = in_class[last_example];
            if (w == 0.0) continue;
            //cerr << "missing: example = " << example << " last_example = "
            //     << last_example << endl;
            class_missing[last_example] += w;
            total_missing += w;
        }

        last_example = example;

        float w = in_class[example];
        if (w == 0.0) continue;

        float val = index[i].value();
        float divisor = index[i].divisor();
        w *= divisor;

        if (MLDB_UNLIKELY(isnan(val))) {
            // We only have NaN values explicitly represented if there is a
            // feature that is both present and missing in the same example.
            // In that case, we need to deal with the missing part here.
            class_missing[example] += w;
            total_missing += w;
            continue;
        }

        int decision;
        try {
            decision = split.apply(val);
        } catch (...) {
            cerr << "exception on split: " << split.print(*data.feature_space())
                 << endl;
            throw;
        }

        switch(decision) {
        case false:
            class_false[example] += w;
            total_false += w;
            break;
        case true:
            class_true[example] += w;
            total_true += w;
            break;
        case MISSING:
            class_missing[example] += w;
            total_missing += w;
            break;
        default:
            throw Exception("split_dataset: bad decision");
        };
    }

    // Any examples we never touched are also missing
    for (++last_example; last_example < nx;  ++last_example) {
        float w = in_class[last_example];
        if (w == 0.0) continue;
        class_missing[last_example] += w;
        total_missing += w;
    }

    /* For validation: make sure that each is in exactly one */
    if (validate) {
        for (unsigned x = 0;  x < nx;  ++x) {
            double w_total = class_true[x] + class_false[x] + class_missing[x];
            double error = in_class[x] - w_total;
            if (abs(error) > 0.000001) {
                cerr << "x = " << x << endl;
                cerr << "orig  = " << in_class[x] << endl;
                cerr << "false = " << class_false[x] << endl;
                cerr << "true  = " << class_true[x] << endl;
                cerr << "miss  = " << class_missing[x] << endl;
                cerr << "total = " << w_total << endl;
                cerr << "error = " << error << endl;
                throw Exception("split_weights: weights don't add up");
            }
        }
    }
}

/// The training code expects weights to be double-indexed, by first example
/// number and second label, so you can do weights[example][label] to extract
/// the label.  In the optimization for the binary symmetric case (where we
/// always have 2 identical weights, one for each label) we have all of the
/// weights in a single array.  This code adds an extra dimension to the
/// weights, allowing us to call AddDimension(weights)[example][label] and
/// have it return weights[example].
struct AddDimension {
    AddDimension(const distribution<float> & vals)
        : vals(vals)
    {
    }

    const distribution<float> & vals;

    const float * operator [] (int n) const { return &vals[n]; };
};

int get_advance(const AddDimension &)
{
    return 0;
}

template<typename W, typename Z>
struct TreeTrainer {
    Tree & tree;
    int max_depth;
    const vector<Feature> & features;
    int advance;
    int nl;
    Feature predicted;
    int trace;
    Stump::Update update_alg;
    std::shared_ptr<const Feature_Space> feature_space;
    bool validate;

    TreeTrainer(Tree & tree,
                int max_depth,
                const vector<Feature> & features,
                int advance,
                Feature predicted,
                int trace,
                Stump::Update update_alg,
                std::shared_ptr<const Feature_Space> feature_space,
                bool validate)
        : tree(tree), max_depth(max_depth),
          features(features), advance(advance),
          predicted(predicted), trace(trace),
          update_alg(update_alg),
          feature_space(feature_space),
          validate(validate)
    {
        nl = feature_space->info(predicted).value_count();
    }

    typedef Stump_Trainer<W, Z> WeightTrainer;

    typedef No_Trace TrainerTracer;
    typedef Tree_Accum<W, Z, Stream_Tracer> Accum;
    typedef Stump_Trainer<W, Z, TrainerTracer> SplitTrainer;

    void do_branch(Tree::Ptr & ptr,
                   Thread_Context & context,
                   const Training_Data & data,
                   const vector<const float *> & weights,
                   const distribution<float> & binary_weights,
                   int advance,
                   const vector<Feature> & features,
                   const distribution<float> & new_in_class,
                   double total_in_class,
                   int new_depth, int max_depth,
                   Tree & tree,
                   MLDB::ThreadPool & tp) const
    {
#if 0
        if (total_in_class > 1024) {
            // Worth multithreading... do it
            if (group_to_wait_on == -1) {
                // Create a new group
                group_to_wait_on = context.worker().get_group(NO_JOB,
                                                              "decision tree",
                                                              context.group());
            }
            Thread_Context child_context = context.child(group_to_wait_on);

            Train_Recursive_Job job(ptr, this, child_context, data, weights,
                                    advance,
                                    features, new_in_class, new_depth, max_depth,
                                    tree);

            context.worker().add(job, "train decision tree branch",
                                 child_context.group());
        }
#else
        if (false) ;
#endif
        else if (total_in_class > 0.0)
            ptr = this->train(context, data, weights, new_in_class, binary_weights,
                              new_depth);
        else {
            // Leaf only
            ptr = tree.new_leaf();
            fillin_leaf(*ptr.leaf(), data, predicted, weights,
                        advance, new_in_class, update_alg, 0.0);
        }
    }

    // This structure needs to be defererenceable and incrementable, and always
    // return 1 when dereferenced.  It's like an iterator into an infinite array
    // of ones.
    struct AlwaysOne {
        float operator * () const { return 1.0f; }
        void operator ++ () {}
    };

    Tree::Ptr
    train(Thread_Context & context,
          const Training_Data & data,
          const std::vector<const float *> & weights,
          const distribution<float> & in_class,
          const distribution<float> & binary_weights,
          int depth) const
    {
        bool debug = false;

#if 0
        int numNonZero = 0;
        for (unsigned i = 0;  i <  in_class.size();  ++i) {
            if (weights[i][0] * in_class[i] != 0)
                ++numNonZero;
            if (i < 10 && false) {
                cerr << i << " p " << weights[i] << " weight " << weights[i][0]
                     << " in_class " << in_class[i] << " total " << weights[i][0] * in_class[i] << endl;
            }
        }
        //cerr << "in class: " << numNonZero << " of " << in_class.size() << endl;
#endif

#if 0
        if (numNonZero * 20 < in_class.size()) {
            cerr << "warning: extremely sparse examples: " << numNonZero << " of "
                 << in_class.size() << endl;
        }
#endif

        if (depth > 100 && max_depth == -1)
            throw Exception("Decision_Tree_Generator::train_recursive(): "
                            "depth of 100 reached");
        if (debug)
            cerr << "train_recursive: depth " << depth << endl;

        double total_weight = in_class.total();

#if 0
        if (debug) {
            cerr << "predicted = " << predicted << endl;
            cerr << "fs = " << data.feature_space()->print() << endl;
            cerr << "data[0] = " << data.feature_space()->print(data[0]) << endl;
            cerr << "data.example_count() = " << data.example_count() << endl;

            cerr << "data.label_count(predicted) = " << data.label_count(predicted)
                 << endl;
            cerr << "data.label_count(model.predicted()) = "
                 << data.label_count(model.predicted())
                 << endl;
        }
#endif

        W default_w(nl);

        /* Check for zero impurity, and return a leaf if we have it. */
        double class_weights[nl];

        // What would we have as a leaf if we were to stop splitting here?
        Tree::Leaf leaf;
        leaf.examples = total_weight;
        double epsilon = xdiv<double>(1.0, total_weight);

        /* Use the stump trainer to accumulate for us. */
        WeightTrainer weightTrainer;

        if (binary_weights.empty()) {
            default_w = weightTrainer.calc_default_w(data, predicted, in_class, weights, advance);
        }
        else default_w = weightTrainer.calc_default_w(data, predicted, in_class, AddDimension(binary_weights),
                                                      advance);

        // Calculate how many classes (labels) have non-zero weight.  We need at least
        // 2 distinct labels for training to make sense; otherwise we bail out

        double maxClassWeight = 0.0;
        double totalClassWeight = 0.0;
        int numNonZeroClasses = 0;
        for (unsigned l = 0;  l < nl;  ++l) {
            double w = 0.0;  // weight for the label class
            for (unsigned j = 0;  j < 3;  ++j)
                w += default_w(l, j, true);
            class_weights[l] = w;
            maxClassWeight = std::max(maxClassWeight, w);
            totalClassWeight += w;
            if (w != 0)
                numNonZeroClasses += 1;
        }

        // Fill in the prediction that this would have made if we stopped
        // training here.  This is used, notably, if we prune the tree and
        // for explanations.
        get_probs(leaf.pred, default_w, update_alg, epsilon);

        // Look for early stopping conditions:
        if (maxClassWeight == 1.0       // minimum impurity; one per class
            || depth == max_depth       // reached maximum depth
            || numNonZeroClasses <= 1   // only one non-zero weighted label left
            || total_weight < 1.0       // split up finer than one example
            || totalClassWeight == 0.0  // weights too small to count
            || in_class.size() == 1     // only one example
            || false) {
            Tree::Leaf * result = tree.new_leaf();
            *result = leaf;
            return result;
        }

        int num_non_zero = std::count_if(in_class.begin(), in_class.end(), [] (auto f) { return f > 0.0; });

        if (debug) {
            cerr << "in_class.size() = " << in_class.size() << " num_non_zero = "
                 << num_non_zero << " total_weight = " << total_weight
                 << endl;
        }

        if (num_non_zero * COMPACT_SPARSENESS_RATIO < in_class.size()) {
            Training_Data new_data(data.feature_space());
            distribution<float> new_in_class;
            distribution<float> new_binary_weights;
            vector<const float *> new_weights;

            compact_dataset(data, in_class, weights, binary_weights, num_non_zero,
                            new_data, new_in_class,
                            new_weights,
                            new_binary_weights,
                            features, predicted);

            /* Restart, with the new training data. */
            return train(context, new_data, new_weights,
                         new_in_class, new_binary_weights, depth);
        }

        Split split;
        float best_z = 0.0;

        Accum accum(*feature_space, nl, trace);
        SplitTrainer splitTrainer;

        if (binary_weights.empty())
            splitTrainer.test_all
                (context, default_w, features, data, predicted,
                 weights, in_class, accum, advance);
        else
            splitTrainer.test_all
                (context, default_w, features, data, predicted,
                 AddDimension(binary_weights), in_class, accum, advance);

        split = accum.split();
        best_z = accum.z();

        if (split.feature() == MISSING_FEATURE) {
            Tree::Leaf * result = tree.new_leaf();
            *result = leaf;
            return result;
        }

        if (debug) {
            cerr << " decision tree training: best split is "
                 << split.print(*feature_space) << endl;
            cerr << "z = " << best_z << endl;
        }

        // We used to not allow the decision tree to learn a perfect split.  Now we allow it
        // but we make sure that the next level down only leaf nodes will be created as there
        // will be only one label.
        if (best_z == 0.0 && false) {
            // No impurity at all
            Tree::Leaf * result = tree.new_leaf();
            *result = leaf;
            return result;
        }

        /* Split these examples based upon what the split said. */
        distribution<float> class_true;
        distribution<float> class_false;
        distribution<float> class_missing;
        double total_true;
        double total_false;
        double total_missing;

        //boost::timer::cpu_timer timer;
        split_dataset(data, split, in_class,
                      class_true, class_false, class_missing,
                      total_true, total_false, total_missing,
                      validate);

        int numClasses = (total_true != 0) + (total_false != 0) + (total_missing != 0);

        // If we classify everything into one class, then we can't split any
        // further.
        if (numClasses < 2) {
            Tree::Leaf * result = tree.new_leaf();
            *result = leaf;
            return result;
        }

        if (debug) {
            //cerr << timer.elapsed().wall << "s split" << endl;

            cerr << " totals: true " << total_true << " false " << total_false
                 << " missing " << total_missing << endl;

            cerr << " totals2: true " << class_true.total()
                 << " false " << class_false.total()
                 << " missing " << class_missing.total() << endl;
        }

        Tree::Node * node = tree.new_node();
        node->split = split;
        node->z = best_z;
        node->examples = total_weight;
        node->pred = leaf.pred;

        MLDB::ThreadPool tp;

        do_branch(node->child_true,
                  context, data, weights, binary_weights, advance, features,
                  class_true, total_true, depth + 1, max_depth,
                  tree, tp);

        do_branch(node->child_false,
                  context, data, weights, binary_weights, advance, features,
                  class_false, total_false, depth + 1, max_depth,
                  tree, tp);

        do_branch(node->child_missing,
                  context, data, weights, binary_weights, advance, features,
                  class_missing, total_missing, depth + 1, max_depth,
                  tree, tp);

        tp.waitForAll();

        return node;
    }
};

} // file scope

struct Decision_Tree_Generator::Train_Recursive_Job {

    Tree::Ptr & ptr;
    const Decision_Tree_Generator * generator;
    Thread_Context context;
    const Training_Data & data;
    const vector<const float *> & weights;
    int advance;
    const vector<Feature> & features;
    const distribution<float> & in_class;
    int depth;
    int max_depth;
    Tree & tree;

    Train_Recursive_Job(Tree::Ptr & ptr,
                        const Decision_Tree_Generator * generator,
                        const Thread_Context & context,
                        const Training_Data & data,
                        const vector<const float *> & weights,
                        int advance,
                        const vector<Feature> & features,
                        const distribution<float> & in_class,
                        int depth, int max_depth,
                        Tree & tree)
        : ptr(ptr), generator(generator), context(context), data(data),
          weights(weights), advance(advance), features(features),
          in_class(in_class), depth(depth), max_depth(max_depth),
          tree(tree)
    {
    }

    void operator () ()
    {
        ptr = generator->train_recursive(context, data, weights, advance,
                                         features, in_class, depth,
                                         max_depth, tree);
    }
};

Tree::Ptr
Decision_Tree_Generator::
train_recursive(Thread_Context & context,
                const Training_Data & data,
                const vector<const float *> & weights,
                int advance,
                const vector<Feature> & features,
                const distribution<float> & in_class,
                int depth, int max_depth,
                Tree & tree) const
{
    if (advance == 0) {
        // binary symmetric, we can use an optimized version
        typedef W_binsym W;
        typedef Z_binsym Z;

        // Pre-calculate a more memory friendly version of in_class
        distribution<float> binary_weights(in_class.size(), 0.0);
        for (unsigned i = 0;  i < in_class.size();  ++i) {
            binary_weights[i] = weights[i][0];
        }

        TreeTrainer<W, Z> trainer(tree, max_depth, features, advance,
                                  predicted, trace, update_alg, feature_space,
                                  validate);
        return trainer.train(context, data, weights, in_class, binary_weights, depth);
    }
    else {
        // generalized version
        typedef W_normal W;
        typedef Z_normal Z;
        TreeTrainer<W, Z> trainer(tree, max_depth, features, advance,
                                  predicted, trace, update_alg, feature_space,
                                  validate);
        return trainer.train(context, data, weights, in_class, {} /* binary weights */, depth);
    }
}

Tree::Ptr
Decision_Tree_Generator::
train_recursive_regression(Thread_Context & context,
                           const Training_Data & data,
                           const vector<float> & weights,
                           const vector<Feature> & features_,
                           const distribution<float> & in_class,
                           int depth, int max_depth,
                           Tree & tree) const
{
    if (depth > 100 && max_depth == -1)
        throw Exception("Decision_Tree_Generator::train_recursive_regression(): "
                        "depth of 100 reached");

    size_t nx = data.example_count();

    Tree::Leaf leaf;
    fillin_leaf_regression(leaf, data, model.predicted(), weights, in_class);

    const vector<Label> & labels = data.index().labels(model.predicted());

    /* Check for all of the labels in the class having the same value. */
    {
        float val_found = NAN;
        bool all_same = true;
        for (unsigned x = 0;  x < nx;  ++x) {
            if (in_class[x] == 0.0) continue;  // not in our class
            if (weights[x] == 0.0) continue;   // no weight; doesn't count
            if (isnan(val_found)) val_found = labels[x].value();
            else if (val_found != labels[x].value()) {
                all_same = false;
                break;
            }
        }

        if (all_same) {
            Tree::Leaf * result = tree.new_leaf();
            *result = leaf;
            return result;
        }
    }

    double total_weight = in_class.total();

    if (depth == max_depth || total_weight < 1.0) {
        Tree::Leaf * result = tree.new_leaf();
        *result = leaf;
        return result;
    }

    int num_non_zero = std::count_if(in_class.begin(), in_class.end(), [] (auto f) { return f > 0.0; });

    //cerr << "in_class.size() = " << in_class.size() << " num_non_zero = "
    //     << num_non_zero << " total_weight = " << total_weight
    //     << endl;

    if (num_non_zero * 16 < in_class.size()) {
        Training_Data new_data(data.feature_space());
        distribution<float> new_in_class;
        vector<float> new_weights;
        distribution<float> new_binary_weights;

        compact_dataset(data, in_class, weights, {} /* binary weights */, num_non_zero,
                        new_data, new_in_class, new_weights, new_binary_weights, features_,
                        model.predicted());

        /* Restart, with the new training data. */
        return train_recursive_regression
            (context, new_data, new_weights, features_, new_in_class, depth,
             max_depth, tree);
    }

    //cerr << "training decision tree with total weight "
    //     << total_weight << " at depth " << depth << endl;

    typedef W_regress W;
    typedef Z_regress Z;

    typedef Tree_Accum<W, Z, Stream_Tracer> Accum;
    //typedef Tree_Accum<W, Z> Accum;
    typedef Stump_Trainer<W, Z> Trainer;

    Accum accum(*model.feature_space(), nl, trace);
    Trainer trainer;

    vector<Feature> features = features_;

    /* We need it in a fixed array like this. */
    boost::multi_array<float, 2> weights2(boost::extents[weights.size()][1]);
    std::copy(weights.begin(), weights.end(), weights2.data());
    trainer.test_all_and_sort(features, data, model.predicted(), weights2,
                              in_class, accum);

    //cerr << "z = " << accum.best_z << endl;
    //cerr << "w = " << endl << accum.best_w.print() << endl;

    if (accum.best_z == 0.0 || accum.best_feature == MISSING_FEATURE) {
        Tree::Leaf * result = tree.new_leaf();
        *result = leaf;
        return result;
    }

    distribution<float> class_true;
    distribution<float> class_false;
    distribution<float> class_missing;

    double total_true;
    double total_false;
    double total_missing;

    //boost::timer::cpu_timer timer;
    split_dataset(data, accum.split(), in_class,
                  class_true, class_false, class_missing,
                  total_true, total_false, total_missing,
                  validate);

    //cerr << timer.elapsed().wall << "s split" << endl;

    //cerr << " totals: true " << total_true << " false " << total_false
    //     << " missing " << total_missing << endl;

    Tree::Node * node = tree.new_node();
    node->split = accum.split();
    node->z = accum.z();
    node->examples = total_weight;
    node->pred = leaf.pred;

    node->child_true
        = train_recursive_regression(context, data, weights, features,
                                     class_true, depth + 1, max_depth,
                                     tree);
    node->child_false
        = train_recursive_regression(context, data, weights, features,
                                     class_false, depth + 1, max_depth,
                                     tree);
    node->child_missing
        = train_recursive_regression(context, data, weights, features,
                                     class_missing, depth + 1, max_depth,
                                     tree);

    return node;
}

/*****************************************************************************/
/* REGISTRATION                                                              */
/*****************************************************************************/

namespace {

Register_Factory<Classifier_Generator, Decision_Tree_Generator>
    DECISION_TREE_REGISTER("decision_tree");

} // file scope

} // namespace ML


/**********************************************************************************/
/* This file is part of spla project                                              */
/* https://github.com/JetBrains-Research/spla                                     */
/**********************************************************************************/
/* MIT License                                                                    */
/*                                                                                */
/* Copyright (c) 2021 JetBrains-Research                                          */
/*                                                                                */
/* Permission is hereby granted, free of charge, to any person obtaining a copy   */
/* of this software and associated documentation files (the "Software"), to deal  */
/* in the Software without restriction, including without limitation the rights   */
/* to use, copy, modify, merge, publish, distribute, sublicense, and/or sell      */
/* copies of the Software, and to permit persons to whom the Software is          */
/* furnished to do so, subject to the following conditions:                       */
/*                                                                                */
/* The above copyright notice and this permission notice shall be included in all */
/* copies or substantial portions of the Software.                                */
/*                                                                                */
/* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR     */
/* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,       */
/* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE    */
/* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER         */
/* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,  */
/* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE  */
/* SOFTWARE.                                                                      */
/**********************************************************************************/
#ifndef SPLA_SPLAREDUCE_HPP
#define SPLA_SPLAREDUCE_HPP

#include <boost/compute/algorithm/reduce.hpp>

#include <compute/metautil/SplaMetaUtil.hpp>

namespace spla {

    namespace detail {

        inline void InplaceReduce(const boost::compute::vector<unsigned char> &values,
                                  std::size_t valueByteSize,
                                  const std::string &functionBody,
                                  boost::compute::command_queue &queue) {
            std::size_t inputSize = values.size() / valueByteSize;
            if (inputSize < 2) {
                return;
            }

            const compute::context &context = queue.get_context();

            const std::size_t blockSize = 64;
            const std::size_t valuesPerThread = 8;
            std::size_t blockCount = inputSize / (blockSize * valuesPerThread);
            if (blockCount * blockSize * valuesPerThread != inputSize)
                blockCount++;

            compute::vector<unsigned char> output(blockCount * valueByteSize, context);

            compute::detail::meta_kernel k("inplace_reduce");
            size_t inputArg = k.add_arg<unsigned char *>(compute::memory_object::global_memory, "input");
            size_t inputSizeArg = k.add_arg<const uint_>("input_size");
            size_t outputArg = k.add_arg<unsigned char *>(compute::memory_object::global_memory, "output");
            size_t scratchArg = k.add_arg<unsigned char *>(compute::memory_object::local_memory, "scratch");

            ReduceOp reduceOp1(k, "inplace_reduce_reduce_1", functionBody, valueByteSize,
                               Visibility::Unspecified,
                               Visibility::Global,
                               Visibility::Unspecified);

            ReduceOp reduceOp2(k, "inplace_reduce_reduce_2", functionBody, valueByteSize,
                               Visibility::Local,
                               Visibility::Local,
                               Visibility::Unspecified);

            k << "const uint gid = get_global_id(0);\n"
              << "const uint lid = get_local_id(0);\n"
              << "const uint values_per_thread = " << uint_(valuesPerThread) << ";\n"
              << "const uint index = gid * values_per_thread;\n"
              << "if(index < input_size){\n"
              << DeclareVal{"sum", valueByteSize} << ";\n"
              << AssignVal{ValVar{"sum"}, ValArrItem{"input", "index", valueByteSize}, valueByteSize} << ";\n"
              << "for(uint i = 1; i < values_per_thread && (index + i) < input_size; i++) {\n"
              << reduceOp1.Apply(ValVar{"sum"}, ValArrItem{"input", "index+i", valueByteSize}, ValVar{"sum"}) << ";\n"
              << "}\n"
              << AssignVal{ValArrItem("scratch", "lid", valueByteSize), ValVar("sum"), valueByteSize} << ";\n"
              << "}\n"
              << "for(uint i = 1; i < get_local_size(0); i <<= 1){\n"
              << "    barrier(CLK_LOCAL_MEM_FENCE);\n"
              << "    uint mask = (i << 1) - 1;\n"
              << "    uint next_index = (gid + i) * values_per_thread;\n"
                 "    if((lid & mask) == 0 && next_index < input_size){\n"
              << reduceOp2.Apply(ValArrItem("scratch", "lid", valueByteSize), ValArrItem("scratch", "lid+i", valueByteSize), ValArrItem("scratch", "lid", valueByteSize))
              << "    }\n"
              << "}\n"
              << "if(lid == 0){\n"
              << AssignVal{ValArrItem("output", "get_group_id(0)", valueByteSize), ValArrItem("scratch", "0", valueByteSize), valueByteSize} << ";\n"
              << "}\n";

            const compute::buffer *inputBuffer = &values.get_buffer();
            const compute::buffer *outputBuffer = &output.get_buffer();

            compute::kernel kernel = k.compile(context);

            while (inputSize > 1) {
                kernel.set_arg(inputArg, *inputBuffer);
                kernel.set_arg(inputSizeArg, static_cast<uint_>(inputSize));
                kernel.set_arg(outputArg, *outputBuffer);
                kernel.set_arg(scratchArg, compute::local_buffer<unsigned char>(blockSize * valueByteSize));

                queue.enqueue_1d_range_kernel(kernel,
                                              0,
                                              blockCount * blockSize,
                                              blockSize);

                inputSize = static_cast<std::size_t>(std::ceil(static_cast<float>(inputSize) / (blockSize * valuesPerThread)));

                blockCount = inputSize / (blockSize * valuesPerThread);
                if (blockCount * blockSize * valuesPerThread != inputSize)
                    blockCount++;

                std::swap(inputBuffer, outputBuffer);
            }

            if (inputBuffer != &values.get_buffer()) {
                ::boost::compute::copy(output.begin(),
                                       output.begin() + static_cast<std::ptrdiff_t>(valueByteSize),
                                       values.begin(),
                                       queue);
            }
        }

        inline std::size_t Reduce(const boost::compute::vector<unsigned char> &values,
                                  std::size_t valueByteSize,
                                  boost::compute::vector<unsigned char> &result,
                                  std::size_t blockSize,
                                  const std::string &functionBody,
                                  boost::compute::command_queue &queue) {
            using namespace boost;

            const compute::context &context = queue.get_context();
            const std::size_t nValues = values.size() / valueByteSize;
            const std::size_t blockCount = nValues / 2 / blockSize;
            const auto totalBlockCount = static_cast<std::size_t>(
                    std::ceil(static_cast<float>(nValues) / 2.f / static_cast<float>(blockSize)));

            if (blockCount != 0) {
                compute::detail::meta_kernel k("block_reduce");
                std::size_t outputArg = k.add_arg<unsigned char *>(compute::memory_object::global_memory, "output");
                std::size_t blockArg = k.add_arg<unsigned char *>(compute::memory_object::local_memory, "block");

                ReduceOp reduceOpGlobal(k, "block_reduce_reduce_global", functionBody, valueByteSize,
                                        Visibility::Global,
                                        Visibility::Global,
                                        Visibility::Local);
                ReduceOp reduceOpLocal(k, "block_reduce_reduce_local", functionBody, valueByteSize,
                                       Visibility::Local,
                                       Visibility::Local,
                                       Visibility::Unspecified);

                k << "const uint gid = get_global_id(0);\n"
                  << "const uint lid = get_local_id(0);\n"
                  << reduceOpGlobal.Apply(ValArrItem(values, "gid*2+0", valueByteSize, k),
                                          ValArrItem(values, "gid*2+1", valueByteSize, k),
                                          ValArrItem("block", "lid", valueByteSize))
                  << "for(uint i = 1; i < " << uint_(blockSize) << "; i <<= 1){\n"
                  << "    barrier(CLK_LOCAL_MEM_FENCE);\n"
                  << "    uint mask = (i << 1) - 1;\n"
                  << "    if((lid & mask) == 0){\n"
                  << reduceOpLocal.Apply(ValArrItem("block", "lid", valueByteSize),
                                         ValArrItem("block", "lid+i", valueByteSize),
                                         ValArrItem("block", "lid", valueByteSize))
                  << "    }\n"
                  << "}\n"
                  << "if(lid == 0)\n"
                  << AssignVal{ValArrItem{"output", "get_group_id(0)", valueByteSize},
                               ValArrItem{"block", "0", valueByteSize},
                               valueByteSize}
                  << ";\n";

                compute::kernel kernel = k.compile(context);
                kernel.set_arg(outputArg, result.get_buffer());
                kernel.set_arg(blockArg, compute::local_buffer<unsigned char>(blockSize * valueByteSize));

                queue.enqueue_1d_range_kernel(kernel,
                                              0,
                                              blockCount * blockSize,
                                              blockSize);
            }

            // serially reduce any leftovers
            if (blockCount * blockSize * 2 < nValues) {
                std::size_t lastBlockStart = blockCount * blockSize * 2;

                compute::detail::meta_kernel k("extra_serial_reduce");
                const std::size_t countArg = k.add_arg<uint_>("count");
                const std::size_t offsetArg = k.add_arg<uint_>("offset");
                const std::size_t outputArg = k.add_arg<unsigned char *>(compute::memory_object::global_memory, "output");
                const std::size_t outputOffsetArg = k.add_arg<uint_>("output_offset");

                ReduceOp reduceOp(k, "leftover_reduce_reduce", functionBody, valueByteSize,
                                  Visibility::Unspecified,
                                  Visibility::Global,
                                  Visibility::Unspecified);

                k << DeclareVal{"result", valueByteSize} << ";\n"
                  << AssignVal{ValVar{"result"}, ValArrItem{values, "offset", valueByteSize, k}, valueByteSize}
                  << "for(uint i = offset + 1; i < count; i++)\n"
                  << reduceOp.Apply(ValVar{"result"},
                                    ValArrItem(values, "i", valueByteSize, k),
                                    ValVar{"result"})
                  << AssignVal{ValArrItem{"output", "output_offset", valueByteSize},
                               ValVar{"result"},
                               valueByteSize}
                  << ";\n";

                compute::kernel kernel = k.compile(context);
                kernel.set_arg(countArg, static_cast<uint_>(nValues));
                kernel.set_arg(offsetArg, static_cast<uint_>(lastBlockStart));
                kernel.set_arg(outputArg, result.get_buffer());
                kernel.set_arg(outputOffsetArg, static_cast<uint_>(blockCount));

                queue.enqueue_task(kernel);
            }

            return totalBlockCount;
        }

        inline boost::compute::vector<unsigned char> BlockReduce(const boost::compute::vector<unsigned char> &values,
                                                                 std::size_t valueByteSize,
                                                                 std::size_t blockSize,
                                                                 const std::string &functionBody,
                                                                 boost::compute::command_queue &queue) {
            using namespace boost;

            const compute::context &ctx = queue.get_context();
            const std::size_t nValues = values.size() / valueByteSize;

            auto totalBlockCount = static_cast<std::size_t>(
                    std::ceil(static_cast<float>(nValues) / 2.f / static_cast<float>(blockSize)));
            compute::vector<unsigned char> resultVector(totalBlockCount * valueByteSize, ctx);

            Reduce(values, valueByteSize, resultVector, blockSize, functionBody, queue);

            return resultVector;
        }

        inline void GenericReduce(const boost::compute::vector<unsigned char> &values,
                                  std::size_t valueByteSize,
                                  boost::compute::vector<unsigned char> &result,
                                  const std::string &functionBody,
                                  boost::compute::command_queue &queue) {
            using namespace boost;

            const std::size_t blockSize = 256;

            compute::vector<unsigned char> results = BlockReduce(values, valueByteSize, blockSize, functionBody, queue);

            if (results.size() / valueByteSize > 1) {
                InplaceReduce(results,
                              valueByteSize,
                              functionBody,
                              queue);
            }

            compute::copy_n(results.begin(), valueByteSize, result.begin(), queue);
        }

    }// namespace detail

    inline boost::compute::vector<unsigned char> Reduce(const boost::compute::vector<unsigned char> &values,
                                                        std::size_t valueByteSize,
                                                        const std::string &reduceOp,
                                                        boost::compute::command_queue &queue) {
        using namespace boost;

        const compute::context &ctx = queue.get_context();
        boost::compute::vector<unsigned char> result(valueByteSize, ctx);
        compute::fill(result.begin(), result.end(), 0, queue);

        detail::GenericReduce(values, valueByteSize, result, reduceOp, queue);

        return result;
    }

}// namespace spla

#endif//SPLA_SPLAREDUCE_HPP


// Copyright 2020 Dow, Inc.
//
// 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 "Analysis.h"

#include <algorithm>
#include <chrono>
#include <iostream>
#include <iterator>
#include <filesystem>
#include <fstream>
#include <memory>
#include <type_traits>
#include <string>
#include <utility>
#include <vector>

// TODO: replace Boost.Accumulators with MKL
// https://software.intel.com/en-us/mkl-ssnotes-computing-quantiles-for-streaming-data

#include <boost/algorithm/string.hpp>
#include <boost/accumulators/accumulators.hpp>
#include <boost/accumulators/statistics/stats.hpp>
#include <boost/accumulators/statistics/mean.hpp>
#include <boost/accumulators/statistics/max.hpp>
#include <boost/accumulators/statistics/variance.hpp>
#include <boost/accumulators/statistics/p_square_quantile.hpp>
#include <boost/accumulators/statistics/rolling_count.hpp>
#include <boost/accumulators/statistics/rolling_mean.hpp>
#include <boost/accumulators/statistics/rolling_sum.hpp>
#include <boost/accumulators/statistics/p_square_cumul_dist.hpp>
#include <boost/accumulators/statistics/density.hpp>
#include <boost/lexical_cast.hpp>

#include <fmt/format.h>

namespace ncpost {

// Utility function to determine if all vectors have equal length.
template <typename... Args>
bool equal_length(const std::vector<Args>&... args) {
    if constexpr (sizeof...(Args) == 0) {
        return true;
    } else {
        return [](const auto& v0, const auto&... rest) {
            return ((v0.size() == rest.size()) && ...);
        }(args...);
    }
}

// EXPECTED DIMENSIONS:
//   rec, arc, net, grp, ave, time, strlen
//
// EXPECTED VARIABLES:
//   double x(rec)
//   double y(rec)
//   double zelev(rec)
//   double zhill(rec)
//   double zflag(rec)
//   int rec(rec)
//   int arc(rec)
//   char arcid(arc, strlen)
//   int net(rec)
//   char netid(net, strlen)
//   char grp(grp, strlen)
//   int ave(ave)
//   int time(time)
//   byte clmsg(time)

analysis::analysis(const std::string& filepath)
    : file_(filepath, ncpp::file::read), ds_(file_)
{
    static auto fn = [](std::size_t){};
    progressfn_ = fn;
}

std::string analysis::library_version()
{
    return ncpp::library_version();
}

std::string analysis::model_version() const
{
    std::string source_string;
    if (ds_.atts.contains("source"))
        ds_.atts["source"].read(source_string);
    return source_string;
}

std::string analysis::model_options() const
{
    std::string options_string;
    if (ds_.atts.contains("options"))
        ds_.atts["options"].read(options_string);
    return options_string;
}

std::string analysis::title() const
{
    std::string title_string;
    if (ds_.atts.contains("title"))
        ds_.atts["title"].read(title_string);
    return title_string;
}

std::vector<output_type_t> analysis::output_types() const
{
    std::vector<output_type_t> result;
    std::array<std::string, 4> vars{"conc", "depos", "ddep", "wdep"};
    for (const auto& var : vars) {
        if (ds_.vars.contains(var)) {
            output_type_t item;
            item.name = var;
            ds_.vars[var].atts["units"].read(item.units);
            boost::algorithm::to_lower(item.units);
            result.push_back(item);
        }
    }
    return result;
}

std::vector<int> analysis::averaging_periods() const
{
    if (ds_.vars.contains("ave"))
        return ds_.vars["ave"].values<int>();
    return {};
}

std::vector<std::string> analysis::source_groups() const
{
    if (ds_.vars.contains("grp"))
        return ds_.vars["grp"].values<std::string>();
    return {};
}

std::vector<std::string> analysis::receptor_arcids() const
{
    if (ds_.vars.contains("arcid"))
        return ds_.vars["arcid"].values<std::string>();
    return {};
}

std::vector<std::string> analysis::receptor_netids() const
{
    if (ds_.vars.contains("netid"))
        return ds_.vars["netid"].values<std::string>();
    return {};
}

std::size_t analysis::receptor_count() const
{
    if (ds_.vars.contains("rec"))
        return ds_.vars["rec"].size();
    return 0;
}

std::vector<receptor_t> analysis::receptors() const
{
    if (receptor_count() == 0)
        return {};

    const auto recs  = ds_.vars["rec"].values<int>();
    const auto x     = ds_.vars["x"].values<double>();
    const auto y     = ds_.vars["y"].values<double>();
    const auto zelev = ds_.vars["zelev"].values<double>();
    const auto zhill = ds_.vars["zhill"].values<double>();
    const auto zflag = ds_.vars["zflag"].values<double>();

    if (!equal_length(recs, x, y, zelev, zhill, zflag))
        throw std::out_of_range("Receptor coordinate length mismatch.");

    const auto arcids = receptor_arcids();
    const auto netids = receptor_netids();

    std::vector<int> arcs;
    if (!arcids.empty() && ds_.vars.contains("arc") && ds_.vars["arc"].size() == recs.size()) {
        arcs = ds_.vars["arc"].values<int>();
    }

    std::vector<int> nets;
    if (!netids.empty() && ds_.vars.contains("net") && ds_.vars["net"].size() == recs.size()) {
        nets = ds_.vars["net"].values<int>();
    }

    std::vector<receptor_t> result;

    result.reserve(recs.size());
    for (std::size_t i = 0; i < recs.size(); ++i)
    {
        std::string arcid;
        std::string netid;

        if (!arcs.empty() && arcs[i] > 0)
            arcid = arcids.at(static_cast<std::size_t>(arcs[i] - 1));
        if (!nets.empty() && nets[i] > 0)
            netid = netids.at(static_cast<std::size_t>(nets[i] - 1));

        result.emplace_back(receptor_t{ recs[i], x[i], y[i], zelev[i], zhill[i], zflag[i], arcid, netid });
    }

    return result;
}

std::size_t analysis::time_step_count() const
{
    if (ds_.vars.contains("time"))
        return ds_.vars["time"].size();
    return 0;
}

std::vector<time_step_t> analysis::time_steps() const
{
    if (time_step_count() == 0)
        return {};

    auto time = ds_.vars["time"].values<date::sys_seconds>();

    std::vector<unsigned char> clmsg;
    if (ds_.vars.contains("clmsg"))
        clmsg = ds_.vars["clmsg"].values<unsigned char>();

    std::vector<time_step_t> result;
    result.reserve(time.size());
    for (std::size_t i = 0; i < time.size(); ++i) {
        unsigned char iclmsg = clmsg.empty() ? 0xff : clmsg[i];
        result.emplace_back(time_step_t{ time[i], iclmsg });
    }

    return result;
}

namespace accumulators {

namespace ba = boost::accumulators;

struct epa_rolling_mean
{
    using rmacc_t = ba::accumulator_set<double, ba::stats<ba::tag::rolling_mean>>;
    using rsacc_t = ba::accumulator_set<double, ba::stats<ba::tag::rolling_sum, ba::tag::rolling_count>>;
    using maxacc_t = ba::accumulator_set<double, ba::features<ba::tag::max>>;

    epa_rolling_mean(int ave, int hours)
        : ave_(ave), value_(0),
          rsacc1_(rsacc_t(ba::tag::rolling_window::window_size = hours / ave)),
          rsacc2_(rsacc_t(ba::tag::rolling_window::window_size = hours / ave)),
          rmacc_(rmacc_t(ba::tag::rolling_window::window_size = hours / ave))
    {
        // This should evaluate true for any short-term averaging periods:
        // (1, 2, 3, 4, 6, 8, 12, 24)
        if (hours % ave != 0)
            throw std::runtime_error("Invalid averaging period.");
    }

    void operator()(double val, unsigned char cmflag)
    {
        if (ave_ == 1) {
            // Apply EPA averaging policy for 1 hour averaging period.
            double cm = (cmflag == 0) ? 0 : 1;
            rsacc1_(val);
            rsacc2_(cm);

            // Current number of elements in the rolling window.
            double n = static_cast<double>(ba::rolling_count(rsacc1_));

            if (n <= 24) {
                // Short-term average. See SUBROUTINE AVER in AERMOD calc2.f
                double numerator = ba::rolling_sum(rsacc1_);
                double denominator = std::max(n - ba::rolling_sum(rsacc2_), std::round(n * 0.75 + 0.4));
                value_ = numerator / denominator;
            }
            else {
                // Long-term average. See SUBROUTINE PERAVE in AERMOD output.f
                double numerator = ba::rolling_sum(rsacc1_);
                double denominator = n - ba::rolling_sum(rsacc2_);
                value_ = numerator / denominator;
            }
        }
        else {
            // Calculate standard rolling mean. Averaging already applied by AERMOD.
            rmacc_(val);
            value_ = ba::rolling_mean(rmacc_);
        }

        // Track the maximum value.
        maxacc_(value_);
    }

    double value() const
    {
        return value_;
    }

    double max() const
    {
        return ba::max(maxacc_);
    }

private:
    int ave_;
    double value_;
    rsacc_t rsacc1_; // rolling sum concentration (ave == 1)
    rsacc_t rsacc2_; // rolling sum valid hours (ave == 1)
    rmacc_t rmacc_;  // rolling mean concentration (ave != 1)
    maxacc_t maxacc_;
};

} // namespace accumulators

void analysis::set_progress_function(const std::function<void(std::size_t)>& fn)
{
    progressfn_ = fn;
}

void analysis::export_time_series(const options::general& opts, const options::tsexport& exopts) const
{
    using rmacc_t = accumulators::epa_rolling_mean;

    // Read the output matrix.
    auto recs = receptors();
    auto arcids = receptor_arcids();
    auto netids = receptor_netids();
    auto times = time_steps();
    auto matrix = output_matrix(opts.averaging_period, opts.source_group, opts.output_type, opts.scale_factor);
    auto allave = ds_.vars["ave"].values<int>();
    int minave = *std::min_element(allave.begin(), allave.end());

    std::ofstream ofs(exopts.output_file);

    // Write CSV header.
    fmt::memory_buffer header;

    if (opts.averaging_period == 1)
        fmt::format_to(header, "arcid,netid,receptor,time,calm_missing,x,y,zelev,zhill,zflag,{}", opts.output_type);
    else
        fmt::format_to(header, "arcid,netid,receptor,time,x,y,zelev,zhill,zflag,{}", opts.output_type);

    if (exopts.rm_windows.size() > 0) {
        for (const auto& w : exopts.rm_windows) {
            fmt::format_to(header, ",{}[RM{}]", opts.output_type, w);
        }
    }

    fmt::format_to(header, "\n");
    ofs << fmt::to_string(header);

    // Write CSV records.
    for (std::size_t j = 0; j < matrix.size2(); ++j) // receptors
    {
        progressfn_(j);

        const auto& rec = recs.at(j);

        // Create the rolling mean accumulators.
        std::vector<rmacc_t> rmaccs;
        rmaccs.reserve(exopts.rm_windows.size());
        for (double w : exopts.rm_windows) {
            rmaccs.emplace_back(rmacc_t(opts.averaging_period, static_cast<int>(w) * 24));
        }

        fmt::memory_buffer buffer1;
        fmt::format_to(buffer1, "{},{},{:d},{:.2f},{:.2f},{:.2f},{:.2f},{:.2f}",
            rec.arcid, rec.netid, rec.id, rec.x, rec.y,
            rec.zelev, rec.zhill, rec.zflag);

        for (std::size_t i = 0; i < matrix.size1(); ++i) // time steps
        {
            std::vector<double> rmvals(exopts.rm_windows.size(), 0);

            const auto ts = times.at(i * opts.averaging_period / minave);
            const double val = matrix.at_element(i, j); // ntime * nrecs, column-major

            for (auto&& rmacc : rmaccs)
                rmacc(val, ts.calm_missing);

            fmt::memory_buffer buffer2;

            if (opts.averaging_period == 1)
                fmt::format_to(buffer2, ",{},{},{:.6g}", date::format("%F %R", ts.time), ts.calm_missing, val);
            else
                fmt::format_to(buffer2, ",{},{:.6g}", date::format("%F %R", ts.time), val);

            for (const auto& rmacc : rmaccs)
                fmt::format_to(buffer2, ",{}", rmacc.value());

            ofs << fmt::to_string(buffer1) << fmt::to_string(buffer2) << "\n";
        }
    }

    ofs.close();
}

void analysis::calc_receptor_stats(const options::general& opts, const options::statistics& statopts, statistics_type& out) const
{
    using namespace boost::accumulators;

    if (!statopts.calc_avg && !statopts.calc_max && !statopts.calc_std &&
        !statopts.percentiles.size() && !statopts.maxrm_windows.size())
        return;

    // Define the statistical accumulator types.
    using avgacc_t = accumulator_set<double, features<tag::mean>>;
    using maxacc_t = accumulator_set<double, features<tag::max>>;
    using varacc_t = accumulator_set<double, features<tag::variance>>;
    using p2acc_t = accumulator_set<double, stats<tag::p_square_quantile>>;
    using rmacc_t = accumulators::epa_rolling_mean;

    // Read the output matrix.
    std::size_t nrecs = receptor_count();
    auto times = time_steps();
    auto matrix = output_matrix(opts.averaging_period, opts.source_group, opts.output_type, opts.scale_factor);
    auto allave = ds_.vars["ave"].values<int>();
    int minave = *std::min_element(allave.begin(), allave.end());

    // Initialize the output vectors for basic statistics.
    if (statopts.calc_avg)
        out.avg.resize(nrecs);
    if (statopts.calc_max)
        out.max.resize(nrecs);
    if (statopts.calc_std)
        out.std.resize(nrecs);

    // Initialize the output vectors for P2 calculations.
    for (double p : statopts.percentiles) {
        std::vector<double> v(nrecs, 0);
        out.p2.push_back(v);
    }

    // Initialize the output vectors for rolling mean calculations.
    for (double w : statopts.maxrm_windows) {
        std::vector<double> v(nrecs, 0);
        out.rm.push_back(v);
    }

    // Analyze the data array for each receptor.
    for (std::size_t j = 0; j < matrix.size2(); ++j) // receptors
    {
        progressfn_(j);

        // Create the accumulators.
        avgacc_t avgacc;
        maxacc_t maxacc;
        varacc_t varacc;

        std::vector<p2acc_t> p2accs;
        p2accs.reserve(statopts.percentiles.size());
        for (double p : statopts.percentiles)
            p2accs.emplace_back(p2acc_t(quantile_probability = p));

        std::vector<rmacc_t> rmaccs;
        rmaccs.reserve(statopts.maxrm_windows.size());
        for (double w : statopts.maxrm_windows)
            rmaccs.emplace_back(rmacc_t(opts.averaging_period, static_cast<int>(w) * 24));

        // Main update loop.
        for (std::size_t i = 0; i < matrix.size1(); ++i) // time steps
        {
            const auto ts = times.at(i * opts.averaging_period / minave);
            const double val = matrix.at_element(i, j); // ntime * nrecs, column-major

            if (statopts.calc_avg)
                avgacc(val);
            if (statopts.calc_max)
                maxacc(val);
            if (statopts.calc_std)
                varacc(val);
            for (auto&& p2acc : p2accs)
                p2acc(val);
            for (auto&& rmacc : rmaccs)
                rmacc(val, ts.calm_missing);
        }

        // Extract results.
        if (statopts.calc_avg)
            out.avg[j] = mean(avgacc);
        if (statopts.calc_max)
            out.max[j] = max(maxacc);
        if (statopts.calc_std)
            out.std[j] = sqrt(variance(varacc));
        for (std::size_t iacc = 0; iacc < p2accs.size(); ++iacc)
            out.p2.at(iacc).at(j) = p_square_quantile(p2accs[iacc]);
        for (std::size_t iacc = 0; iacc < p2accs.size(); ++iacc)
            out.rm.at(iacc).at(j) = rmaccs[iacc].max();
    }
}

void analysis::calc_histogram(const options::general& opts, const options::histogram& histopts, histogram_type& out) const
{
    using namespace boost::accumulators;

    using cdfacc_t = accumulator_set<double, stats<tag::p_square_cumulative_distribution>>;
    using pdfacc_t = accumulator_set<double, stats<tag::density>>;

    if (!histopts.calc_cdf && !histopts.calc_pdf)
        return;

    // Read the data array.
    auto matrix = output_matrix(opts.averaging_period, opts.source_group, opts.output_type, opts.scale_factor);

    // Create the accumulators.
    cdfacc_t cdfacc(tag::p_square_cumulative_distribution::num_cells = static_cast<std::size_t>(histopts.cdf_bins));

    // If selected cache size exceeds sample count, use all samples.
    std::size_t cache_size = std::max(static_cast<std::size_t>(histopts.pdf_cache_size), matrix.data().size());

    pdfacc_t pdfacc(tag::density::num_bins = static_cast<std::size_t>(histopts.pdf_bins),
                    tag::density::cache_size = cache_size);

    // Analyze the data array.
    for (std::size_t j = 0; j < matrix.size2(); ++j) // receptors
    {
        progressfn_(j);

        for (std::size_t i = 0; i < matrix.size1(); ++i) { // time steps
            const double val = matrix.at_element(i, j); // ntime * nrecs, column-major
            if (histopts.calc_cdf)
                cdfacc(val);
            if (histopts.calc_pdf)
                pdfacc(val);
        }
    }

    // Extract results.
    typedef boost::iterator_range<std::vector<std::pair<double, double>>::iterator> histogram_t;
    if (histopts.calc_cdf) {
        histogram_t cdf = p_square_cumulative_distribution(cdfacc);
        out.cdf.insert(out.cdf.end(), std::make_move_iterator(cdf.begin()), std::make_move_iterator(cdf.end()));
    }
    if (histopts.calc_pdf) {
        histogram_t pdf = density(pdfacc);
        out.pdf.insert(out.pdf.end(), std::make_move_iterator(pdf.begin()), std::make_move_iterator(pdf.end()));
    }
}

matrix_t analysis::output_matrix(int ave, const std::string& grp, const std::string& var, double sf) const
{
    auto slice = ds_.vars[var].select(
        ncpp::selection<int>{"ave", ave, ave},
        ncpp::selection<std::string>{"grp", grp, grp});

    // Get full matrix dimensions.
    std::size_t ntime = ds_.vars["time"].size();
    std::size_t nrecs = ds_.vars["rec"].size();
    auto allave = ds_.vars["ave"].values<int>();
    int minave = *std::min_element(allave.begin(), allave.end());

    // Recalculate ntime to exclude fill values for multiple averaging periods.
    ntime = ntime / static_cast<std::size_t>(ave / minave);

    // Read the data array and remove fill values.
    auto values = slice.values<double, aligned_allocator_t>();
    values.erase(std::remove(values.begin(), values.end(), NC_FILL_DOUBLE), values.end());
    if (values.size() != ntime * nrecs)
        throw std::runtime_error("Data array has unexpected missing values.");

    matrix_t result(ntime, nrecs); // column-major
    std::copy(values.begin(), values.end(), result.data().begin());

    //matrix_t result = slice.matrix<double, aligned_allocator_t>();
    //auto end = std::remove(result.data().begin(), result.data().end(), NC_FILL_DOUBLE);
    //if (std::distance(result.data().begin(), end) != ntime * nrecs)
    //    throw std::runtime_error("Data array has unexpected missing values.");
    //result.resize(ntime, nrecs); // column-major

    // Apply scale factor.
    if (sf != 1.0)
        result *= sf;

    return result;
}

} // namespace ncpost