/* -*-C++-*-
   (c) Copyright 2005-2008, Hewlett-Packard Development Company, LP

   See the file named COPYING for license details
*/

/** @file
    StatsQuantile implementation
*/

#include <algorithm>
#include <sstream>
#include <cstring>

#include <boost/format.hpp>

#include <Lintel/AssertBoost.hpp>
#include <Lintel/Double.hpp>
#include <Lintel/HashMap.hpp>
#include <Lintel/LintelLog.hpp>
#include <Lintel/StatsQuantile.hpp>

using namespace std;
using boost::format;

namespace {
    struct ErrorNbound {
	double quantile_error;
	int64_t nbound;
	ErrorNbound() : quantile_error(0), nbound(0) { }
	ErrorNbound(double a, int64_t b) : quantile_error(a), nbound(b) { }

	uint32_t hash() const {
	    uint32_t a = static_cast<uint32_t>(quantile_error * 1e9);
	    uint32_t b = nbound >> 32;
	    uint32_t c = nbound & 0xFFFFFFFFU;
	    return lintel::BobJenkinsHashMix3(a, b, c);
	}

	bool operator == (const ErrorNbound &rhs) const {
	    return quantile_error == rhs.quantile_error && nbound == rhs.nbound;
	}
    };

    struct BK {
	uint32_t b, k;
	BK() : b(0), k(0) { }
	BK(uint32_t in_b, uint32_t in_k) : b(in_b), k(in_k) { }
    };

    HashMap<ErrorNbound, BK> bk_cache;

    double fact(double a) {
	double ret = 1.0;
	for(double i = 2;i<=a;i++) {
	    ret *= i;
	}
	return ret;
    }

    double choose(int b, int a) { // b choose a
	double ret = fact(b) / (fact(b-a) * fact(a));
	return ret;
    }

    class doubleOrder {
    public:
	inline bool operator() (double a, double b) const {
	    return a < b;
	}
    };

}

StatsQuantile::StatsQuantile(double _quantile_error, int64_t in_nbound, int _print_nrange, bool lazy)
    : quantile_error(_quantile_error), Nbound(in_nbound < 100 ? 100 : in_nbound),
      print_nrange(_print_nrange), lazy(lazy)
{
    // sanity check that we haven't been called in a weird/wrong way.
    INVARIANT(quantile_error < 0.2, format("whoa, quantile_error %.3g >= 0.2??") % quantile_error);
    double tmp_nbound = Nbound; // easy for calculations
    INVARIANT(in_nbound > 1, format("Usage error? StatsQuantile nbound set to %d?") % in_nbound);

    int32_t best_b = -1, best_k = -1;

    // If we are making lots of smallish StatsQuantiles, we can spend
    // too much of our time doing this, so cache the sizes.
    BK *cache = bk_cache.lookup(ErrorNbound(quantile_error, Nbound));

    if (cache != NULL) {
	best_b = cache->b;
	best_k = cache->k;
    } else {
	// The hunt for b = nbuffers, and k = buffer_size is on, following
	// section 4.5 in the paper

	const int32_t max_b = 30;
	const int32_t max_h = 40;
	for(int32_t b=2;b<max_b;b++) {
	    int32_t h = 3;
	    for(;h < max_h;h++) {
		double v = (h-2) * choose(b+h-2,h-1) - choose(b+h-3,h-3) + choose(b+h-3,h-2);
		if (v > 2 * quantile_error * tmp_nbound) {
		    --h; // this one is too high
		    break;
		}
	    }
	    double k = ceil(tmp_nbound / choose(b+h-2,h-1));
	    if (b * k < 1.0e9) {
		if (best_b == -1 || b * k < best_b * best_k) {
		    best_b = b;
		    best_k = (int)k;
		}
	    }
	}
	if (bk_cache.size() > 1000) { // keep the cache small
	    bk_cache.clear();
	}
	bk_cache[ErrorNbound(quantile_error, Nbound)] = BK(best_b, best_k);
    }
    INVARIANT(best_b > 0 && best_k > 0,
	      format("unable to find b/k for %.8g %.8g") % quantile_error % tmp_nbound);

    // lower bound sanity value
    if (best_k < 10) best_k = 10;
    // First requirement makes the logic setting the buffer_level faster
    // second one just seems sane.
    init(best_b, best_k);
}

// fake Nbound, big enough for tests
StatsQuantile::StatsQuantile(const string &,
			     int _nbuffers, int _buffer_size, int _print_nrange)
    : quantile_error(Double::NaN), Nbound(5000000), print_nrange(_print_nrange), lazy(true)
{
    init(_nbuffers, _buffer_size);
}

void StatsQuantile::init(int _nbuffers, int _buffer_size) {
    nbuffers = _nbuffers;
    buffer_size = _buffer_size;
    INVARIANT(nbuffers > 1,
	      "need at least two buffers for code to work correctly");
    INVARIANT(buffer_size >= 10,
	      "buffers smaller than 10 don't make sense.");
    all_buffers = new one_buffer[nbuffers];
    for(int i = 0; i < nbuffers; ++i) {
	all_buffers[i] = NULL;
    }
    buffer_weight = new int64_t[nbuffers];
    buffer_level = new int[nbuffers];
    buffer_sorted = new bool[nbuffers];
    if (!lazy) {
	tmp_buffer = new double[buffer_size];
    } else {
	tmp_buffer = NULL;
    }
    collapse_pos = new int[buffer_size];
    init_buffers();
}

void StatsQuantile::init_buffers() {
    for(int i=0;i<nbuffers;i++) {
	if (all_buffers[i] == NULL && ! lazy) {
	    all_buffers[i] = new double[buffer_size];
	}
	// touch all the space to make sure it is allocated; necessary
	// to avoid demand allocation (which can create long delays
	// when used in Buttress2)
	if (all_buffers[i] != NULL) {
	    for(int j=0;j<buffer_size;j+=32) {
		all_buffers[i][j] = 0.0;
	    }
	}
	buffer_weight[i] = -1;
	buffer_level[i] = -1;
	buffer_sorted[i] = false;
    }

    cur_buffer = 0;
    cur_buffer_pos = 0;
    buffer_weight[cur_buffer] = 1;
    buffer_level[cur_buffer] = 0;
    collapse_even_low = true;
    if (lazy) {
	cur_buffer = -1;
	cur_buffer_pos = buffer_size+1;
    }
}

StatsQuantile::~StatsQuantile() {
    for(int i=0;i<nbuffers;i++) {
	delete[] all_buffers[i];
    }
    delete[] all_buffers;
    delete[] buffer_weight;
    delete[] buffer_level;
    delete[] buffer_sorted;
    delete[] tmp_buffer;
    delete[] collapse_pos;
}

void StatsQuantile::reset() {
    Stats::reset();
    init_buffers();
}

void StatsQuantile::addQuantile(const double value) {
    if (cur_buffer_pos >= buffer_size) {
	cur_buffer += 1;
	cur_buffer_pos = 0;
	if (cur_buffer == nbuffers) {
	    collapse();
	} else if (cur_buffer>0) {
#if 1
	    // do the sort here to amortize the work
	    sort(all_buffers[cur_buffer - 1],
		      all_buffers[cur_buffer - 1]+buffer_size,
		      doubleOrder());
	    buffer_sorted[cur_buffer-1] = true;
#endif
	}
        SINVARIANT(cur_buffer < nbuffers);
	if (lazy && all_buffers[cur_buffer] == NULL) {
	    all_buffers[cur_buffer] = new double[buffer_size];
	}
	buffer_weight[cur_buffer] = 1;
	buffer_level[cur_buffer] = 0;
	if (cur_buffer == nbuffers - 1) {
	    buffer_level[cur_buffer] = buffer_level[cur_buffer - 1];
	}
    }
    buffer_sorted[cur_buffer] = false;
    all_buffers[cur_buffer][cur_buffer_pos] = value;
    cur_buffer_pos += 1;
}

void StatsQuantile::add(const double value) {
    Stats::add(value);
    addQuantile(value);
}

void StatsQuantile::add(const Stats &_stat) {
    const StatsQuantile *stat = dynamic_cast<const StatsQuantile *>(&_stat);
    INVARIANT(stat != NULL, "need another StatsQuantile for add(Stats) to work");

    Stats::add(_stat);

    if ((cur_buffer == -1 ||
	 (cur_buffer == 0 && cur_buffer_pos == 0)) && // we are empty
	buffer_size == stat->buffer_size && // we are compatible
	nbuffers == stat->nbuffers) {
	INVARIANT(countll() == _stat.countll(), "??");
	INVARIANT(stat->cur_buffer < nbuffers, "??");
	for(int i=0; i <= stat->cur_buffer; ++i) {
	    // Will copy some useless stuff in last buffer.
	    if (all_buffers[i] == NULL) {
		all_buffers[i] = new double[buffer_size];
	    }
	    memcpy(all_buffers[i], stat->all_buffers[i],
		   sizeof(double) * buffer_size);
	    buffer_weight[i] = stat->buffer_weight[i];
	    buffer_level[i] = stat->buffer_level[i];
	    buffer_sorted[i] = stat->buffer_sorted[i];
	}
	cur_buffer = stat->cur_buffer;
	cur_buffer_pos = stat->cur_buffer_pos;
	// TODO: determine if we really have to copy collapse_pos
	memcpy(collapse_pos, stat->collapse_pos, sizeof(int) * buffer_size);
	collapse_even_low = stat->collapse_even_low;
	return;
    }

    // There should be a faster way to do this, but this should work
    // correctly.  The faster way should be possible if for no other
    // reason we are inserting the values in order.  In essence we
    // pretent that stat was made up of repeated insertions of the
    // values kept.  Since it's contents are indistinguishable from
    // that actually happening, the merge should be pretty close to
    // correct; this is of course not a formal proof, and I suspect
    // that it is wrong in some special case.

    // TODO: special case bulk loading of lots of the same values;
    // should be able to greatly simplify the addQuantile code to just
    // copy a swath of values in.  Probably can even sort the buffer
    // first, and if we are in order, then just drop everything in and
    // leave the buffer marked sorted.

    for(int i=0; i < stat->cur_buffer; ++i) {
	// cur_buffer is the one we are filling, so all previous ones are full
	double *abuf = stat->all_buffers[i];
	for(int j=0; j < stat->buffer_size; ++j) {
	    double v = abuf[j];
	    for(int64_t k=0; k < stat->buffer_weight[i]; ++k) {
		addQuantile(v);
	    }
	}
    }
    double *abuf = stat->all_buffers[stat->cur_buffer];
    for (int j=0; j < stat->cur_buffer_pos; ++j) {
	double v = abuf[j];
	for(int64_t k=0; k < stat->buffer_weight[stat->cur_buffer]; ++k) {
	    addQuantile(v);
	}
    }
}

// getQuantile is quite expensive for the upper quantiles; in theory
// we ought to be able to do some sort of binary-like search in that case,
// skipping by a lot of the lower quantiles in a batch rather than
// ratcheting through them one by one as in the paper, and the current
// implementation

// const because the only operation on class variables is a sort of buffers
double StatsQuantile::getQuantile(double quantile, bool allow_invalid_nbound) const {
    int64_t count = countll();
    if (count == 0) {
	return Double::NaN;
    }
    INVARIANT(count <= Nbound || allow_invalid_nbound,
	      format("Error: %d (# entries in quantile) > %d (Nbound on quantile)")
	      % count % Nbound);
    INVARIANT(quantile >= 0.0 && quantile <= 1.0, "quantile out of bounds");
    // this is slightly different than the algorithm in the paper; the
    // goal is to allow people to calculate quantiles while the algorithm
    // is running without having to add in and remove the +-\Inf entries
    // that would be needed to make the last buffer have integral size
    // I think this change is correct

    for(int i = 0; i <= cur_buffer; i++) {
	collapse_pos[i] = 0;
	// TODO: add a check that the total weight covered by all the
	// current buffers is exactly the number of elements that we
	// have.  Have to be a little careful on the last buffer to
	// only count the values that we have actually inserted.
	if (!buffer_sorted[i]) {
	    if (i < cur_buffer) {
		sort(all_buffers[i],all_buffers[i] + buffer_size,
			  doubleOrder());
	    } else {
		// only sort part that has values in it.
		sort(all_buffers[i],all_buffers[i] + cur_buffer_pos,
			  doubleOrder());
		// Fill the rest with Inf's so that we don't have to
		// do any special logic when choosing the "smallest"
		// next value, all of the Inf's will come at the end, and
		// if we get to an Inf, then we are done anyway and it
		// doesn't matter if we take one that doesn't actually exist.
		for(int j=cur_buffer_pos;j<buffer_size;j++) {
		    all_buffers[i][j] = Double::Inf;
		}
	    }
	    buffer_sorted[i] = true;
	}
    }
    // the 1e-15 accounts for a minor rounding error that seems to occur
    // when doing division, e.g. ceil(150000.0 * (131058.0 / 150000.0)) on
    // linux = 131059.0; I'd imagine we'll never see 1e+15 values
    // wherein this adjustment would make a difference
    double nentries = countll();
    INVARIANT(nentries < 1e+14,
	      "Rounding error adjustment may start to do something wrong around 1e+14 entries.\n"
	      "May be safe to decrease rounding adjustment to 0.5e-15, but will be getting very\n"
	      "close to the limit of precision in floating point.  How did you get this many\n"
	      "entries into the table??");
    // Subtract 1 because the quantile positions count from 1..n, but
    // C++ arrays index from 0..n-1
    int64_t target_index
	= static_cast<int64_t>(ceil(nentries * (quantile - 1e-15))) - 1;
    if (target_index < 0) target_index = 0;
    if (target_index == nentries) {
	target_index = (long long)(nentries - 1);
    }
    // can be switched back to getQuantileByIndex() if there are difficulties.
    return getQuantileByBinSearchIndex(target_index);
}

int StatsQuantile::collapseFindFirstBuffer() {
    double nentries = countll();
    INVARIANT(nentries < 1e+14, "getQuantile will fail now.");

    int first_buffer;
    for(first_buffer = nbuffers - 1;first_buffer > 0; --first_buffer) {
	if (buffer_level[first_buffer - 1] != buffer_level[first_buffer]) {
	    break;
	}
    }
    INVARIANT(first_buffer < nbuffers - 1,
	      "Whoa, collapse should always operate on at least two buffers");
    INVARIANT(buffer_level[first_buffer] >= 0,
	      "Whoa, buffer level should be positive");
    return first_buffer;
}

// TODO-2010-04-26: remove this and the associated test case; if it's
// lasted this long, then we're not having problems with the new
// binary search version.
double StatsQuantile::getQuantileByIndex(uint64_t target_index) const {
    uint64_t cur_index = 0;

    PriorityQueue<std::pair<double, int>, pairCmp> tmp_pq(nbuffers);
    for(int i=0; i < cur_buffer; ++i) {
	collapse_pos[i] = 0;
	tmp_pq.push(make_pair(collapseVal(i), i));
    }
    if (cur_buffer_pos > 0) {
	collapse_pos[cur_buffer] = 0;
	tmp_pq.push(make_pair(collapseVal(cur_buffer), cur_buffer));
    }

#if LINTEL_DEBUG
    double prev_val = -Double::Inf;
#endif
    while(!tmp_pq.empty()) {
	double min_val = tmp_pq.top().first;
	int min_buffer = tmp_pq.top().second;

#if LINTEL_DEBUG
	SINVARIANT(collapse_pos[min_buffer] < buffer_size &&
		   (min_buffer < cur_buffer || collapse_pos[min_buffer] < cur_buffer_pos));
	INVARIANT(min_val >= prev_val, "Whoa, sort order error?!");
	prev_val = min_val;
#endif

	collapse_pos[min_buffer] += 1;

	// this entry occupies sorted order positions
	// [cur_index .. cur_index + buffer_weight[min_buffer] - 1]
	// so the check is for cur_index > target_index
	// need to think about the logic for this in collapse()
	cur_index += buffer_weight[min_buffer];
	if (cur_index > target_index) {
	    return min_val;
	}
	int max_pos = min_buffer < cur_buffer ? buffer_size : cur_buffer_pos;
	if (collapse_pos[min_buffer] < max_pos) {
	    tmp_pq.replaceTop(make_pair(collapseVal(min_buffer), min_buffer));
	} else {
	    tmp_pq.pop();
	}
    }
    return Double::NaN;
}

namespace {
    struct BinSearchState {
	uint32_t min, max, orig_max, split_lower, split_upper;
	explicit BinSearchState(uint32_t b)
	    : min(0), max(b), orig_max(b)
	{ }
	uint32_t mid() const { return (max + min) / 2; }
	uint32_t remain() const { return max - min; }

	void calculateBounds(double target, double *buffer) {
	    double *lower = lower_bound(buffer + min, buffer + max, target);
	    split_lower = lower - buffer;
	    double *upper = upper_bound(buffer + min, buffer + max, target);
	    split_upper = upper - buffer;
	}
    };

    uint32_t biggestRemaining(const vector<BinSearchState> &state) {
	uint32_t biggest = 0;
	uint32_t remain = state[0].remain();
	for(uint32_t i = 1; i < state.size(); ++i) {
	    if (state[i].remain() > remain) {
		biggest = i;
		remain = state[i].remain();
	    }
	}
	return biggest;
    }
}

double StatsQuantile::getQuantileByBinSearchIndex(uint64_t target_index) const {
    vector<BinSearchState> state;

    // one last buffer if the last one is present
    state.resize(cur_buffer + (cur_buffer_pos > 0 ? 1 : 0), BinSearchState(buffer_size));
    if (cur_buffer_pos > 0) {
	state[cur_buffer].max = state[cur_buffer].orig_max = cur_buffer_pos;
    }

    uint64_t entry_count = 0;
    for(uint32_t i = 0; i < static_cast<uint32_t>(cur_buffer); ++i) {
	entry_count += static_cast<uint64_t>(buffer_weight[i]) * buffer_size;
    }
    entry_count += cur_buffer_pos;
    SINVARIANT(entry_count == countll());

    LintelLogDebug("StatsQuantile::getQuantile", format("entries %d") % entry_count);
    while(true) {
	uint32_t biggest = biggestRemaining(state);
	double split_val = all_buffers[biggest][state[biggest].mid()];
	LintelLogDebug("StatsQuantile::getQuantile", format("biggest %d, split @%.12g")
		       % biggest % split_val);

	uint64_t count_lower = 0;
	uint64_t count_match = 0;
	uint64_t count_upper = 0;
	for(uint32_t i = 0; i < state.size(); ++i) {
	    state[i].calculateBounds(split_val, all_buffers[i]);
	    LintelLogDebug("StatsQuantile::getQuantile",
			   format("  buffer %d wt=%d min=%d, sl=%d, su=%d, max=%d")
			   % i % buffer_weight[i] % state[i].min
			   % state[i].split_lower % state[i].split_upper % state[i].max);
	    count_lower += buffer_weight[i] * state[i].split_lower; // 0 .. split_lower - 1
	    count_match += buffer_weight[i]
		* (state[i].split_upper - state[i].split_lower); // split_lower .. split_upper -1
	    count_upper += buffer_weight[i]
		* (state[i].orig_max - state[i].split_upper); // split_upper .. buffer_size - 1
	}
	LintelLogDebug("StatsQuantile::getQuantile", format("cl=%d, cm=%d, cu=%d")
		       % count_lower % count_match % count_upper);
	SINVARIANT(count_lower + count_match + count_upper == entry_count);

	if (target_index < count_lower) {
	    LintelLogDebug("StatsQuantile::getQuantile", "in-lower");
	    bool any_changed = false;
	    for(uint32_t i = 0; i < state.size(); ++i) {
		if (state[i].max != state[i].split_lower) any_changed = true;
		state[i].max = state[i].split_lower;
		SINVARIANT(state[i].min <= state[i].max);
	    }
	    SINVARIANT(any_changed);
	} else if (target_index < count_lower + count_match) {
	    LintelLogDebug("StatsQuantile::getQuantile", "***match");
	    return split_val; // even if we didn't find the exact match, we're good enough
	} else if (target_index < count_lower + count_match + count_upper) {
	    LintelLogDebug("StatsQuantile::getQuantile", "in-upper");
	    bool any_changed = false;
	    for(uint32_t i = 0; i < state.size(); ++i) {
		if (state[i].min != state[i].split_upper) any_changed = true;
		state[i].min = state[i].split_upper;
		SINVARIANT(state[i].min <= state[i].max);
	    }
	    SINVARIANT(any_changed);
	}
    }
    FATAL_ERROR("unimplemented");
}

int64_t StatsQuantile::collapseSortBuffers(int first_buffer) {
    int64_t total_weight = 0;
    for(int i=first_buffer;i<nbuffers;i++) {
	total_weight += buffer_weight[i];
	INVARIANT(buffer_weight[i] > 0, "Whoa, buffer_weight should be at least 1");
	collapse_pos[i] = 0;
	if (buffer_sorted[i]) {
	    // could re-verify sortedness here, but we'll probably catch it later
	} else {
	    INVARIANT(i == (nbuffers - 1) || buffer_weight[i] == 1,
		      format("Huh %d != %d // %d") % i % (nbuffers - 1) % buffer_weight[i]);
	    // sort the buffer
	    sort(all_buffers[i], all_buffers[i] + buffer_size, doubleOrder());
	    buffer_sorted[i] = true;
	}
    }
    return total_weight;
}

int64_t StatsQuantile::collapseNextQuantileOffset(int64_t total_weight) {
    int64_t next_quantile_offset = 0;
    if ((total_weight % 2) == 0) {
	if (collapse_even_low) {
	    next_quantile_offset = total_weight / 2;
	    collapse_even_low = false;
	} else {
	    next_quantile_offset = (total_weight + 2) / 2;
	    collapse_even_low = true;
	}
    } else if ((total_weight % 2) == 1) {
	next_quantile_offset = (total_weight + 1) / 2;
    } else {
	FATAL_ERROR("Huh?");
    }
    return next_quantile_offset;
}

void StatsQuantile::collapse() {
    // could do a more "in-place" collapse by first doing the walk, and
    // for every entry which is not going to go into the final collation,
    // replace it with a NaN; then move all the entries in the output
    // bucket which are not NaN's to the end, then collate again in order
    // filling into the output location; this does in-place, but requires
    // two walks over the larger data, and so may very well be slower
    // it's also a lot more complex

    int first_buffer = collapseFindFirstBuffer();

    // first, sort each of the unsorted input buffers
    int64_t total_weight = collapseSortBuffers(first_buffer);

    int64_t next_quantile_offset = collapseNextQuantileOffset(total_weight);

    // Since we start counting at offset 0 (the paper starts at 1), we
    // subtract one from the starting offset
    next_quantile_offset -= 1;
    DEBUG_SINVARIANT(next_quantile_offset >= 0);

    int64_t cur_quantile_offset = 0;
    int64_t next_output_pos = 0;
    IF_LINTEL_DEBUG(double prev_val = -Double::Inf;)

    // Because a double and int are small, it's faster to store the
    // pair rather than do the double dereference each time to
    // translate a buffer number into its value.
    PriorityQueue<pair<double, int>, pairCmp> pq(nbuffers);
    for(int i=first_buffer; i < nbuffers; ++i) {
	pq.push(make_pair(collapseVal(i), i));
    }
    if (tmp_buffer == NULL) {
	tmp_buffer = new double[buffer_size];
    }
    while(next_output_pos < buffer_size) {
	int min_buffer = pq.top().second;
	double min_val = pq.top().first;

	DEBUG_INVARIANT(min_val >= prev_val, "Whoa, sort order error?!");
	IF_LINTEL_DEBUG(prev_val = min_val;)
	collapse_pos[min_buffer] += 1;
	// same logic as for output(), this entry is in positions
	// [cur_quantile_offset .. cur_quantile_offset + bw[mb] - 1]
	// so if after adding in bw[mb], cqo > nqo then we just passed
	// nqo, so it should be used as the value to fit into the
	// new list
	cur_quantile_offset += buffer_weight[min_buffer];
	if (cur_quantile_offset > next_quantile_offset) {
	    tmp_buffer[next_output_pos] = min_val;
	    next_output_pos += 1;
	    next_quantile_offset += total_weight;
	}

	if (collapse_pos[min_buffer] < buffer_size) {
	    pq.replaceTop(make_pair(collapseVal(min_buffer), min_buffer));
	} else {
	    pq.pop();
	    DEBUG_SINVARIANT(next_output_pos == buffer_size || !pq.empty());
	}
    }
    INVARIANT(next_quantile_offset / total_weight == buffer_size,
	      format("Huh, didn't get to correct quantile offset %lld/%d != %d?!")
	      % next_quantile_offset % total_weight % buffer_size);

    // update all_buffers[first_buffer]
    memcpy(all_buffers[first_buffer],tmp_buffer, buffer_size * sizeof(double));
    buffer_level[first_buffer] = buffer_level[first_buffer] + 1;
    buffer_weight[first_buffer] = total_weight;
    buffer_sorted[first_buffer] = true;

    // clean up buffer weights
    for(int i=first_buffer + 1;i<nbuffers;i++) {
	buffer_weight[i] = -1;
	buffer_level[i] = -1;
	buffer_sorted[i] = false;
    }
    cur_buffer = first_buffer + 1;
}

void StatsQuantile::dumpState() {
    printf("%d buffers, each containing %d entries\n",
	   nbuffers,buffer_size);
    printf("cur buffer is %d, cur buffer pos is %d\n",
	   cur_buffer,cur_buffer_pos);

    for(int i=0;i<=cur_buffer;i++) {
	cout << format("  buffer %d, level %d, weight %d:\n    ")
	    % i % buffer_level[i] % buffer_weight[i];
	int max = i == cur_buffer ? cur_buffer_pos : buffer_size;
	for(int j=0;j<max;j++) {
	    printf("%.4g, ",all_buffers[i][j]);
	    if ((j%10) == 9) {
		printf("\n    ");
	    }
	}
	printf("\n");
    }
}

void StatsQuantile::printRome(int depth, ostream &out) const {
    Stats::printRome(depth,out);
    string spaces;
    for(int i = 0; i < depth; i++) {
	spaces += " ";
    }
    double nentries = countll();
    if (nentries > 0) {
	out << spaces << "{ quantiles (\n";
	double step = 1.0 / (double)print_nrange;
	int nquantiles = 0;
	for(double quantile = step;Double::lt(quantile,1.0);quantile += step) {
	    double quantile_value = getQuantile(quantile);
	    out << spaces << "  { " << quantile*100 << " " << quantile_value << " }\n";
	    ++nquantiles;
	}
	for(double tail_frac = 0.1; (tail_frac * nentries) >= 1.0;) {
	    double quantile_value = getQuantile(1-tail_frac);
	    out << spaces << "  { " << 100*(1-tail_frac) << " " << quantile_value << " }\n";
	    tail_frac /= 2.0;
	    quantile_value = getQuantile(1-tail_frac);
	    out << spaces << "  { " << 100*(1-tail_frac) << " " << quantile_value << " }\n";
	    tail_frac /= 5.0;
	}
	out << spaces << ") }\n";
    }
}

void StatsQuantile::printFile(FILE *out, int nranges) {
    ostringstream tmp;
    printTextRanges(tmp, nranges);
    fwrite(tmp.str().data(), tmp.str().size(), 1, out);
}

void StatsQuantile::printTextRanges(ostream &out, int nranges, double multiplier) const {
    nranges = (nranges == -1 ? print_nrange : nranges);
    out << format("%lld data points, mean %.6g +- %.6g [%.6g,%.6g]\n")
	% countll() % (multiplier * mean()) % (multiplier * stddev())
	% (multiplier * min()) % (multiplier * max());
    if (countll() == 0) return;
    out << format("    quantiles about every %.0f data points:")
	% ((double)countll()/(double)nranges);
    double step = 1.0 / (double)nranges;
    int nquantiles = 0;
    for(double quantile = step;Double::lt(quantile,1.0);quantile += step) {
	if ((nquantiles % 10) == 0) {
	    out << format("\n    %.4g%%: ") % (quantile * 100);
	} else {
	    out << ", ";
	}

	out << format("%.8g") % (multiplier * getQuantile(quantile));
	++nquantiles;
    }
    out << "\n";
}

void StatsQuantile::printTail(FILE *out) {
    ostringstream tmp;
    printTextTail(tmp);
    fwrite(tmp.str().data(), tmp.str().size(), 1, out);
}

// TODO: should this print tails if we have very few data points,
// e.g. 22 data points, should we get the 90%,95% tails?  This happens
// in the dataseries groupby regression test, but may be otherwise
// irrelevant.
void StatsQuantile::printTextTail(ostream &out, double multiplier) const {
    double nentries = countll();
    out << "  tails: ";
    for(double tail_frac = 0.1; (tail_frac * nentries) >= 10.0;) {
	if (tail_frac < 0.05) {
	    out << ", ";
	}
	out << format("%.12g%%: %.8g")
	    % (100*(1-tail_frac)) % (multiplier * getQuantile(1-tail_frac));
	tail_frac /= 2.0;
	out << format(", %.12g%%: %.8g")
	    % (100*(1-tail_frac)) % (multiplier * getQuantile(1-tail_frac));
	tail_frac /= 5.0;
    }
    out << "\n";
}

void StatsQuantile::printText(ostream &out) const {
    printTextRanges(out);
    printTextTail(out);
}

size_t StatsQuantile::memoryUsage() const {
    // primary data (init_buffers) + secondary metadata (init)
    return nbuffers * buffer_size * sizeof(double) // primary
	+ nbuffers * (sizeof(int64_t) + sizeof(int) + sizeof(bool)
		      + sizeof(double) + sizeof(int)) // secondary
	+ sizeof(StatsQuantile);
}


/*
 * LineIterator.hpp
 *
 *  Created on: Nov 13, 2014
 *      Author: Péter Fankhauser
 *   Institute: ETH Zurich, Autonomous Systems Lab
 */

#pragma once

#include "grid_map_core/GridMap.hpp"
#include "grid_map_core/iterators/SubmapIterator.hpp"

// Eigen
#include <Eigen/Core>

namespace grid_map {

/*!
 * Iterator class to iterate over a line in the map.
 * Based on Bresenham Line Drawing algorithm.
 */
class LineIterator
{
public:

  /*!
   * Constructor.
   * @param gridMap the grid map to iterate on.
   * @param start the starting index of the line.
   * @param end the ending index of the line.
   */
  LineIterator(const grid_map::GridMap& gridMap, const Eigen::Array2i& start, const Eigen::Array2i& end);

  /*!
   * Constructor.
   * @param gridMap the grid map to iterate on.
   * @param start the starting point of the line.
   * @param end the ending point of the line.
   */
  LineIterator(const grid_map::GridMap& gridMap, const Eigen::Vector2d& start, const Eigen::Vector2d& end);

  /*!
   * Assignment operator.
   * @param iterator the iterator to copy data from.
   * @return a reference to *this.
   */
  LineIterator& operator =(const LineIterator& other);

  /*!
   * Compare to another iterator.
   * @return whether the current iterator points to a different address than the other one.
   */
  bool operator !=(const LineIterator& other) const;

  /*!
   * Dereference the iterator with const.
   * @return the value to which the iterator is pointing.
   */
  const Eigen::Array2i& operator *() const;

  /*!
   * Increase the iterator to the next element.
   * @return a reference to the updated iterator.
   */
  LineIterator& operator ++();

  /*!
   * Indicates if iterator is past end.
   * @return true if iterator is out of scope, false if end has not been reached.
   */
  bool isPastEnd() const;

private:

  // TODO
  void initializeParameters();

  //! Current index.
  Eigen::Array2i index_;

  //! Starting index of the line.
  Eigen::Array2i start_;

  //! Ending index of the line.
  Eigen::Array2i end_;

  //! Current cell number.
  unsigned int iCell_;

  //! Number of cells in the line.
  unsigned int nCells_;

  //! Helper variables for Bresenham Line Drawing algorithm.
  Eigen::Array2i increment1_, increment2_;
  int denominator_, numerator_, numeratorAdd_;

  //! Map information needed to get position from iterator.
  Eigen::Array2d mapLength_;
  Eigen::Vector2d mapPosition_;
  double resolution_;
  Eigen::Array2i bufferSize_;
  Eigen::Array2i bufferStartIndex_;
};

} /* namespace */


// Sampling code

#include <ctime>
#include <iostream>
#include <fstream>
#include <sstream>
#include <string>
#include <Eigen/Dense>
#include "state.hpp"

#define USAGE "./infer directory [samples [burn [thin [proposals]]]]"

// Default params
#define SAMPLES 24000
#define BURN 6000
#define THIN 6
#define PROPOSALS 6

using namespace Eigen;
using namespace std;

///////////////////////
// Sampling function //
///////////////////////

// Given the specified sampling parameters, performs MCMC updates and
// computes the posterior mean of object coocurrence/strengths
void infer(State &state, int samples, int burn, int thin, int proposals,
            ArrayXXd &avg_cooc, ArrayXXd &avg_strengths)
{
    int N = state.strengths.cols();
    int A = state.strengths.rows();
    avg_cooc = MatrixXd::Identity(N,N); // Only matrix class has Identity
    avg_strengths = ArrayXXd::Zero(A,N);

    // Burn-in
    for (int i=0;i<burn;i++){
        state.update(proposals);
    }

    // Thinned samples from posterior
    for (int i=0;i<samples-burn;i++){
        state.update(proposals);
        if (i % thin == 0){
            // Update avg_cooc
            for (int j=0;j<N;j++) {
            for (int k=j+1;k<N;k++) {
                bool same = state.z(j) == state.z(k);
                avg_cooc(k,j) += (same - avg_cooc(k,j))/(i+1);
            }
            }

            // Update avg_strengths
            ArrayXXd diff = state.strengths - avg_strengths;
            avg_strengths += diff/(i+1);
        }
    }

    // Fill in other side of cooc
    for (int j=0;j<N;j++) {
    for (int k=j+1;k<N;k++) {
        avg_cooc(j,k) = avg_cooc(k,j);
    }
    }
}

//////////////////
// IO Functions //
//////////////////

void events_from_file(string file, ArrayXXi &pos, ArrayXXi &neg)
{
    // First we compute the size of the file (we assume correct formatting)
    ifstream in(file.c_str());
    string line,field;
    getline(in,line);

    // Get the first line to determine the number of cols
    int cols = 0;
    stringstream firstrow(line);
    while (getline(firstrow,field,',')){
        cols++;
    }

    // Get the rest of the lines to determine the number of rows
    int rows = 1;
    while (getline(in,line)){
        rows++;
    }
    in.close();

    // Now we rescan the file, and split the data into pos/neg event counts
    // (Note that data is stored tranposed relative to CSV layout).
    in.open(file.c_str());
    pos.resize(cols/2,rows);
    neg.resize(cols/2,rows);
    for (int r=0;r<rows;r++){
        getline(in,line);
        stringstream stream(line);
        bool even = true;

        for (int c=0;c<cols;c++){
            getline(stream,field,',');
            if (even){
                pos(c/2,r) = atoi(field.c_str());
            } else {
                neg(c/2,r) = atoi(field.c_str());
            }
            even = !even;
        }
    }
    in.close();
}

// Writes an coccurence array to standard error
void array_to_file(string file, const Ref<const ArrayXXd> &arr)
{
    ofstream out(file.c_str());
    int C = arr.cols();
    int R = arr.rows();

    for (int c=0;c<C;c++){
        for (int r=0;r<R;r++){
            out << arr(r,c);
            if ( R-r > 1 ){
                out << ',';
            }
        }
        out << endl;
    }
}

///////////////////
// Main function //
///////////////////

int main(int argc, char *argv[])
{
    // Parse the options
    if (argc < 2 || argc > 6) {
        cout << USAGE << endl;
        return 1;
    }

    // Sampling parameters
    int samples = SAMPLES;
    int burn = BURN;
    int thin = THIN;
    int proposals = PROPOSALS;
    if (argc > 2){
        samples = atoi(argv[2]);
    if (argc > 3){
        burn = atoi(argv[3]);
    if (argc > 4){
        thin = atoi(argv[4]);
    if (argc > 5){
        proposals = atoi(argv[5]);
    }}}}

    // Directory location
    string dir = argv[1];
    if (dir[dir.length()] != '/'){ // Kind of a hack...
        dir += '/';
    }
    string events_file = dir + "events.csv";

    // Import the data
    ArrayXXi pos,neg;
    events_from_file(events_file,pos,neg);

    // Create the particle
    gsl_rng *rng = gsl_rng_alloc(gsl_rng_default);
    State particle(rng,pos,neg);

    // Sample and compute the posterior co-occurence/strengths
    ArrayXXd avg_cooc,avg_strengths;
    clock_t begin = clock();
    infer(particle,samples,burn,thin,proposals,avg_cooc,avg_strengths);
    double duration = float(clock() - begin) / CLOCKS_PER_SEC;
    cout << "Took : " << duration << " seconds" << endl;
    cout << "Final state: " << endl;
    particle.print_vars(true);

    // Write the posterior averages to disk
    string cooc_file = dir + "cooc.csv";
    array_to_file(cooc_file,avg_cooc);

    string strengths_file = dir + "strengths.csv";
    array_to_file(strengths_file,avg_strengths);

    gsl_rng_free(rng);
    return 0;
}


#include "msl_raptor_backend.h"
#include "opencv2/highgui/highgui.hpp"
#include "opencv2/imgproc/imgproc.hpp"
#include <Eigen/Core>
#include <Eigen/Dense>
#include <chrono>
#include <iostream>
#include <opencv2/core/eigen.hpp>
#include <type_traits>

int main(int argc, char **argv)
{
    // Define if we want to used an aligned bb (4 parameters) or an angled one (5 parameters).
    const bool aligned_bb = false;

    // Define camera parameters (camera matrix, extrinsics, distorsion coefficients) and create camera params object.
    std::vector<float> cam_mat{304.7262878417969, 0.0, 214.2415313720703,
                               0.0, 304.7262878417969, 121.70726013183594,
                               0.0, 0.0, 1.0};
    std::vector<float> rvec(3, 0.0);
    std::vector<float> tvec(3, 0.0);
    std::vector<float> dist_coeffs(4, 0);
    msl_raptor_backend::CameraParams cam_params(cam_mat, rvec, tvec, dist_coeffs);

    // Define object parameters (UKF initial covariances, process noise, and measurement noise, dimensions) and create object params object.
    std::vector<double> sigma{0.0000000000001, 0.000001, 0.000001,  // position
                              0.00000000000005, 0.000005, 0.000005, //  linear velocity
                              0.000003, 0.000003, 0.000003,         // angular velocity
                              0.000005, 0.000005, 0.00005};         // orientation

    std::vector<double> process_noise{0.002, 0.0002, 0.0002,     // position
                                      0.0001, 0.0000001, 0.0001, // linear velocity
                                      0.0006, 0.000006, 0.00006, // angular velocity
                                      0.00001, 0.0001, 0.00001}; // orientation

    // measurement vector for angle bb is center x, center y, width, height, angle.
    std::vector<double> meas_noise{1, 1, 2, 1, 0.5};

    // Parameters for optimisation-based pose estimation from bb
    double momentum = 0.9;                                         // Momentum parameter
    int max_steps = 500;                                           // Maximum number of steps
    int conv_steps = 4;                                            // Steps to check convergence
    int period_lower_lr = 20;                                      // How often to reduce learning rate
    std::vector<msl_raptor_backend::PoseVec> bb_init_pose_guesses; // Contains position vector and quaternion for orientation
    // Add a grid of initial guesses
    for (double i = -2; i < 2; i += 1)
    {
        for (double j = -2; j < 2; j += 1)
        {
            for (double k = 2; k < 8; k += 1)
            {
                bb_init_pose_guesses.push_back(msl_raptor_backend::PoseVec(i, j, k, 0, 0, 0, 1)); // Position vector and orientation quaternion
                // std::cout << msl_raptor_backend::PoseVec(i, j, k, 0, 0, 0, 1) << std::endl;
            }
        }
    }
    std::vector<double> bb_init_step_size{0.001, 0.001, 0.01, 0.001, 0.001, 0.001};    // Step sizes along position and Euler axis for orientation
    std::vector<double> bb_init_lr{0.0005, 0.0005, 0.0001, 0.00005, 0.00005, 0.00005}; // Step sizes for gradient approximation along position and Euler axis for orientation
    msl_raptor_backend::ObjPoseInitParams obj_pose_init_params(bb_init_step_size, bb_init_lr, bb_init_pose_guesses, momentum, max_steps, conv_steps, period_lower_lr);

    // Provides width, height and length of object in meters, from which corners of a 3D box are used to approximate the object's shape.
    msl_raptor_backend::ObjParams obj_params(0.3, 0.2, 0.5, sigma, process_noise, meas_noise, obj_pose_init_params);

    // Initial state, here places the object 5 meters away in front with no velocity.
    MSLRaptorUKF<aligned_bb>::StateVec true_state, state_approx_heuristic, state_approx_optim;
    true_state << 0, 0, 5.2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0;
    true_state.quat(0).setIdentity();

    // Initialize MSL-RAPTOR back-end UKF using set object parameters, camera parameters, and true state.
    MSLRaptorUKF<aligned_bb> msl_raptor_ukf(obj_params, cam_params, true_state);
    std::cout << "True pose is " << std::endl
              << msl_raptor_ukf.stateToPose(true_state) << std::endl
              << std::endl;

    // Check the measurement process applied to the current UKF state with no input, which returns a bounding box prediction.
    MSLRaptorUKF<aligned_bb>::MeasureVec m = msl_raptor_ukf.H(true_state, Eigen::Affine3d());

    std::cout << "Measurement at true state" << std::endl
              << m << std::endl
              << std::endl;

    // Example of approximating a state from a 2D bounding box with heuristics
    state_approx_heuristic = msl_raptor_ukf.approxStateFromBbHeuristic(m);
    std::cout << "Pose approximated only from 2D bounding box with a heuristic " << std::endl
              << msl_raptor_ukf.stateToPose(state_approx_heuristic) << std::endl
              << std::endl;

    // Example of approximating a state from a 2D bounding box with optimisation
    double optim_error;
    std::tie(optim_error, state_approx_optim) = msl_raptor_ukf.approxStatePoseOptim(m);
    std::cout << "Pose approximated only from 2D bounding box with optimisation (err " << optim_error << ") "
              << std::endl
              << msl_raptor_ukf.stateToPose(state_approx_optim) << std::endl
              << std::endl;

    // Modify the measurement width to make it appear closer
    m(2) *= 4;
    m(3) *= 4;
    std::cout << "What if we observed a bounding box with a larger size. This should indicate a closer object." << std::endl
              << m << std::endl
              << std::endl;

    // Update the UKF with the modified measurement, which should make the pose appear ...
    msl_raptor_ukf.update(0.03, m);

    // Check the new state
    std::cout << "Predicted pose after measurement update is " << std::endl
              << msl_raptor_ukf.stateToPose(msl_raptor_ukf.getState()) << "\n\n";

    // Create a fake input (could be from odometry in practice)
    Eigen::Affine3d input = Eigen::AngleAxisd(0.05, Eigen::Vector3d::UnitZ()) *
                            Eigen::Translation3d(0.0, 0.0, 0.2);
    std::cout << "What if we apply a translation and orientation input with the measurement \n \n";

    // Update the UKF with the modified measurement and fake input
    msl_raptor_ukf.update(0.03, m, input);

    // Check the new state
    std::cout << "The new pose is " << std::endl
              << msl_raptor_ukf.stateToPose(msl_raptor_ukf.getState()) << "\n\n";
}


/* ----------------------------------------------------------------------- *//**
 *
 * @file multilogistic.cpp
 *
 * @brief Multinomial Logistic-Regression functions
 *
 * We implement the iteratively-reweighted-least-squares method.
 *
 *//* ----------------------------------------------------------------------- */

#include <dbconnector/dbconnector.hpp>

#include <modules/shared/HandleTraits.hpp>
#include <modules/prob/boost.hpp>
#include <boost/math/distributions.hpp>
#include <modules/prob/student.hpp>
#include "multilogistic.hpp"

#include <vector>
#include <fstream>

namespace madlib {

// Use Eigen
using namespace dbal::eigen_integration;

namespace modules {

// Import names from other MADlib modules
using dbal::NoSolutionFoundException;

namespace regress {

// Internal functions
AnyType mLogstateToResult(
    const Allocator &inAllocator,
    int ref_category,
    const HandleMap<const ColumnVector, TransparentHandle<double> >& inCoef,
    const ColumnVector &diagonal_of_heissian,
    double logLikelihood,
    double conditionNo);

/**
 * @brief Logistic function
 */
inline double sigma(double x) {
    return 1. / (1. + std::exp(-x));
}

/**
 * @brief Inter- and intra-iteration state for iteratively-reweighted-least-
 *        squares method for logistic regression
 *
 * TransitionState encapsualtes the transition state during the
 * logistic-regression aggregate function. To the database, the state is
 * exposed as a single DOUBLE PRECISION array, to the C++ code it is a proper
 * object containing scalars, a vector, and a matrix.
 *
 * Note: We assume that the DOUBLE PRECISION array is initialized by the
 * database with length at least 4, and all elemenets are 0.
 */
template <class Handle>
class MLogRegrIRLSTransitionState {

    // By §14.5.3/9: "Friend declarations shall not declare partial
    // specializations." We do access protected members in operator+=().
    template <class OtherHandle>
    friend class MLogRegrIRLSTransitionState;

public:

    MLogRegrIRLSTransitionState(const AnyType &inArray)
        : mStorage(inArray.getAs<Handle>()) {
        rebind(static_cast<uint16_t>(mStorage[0]),static_cast<uint16_t>(mStorage[1]));
    }

    /**
     * @brief Convert to backend representation
     *
     * We define this function so that we can use State in the
     * argument list and as a return type.
     */
    inline operator AnyType() const {
        return mStorage;
    }

    /**
     * @brief Initialize the iteratively-reweighted-least-squares state.
     *
     * This function is only called for the first iteration, for the first row.
     */
    inline void initialize(
        const Allocator &inAllocator,
        uint16_t inWidthOfX,
        uint16_t inNumCategories, uint16_t inRefCategory) {

        mStorage = inAllocator.allocateArray<double, dbal::AggregateContext,
            dbal::DoZero, dbal::ThrowBadAlloc>(arraySize(inWidthOfX, inNumCategories));
        rebind(inWidthOfX, inNumCategories);
        widthOfX = inWidthOfX;
        numCategories = inNumCategories;
        ref_category = inRefCategory;
    }

    /**
     * @brief We need to support assigning the previous state
     */
    template <class OtherHandle> MLogRegrIRLSTransitionState &operator=(
        const MLogRegrIRLSTransitionState<OtherHandle> &inOtherState) {

        for (size_t i = 0; i < mStorage.size(); i++)
            mStorage[i] = inOtherState.mStorage[i];
        return *this;
    }

    /**
     * @brief Merge with another State object by copying the intra-iteration
     *     fields
     */
    template <class OtherHandle> MLogRegrIRLSTransitionState &operator+=(
        const MLogRegrIRLSTransitionState<OtherHandle> &inOtherState) {
        if (mStorage.size() != inOtherState.mStorage.size() ||
            widthOfX != inOtherState.widthOfX)
            throw std::logic_error("Internal error: Incompatible transition "
                "states");

        numRows += inOtherState.numRows;
        gradient += inOtherState.gradient;
        X_transp_AX += inOtherState.X_transp_AX;
        logLikelihood += inOtherState.logLikelihood;
        return *this;
    }

    /**
     * @brief Reset the inter-iteration fields.
     */
    inline void reset() {
        numRows = 0;
        gradient.fill(0);
        X_transp_AX.fill(0);
        logLikelihood = 0;
    }

private:
    static inline uint32_t arraySize(const uint16_t inWidthOfX,
        const uint16_t inNumCategories) {
        return 5 + inWidthOfX * inWidthOfX * inNumCategories * inNumCategories
                                 + 2 * inWidthOfX * inNumCategories;
    }

    /**
     * @brief Rebind to a new storage array
     *
     * @param inWidthOfX             The number of independent variables.
     * @param inNumCategories The number of categories of the dependant var

     * Array layout (iteration refers to one aggregate-function call):
     * Inter-iteration components (updated in final function):
     * - 0: widthOfX (number of independant variables)
     * - 1: numCategories (number of categories)
     * - 2: coef (vector of coefficients)
     *
     * Intra-iteration components (updated in transition step):
     * - 2 + widthOfX*numCategories: numRows (number of rows already processed in this iteration)
     * - 3 + widthOfX*numCategories: gradient (X^T A z)
     * - 3 + 2 * widthOfX * inNumCategories: X_transp_AX (X^T A X)
     * - 3 + widthOfX^2*numCategories^2
                         + 2 * widthOfX*numCategories: logLikelihood ( ln(l(c)) )
     * - 4 + widthOfX^2*numCategories^2
                         + 2 * widthOfX*numCategories: ref_category
     */
    void rebind(uint16_t inWidthOfX = 0, uint16_t inNumCategories = 0) {
        widthOfX.rebind(&mStorage[0]);
        numCategories.rebind(&mStorage[1]);
        coef.rebind(&mStorage[2], inWidthOfX*inNumCategories);

        numRows.rebind(&mStorage[2 + inWidthOfX*inNumCategories]);

        gradient.rebind(&mStorage[3 + inWidthOfX*inNumCategories],inWidthOfX*inNumCategories);
        X_transp_AX.rebind(&mStorage[3 + 2 * inWidthOfX*inNumCategories],
            inNumCategories*inWidthOfX, inWidthOfX*inNumCategories);
        logLikelihood.rebind(&mStorage[3 +
             inNumCategories*inNumCategories*inWidthOfX*inWidthOfX
             + 2 * inWidthOfX*inNumCategories]);
        ref_category.rebind(&mStorage[4 +
             inNumCategories*inNumCategories*inWidthOfX*inWidthOfX
             + 2 * inWidthOfX*inNumCategories]);
    }

    Handle mStorage;

public:
    typename HandleTraits<Handle>::ReferenceToUInt16 widthOfX;
    typename HandleTraits<Handle>::ReferenceToUInt16 numCategories;

    typename HandleTraits<Handle>::ColumnVectorTransparentHandleMap coef;
    typename HandleTraits<Handle>::ReferenceToUInt64 numRows;
    typename HandleTraits<Handle>::ColumnVectorTransparentHandleMap gradient;
    typename HandleTraits<Handle>::MatrixTransparentHandleMap X_transp_AX;
    typename HandleTraits<Handle>::ReferenceToDouble logLikelihood;
    typename HandleTraits<Handle>::ReferenceToUInt16 ref_category;
};

/**
 * @brief Inter- and intra-iteration state for robust variance calculations
 *
 * TransitionState encapsualtes the transition state during the
 * logistic-regression robust variance calculation. To the database, the state is
 * exposed as a single DOUBLE PRECISION array, to the C++ code it is a proper
 * object containing scalars, a vector, and a matrix.
 *
 * Note: We assume that the DOUBLE PRECISION array is initialized by the
 * database with length at least 4, and all elemenets are 0.
 */
template <class Handle>
class MLogRegrRobustTransitionState {

    // By §14.5.3/9: "Friend declarations shall not declare partial
    // specializations." We do access protected members in operator+=().
    template <class OtherHandle>
    friend class MLogRegrRobustTransitionState;

public:
    MLogRegrRobustTransitionState(const AnyType &inArray)
        : mStorage(inArray.getAs<Handle>()) {
        rebind(static_cast<uint16_t>(mStorage[0]),static_cast<uint16_t>(mStorage[1]));
    }

    /**
     * @brief Convert to backend representation
     *
     * We define this function so that we can use State in the
     * argument list and as a return type.
     */
    inline operator AnyType() const {
        return mStorage;
    }

    /**
     * @brief Initialize the iteratively-reweighted-least-squares state.
     *
     * This function is only called for the first iteration, for the first row.
     */
    inline void initialize(const Allocator &inAllocator,
        uint16_t inWidthOfX, uint16_t inNumCategories, uint16_t inRefCategory) {
        mStorage = inAllocator.allocateArray<double, dbal::AggregateContext,
            dbal::DoZero, dbal::ThrowBadAlloc>(arraySize(inWidthOfX, inNumCategories));
        rebind(inWidthOfX, inNumCategories);
        widthOfX = inWidthOfX;
        numCategories = inNumCategories;
        ref_category = inRefCategory;
    }

    /**
     * @brief We need to support assigning the previous state
     */
    template <class OtherHandle> MLogRegrRobustTransitionState &operator=(
        const MLogRegrRobustTransitionState<OtherHandle> &inOtherState) {
        for (size_t i = 0; i < mStorage.size(); i++)
            mStorage[i] = inOtherState.mStorage[i];
        return *this;
    }

    /**
     * @brief Merge with another State object by copying the intra-iteration
     *     fields
     */
    template <class OtherHandle> MLogRegrRobustTransitionState &operator+=(
        const MLogRegrRobustTransitionState<OtherHandle> &inOtherState) {
        if (mStorage.size() != inOtherState.mStorage.size() ||
            widthOfX != inOtherState.widthOfX)
            throw std::logic_error("Internal error: Incompatible transition "
                "states");

        numRows += inOtherState.numRows;
        X_transp_AX += inOtherState.X_transp_AX;
        meat += inOtherState.meat;
        return *this;
    }

    /**
     * @brief Reset the inter-iteration fields.
     */
    inline void reset() {
        numRows = 0;
        meat.fill(0);
        X_transp_AX.fill(0);
    }

private:
    static inline uint32_t arraySize(const uint16_t inWidthOfX,
        const uint16_t inNumCategories) {
        return 4 + 2*inWidthOfX * inWidthOfX * inNumCategories * inNumCategories
                                 + inWidthOfX * inNumCategories;
    }

    /**
     * @brief Rebind to a new storage array
     *
     * @param inWidthOfX             The number of independent variables.
     * @param inNumCategories The number of categories of the dependant var

     * Array layout (iteration refers to one aggregate-function call):
     * Inter-iteration components (updated in final function):
     * - 0: widthOfX (number of independant variables)
     * - 1: numCategories (number of categories)
     * - 2: ref_category
     * - 3: coef (vector of coefficients)
     *
     * Intra-iteration components (updated in transition step):
     * - 3 + widthOfX*numCategories: numRows (number of rows already processed in this iteration)
     * - 4 + widthOfX * inNumCategories: X_transp_AX (X^T A X).
     * - 4 + widthOfX^2*numCategories^2: meat  (The meat matrix)
     */
    void rebind(uint16_t inWidthOfX = 0, uint16_t inNumCategories = 0) {
        widthOfX.rebind(&mStorage[0]);
        numCategories.rebind(&mStorage[1]);
        ref_category.rebind(&mStorage[2]);
        coef.rebind(&mStorage[3], inWidthOfX * inNumCategories);
        numRows.rebind(&mStorage[3 + inWidthOfX * inNumCategories]);
        X_transp_AX.rebind(&mStorage[4 + inWidthOfX * inNumCategories],
            inNumCategories * inWidthOfX, inWidthOfX * inNumCategories);
        meat.rebind(&mStorage[4 +
             inNumCategories * inNumCategories * inWidthOfX * inWidthOfX
             + inWidthOfX * inNumCategories], inWidthOfX * inNumCategories, inWidthOfX * inNumCategories);
    }

    Handle mStorage;

public:
    typename HandleTraits<Handle>::ReferenceToUInt16 widthOfX;
    typename HandleTraits<Handle>::ReferenceToUInt16 numCategories;

    typename HandleTraits<Handle>::ColumnVectorTransparentHandleMap coef;
    typename HandleTraits<Handle>::ReferenceToUInt64 numRows;
    typename HandleTraits<Handle>::MatrixTransparentHandleMap X_transp_AX;
    typename HandleTraits<Handle>::MatrixTransparentHandleMap meat;
    typename HandleTraits<Handle>::ReferenceToUInt16 ref_category;
};

/**
 * @brief IRLS Transition
 * @param args
 *
 * Arguments (Matched with PSQL wrapped)
 * - 0: Current State
 * - 1: y value (Integer)
 * - 2: numCategories (Integer)
 * - 3: ref_category (Integer)
 * - 4: X value (Column Vector)
 * - 5: Previous State

 */
AnyType
__mlogregr_irls_step_transition::run(AnyType &args) {
    MLogRegrIRLSTransitionState<MutableArrayHandle<double> > state = args[0];

    // Get x as a vector of double
    MappedColumnVector x = args[4].getAs<MappedColumnVector>();

    // Get the category & numCategories as integer
    int32_t category = args[1].getAs<int32_t>();
    // Number of categories after pivoting (We pivot around the first category)
    int32_t numCategories = (args[2].getAs<int32_t>() - 1);
    int32_t ref_category = args[3].getAs<int32_t>();

    // The following check was added with MADLIB-138.
    if (!x.is_finite())
            throw std::domain_error("Design matrix is not finite.");

    if (state.numRows == 0) {
        if (x.size() > std::numeric_limits<uint16_t>::max())
                throw std::domain_error("Number of independent variables cannot be "
                        "larger than 65535.");

        if (numCategories < 1)
                throw std::domain_error("Number of cateogires must be at least 2");

        // Init the state (requires x.size() and category.size())
        state.initialize(*this,
            static_cast<uint16_t>(x.size()) ,
            static_cast<uint16_t>(numCategories),
            static_cast<uint16_t>(ref_category));

        if (!args[5].isNull()) {
                MLogRegrIRLSTransitionState<ArrayHandle<double> >
                    previousState = args[5];
                state = previousState;
                state.reset();
        }
    }

    /*
     * This check should be done for each iteration. Only checking the first
     * run is not enough.
     */
    if (category > numCategories || category < 0)
        throw std::domain_error("Invalid category. Categories must be integer "
            "values between 0 and (number of categories - 1).");

    if (ref_category > numCategories || ref_category < 0)
        throw std::domain_error("Invalid reference category. Reference category must be integer "
            "value between 0 and (number of categories - 1).");

    // Now do the transition step
    state.numRows++;
    /*     Get y: Convert to 1/0 boolean vector
            Example: Category 4 : 0 0 0 1 0 0
            Storing it in this forms helps us get a nice closed form expression
    */

    //To pivot around the specified reference category
    ColumnVector y(numCategories);
    y.fill(0);
    if (category > ref_category) {
        y(category - 1) = 1;
    } else if (category < ref_category) {
        y(category) = 1;
    }

    /*
    Compute the parameter vector (the 'pi' vector in the documentation)
    for the data point being processed.
    Casting the coefficients into a matrix makes the calculation simple.
    */
    Matrix coef = state.coef;
    coef.resize(numCategories, state.widthOfX);

    //Store the intermediate calculations because we'll reuse them in the LLH
    ColumnVector t1 = x; //t1 is vector of size state.widthOfX
    t1 = coef*x;
    /* Note: The above 2 lines could have been written as:
        ColumnVector t1 = -coef*x;

        but this creates warnings. These warnings are somehow related to the factor
        that x is of an immutable type. The following alternative could resolve
        the warnings (although not necessary the best alternative):
    */

    ColumnVector t2 = t1.array().exp();
    double t3 = 1 + t2.sum();
    ColumnVector pi = t2/t3;

    //The gradient matrix has numCategories rows and widthOfX columns
    Matrix grad = -y*x.transpose() + pi*x.transpose();
    //We cast the gradient into a vector to make the Newton step calculations much easier.
    grad.resize(numCategories*state.widthOfX,1);

    /*
         a is a matrix of size JxJ where J is the number of categories
         a_j1j2 = -pi(j1)*(1-pi(j2))if j1 == j2
         a_j1j2 =  pi(j1)*pi(j2) if j1 != j2
    */
    Matrix a(numCategories,numCategories);
    // Compute the 'a' matrix.
    Matrix piDiag = pi.asDiagonal();
    a = pi * pi.transpose() - piDiag;
    state.gradient.noalias() += grad;

    //Start the Hessian calculations
    Matrix X_transp_AX(numCategories * state.widthOfX, numCategories * state.widthOfX);

    /*
        Again: The following 3 lines could have been written as
        Matrix XXTrans = x * x.transpose();
        but it creates warnings related to the type of x. Here is an easy fix
    */
    Matrix cv_x = x;
    Matrix XXTrans = trans(cv_x);
    XXTrans = cv_x * XXTrans;

    //Eigen doesn't supported outer-products for matrices, so we have to do our own.
    //This operation is also known as a tensor-product.
    for (int i1 = 0; i1 < state.widthOfX; i1++){
         for (int i2 = 0; i2 <state.widthOfX; i2++){
            int rowOffset = numCategories * i1;
            int colOffset = numCategories * i2;

            X_transp_AX.block(rowOffset, colOffset, numCategories,  numCategories) = XXTrans(i1,i2)*a;
        }
    }

    triangularView<Lower>(state.X_transp_AX) += X_transp_AX;

    state.logLikelihood += y.transpose()*t1 - log(t3);

    return state;

}

/**
 * @brief Perform the perliminary aggregation function: Merge transition states
     This merge step is the same as that of logistic regression.
 */
AnyType
__mlogregr_irls_step_merge_states::run(AnyType &args) {
    MLogRegrIRLSTransitionState<MutableArrayHandle<double> > stateLeft = args[0];
    MLogRegrIRLSTransitionState<ArrayHandle<double> > stateRight = args[1];

    // We first handle the trivial case where this function is called with one
    // of the states being the initial state
    if (stateLeft.numRows == 0)
        return stateRight;
    else if (stateRight.numRows == 0)
        return stateLeft;

    // Merge states together and return
    stateLeft += stateRight;
    return stateLeft;
}

/**
 * @brief Perform the logistic-regression final step
 */
AnyType
__mlogregr_irls_step_final::run(AnyType &args) {
    // We request a mutable object.
    // Depending on the backend, this might perform a deep copy.
    MLogRegrIRLSTransitionState<MutableArrayHandle<double> > state = args[0];

    // Aggregates that haven't seen any data just return Null.
    if (state.numRows == 0)
        return Null();

    // See MADLIB-138. At least on certain platforms and with certain versions,
    // LAPACK will run into an infinite loop if pinv() is called for non-finite
    // matrices. We extend the check also to the dependent variables.
    if (!state.X_transp_AX.is_finite() || !state.gradient.is_finite())
        throw NoSolutionFoundException("Over- or underflow in intermediate "
            "calulation. Input data is likely of poor numerical condition.");

    SymmetricPositiveDefiniteEigenDecomposition<Matrix> decomposition(
        -1 * state.X_transp_AX, EigenvaluesOnly, ComputePseudoInverse);

    // Precompute (X^T * A * X)^-1
    Matrix heissianInver = -1 * decomposition.pseudoInverse();

    state.coef.noalias() += heissianInver * state.gradient;

    if(!state.coef.is_finite())
        throw NoSolutionFoundException("Over- or underflow in Newton step, "
            "while updating coefficients. Input data is likely of poor "
            "numerical condition.");

    // We use the intra-iteration field gradient for storing the diagonal
    // of X^T A X, so that we don't have to recompute it in the result function.
    // Likewise, we store the condition number.
    // FIXME: This feels a bit like a hack.
    state.gradient = -1 * heissianInver.diagonal();
    state.X_transp_AX(0,0) = decomposition.conditionNo();

    return state;
}

AnyType
mlogregr_robust_step_transition::run(AnyType &args) {
    using std::endl;

	MLogRegrRobustTransitionState<MutableArrayHandle<double> > state = args[0];
    // Get x as a vector of double
    MappedColumnVector x = args[4].getAs<MappedColumnVector>();
    // Get the category & numCategories as integer
    int16_t category = args[1].getAs<int>();
    // Number of categories after pivoting (We pivot around the first category)
    int16_t numCategories = (args[2].getAs<int>() - 1);
    int32_t ref_category = args[3].getAs<int32_t>();
	MappedColumnVector coefVec = args[5].getAs<MappedColumnVector>();

    // The following check was added with MADLIB-138.
    if (!x.is_finite())
            throw std::domain_error("Design matrix is not finite.");

    if (state.numRows == 0) {
        if (x.size() > std::numeric_limits<uint16_t>::max())
                throw std::domain_error("Number of independent variables cannot be "
                        "larger than 65535.");
        if (numCategories < 1)
                throw std::domain_error("Number of cateogires must be at least 2");
        if (category > numCategories)
                throw std::domain_error("You have entered a category > numCategories"
                    "Categories must be of values {0,1... numCategories-1}");

        // Init the state (requires x.size() and category.size())
        state.initialize(*this,
            static_cast<uint16_t>(x.size()) ,
            static_cast<uint16_t>(numCategories),
            static_cast<uint16_t>(ref_category));

        state.coef = coefVec;
    }

    // Now do the transition step
    state.numRows++;
    /*
        Get y: Convert to 1/0 boolean vector
        Example: Category 4 : 0 0 0 1 0 0
        Storing it in this forms helps us get a nice closed form expression
    */

    //To pivot around the specified reference category
    ColumnVector y(numCategories);
    y.fill(0);
    if (category > ref_category) {
        y(category - 1) = 1;
    } else if (category < ref_category) {
        y(category) = 1;
    }

    /*
    Compute the parameter vector (the 'pi' vector in the documentation)
    for the data point being processed.
    Casting the coefficients into a matrix makes the calculation simple.
    */

    Matrix coef = state.coef;
    coef.resize(numCategories, state.widthOfX);

    //Store the intermediate calculations because we'll reuse them in the LLH
    ColumnVector t1 = x; //t1 is vector of size state.widthOfX
    t1 = coef*x;
    /*
        Note: The above 2 lines could have been written as:
        ColumnVector t1 = -coef*x;

        but this creates warnings. These warnings are somehow related to the factor
        that x is an immutable type.
    */

    ColumnVector t2 = t1.array().exp();
    double t3 = 1 + t2.sum();
    ColumnVector pi = t2/t3;

	//The gradient matrix has numCategories rows and widthOfX columns
    Matrix grad = -y * x.transpose() + pi * x.transpose();
    //We cast the gradient into a vector to make the math easier.
    grad.resize(numCategories * state.widthOfX, 1);

    Matrix GradGradTranspose;
    GradGradTranspose = grad * grad.transpose();
	state.meat += GradGradTranspose;

    /*
         a is a matrix of size JxJ where J is the number of categories
         a_j1j2 = -pi(j1)*(1-pi(j2))if j1 == j2
         a_j1j2 =  pi(j1)*pi(j2) if j1 != j2
    */
    // Compute the 'a' matrix.
    Matrix a(numCategories,numCategories);
    Matrix piDiag = pi.asDiagonal();
    a = pi * pi.transpose() - piDiag;

    //Start the Hessian calculations
    Matrix X_transp_AX(numCategories * state.widthOfX, numCategories * state.widthOfX);

    /*
        Again: The following 3 lines could have been written as
        Matrix XXTrans = x * x.transpose();
        but it creates warnings related to the type of x. Here is an easy fix
    */
    Matrix cv_x = x;
    Matrix XXTrans = trans(cv_x);
    XXTrans = cv_x * XXTrans;

    //Eigen doesn't supported outer-products for matrices, so we have to do our own.
    //This operation is also known as a tensor-product.
    for (int i1 = 0; i1 < state.widthOfX; i1++){
         for (int i2 = 0; i2 <state.widthOfX; i2++){
            int rowOffset = numCategories * i1;
            int colOffset = numCategories * i2;

            X_transp_AX.block(rowOffset, colOffset, numCategories,  numCategories) = XXTrans(i1, i2) * a;
        }
    }

    triangularView<Lower>(state.X_transp_AX) += X_transp_AX;

    return state;

}

/**
 * @brief Perform the perliminary aggregation function: Merge transition states
     This merge step is the same as that of logistic regression.
 */
AnyType
mlogregr_robust_step_merge_states::run(AnyType &args) {
    MLogRegrRobustTransitionState<MutableArrayHandle<double> > stateLeft = args[0];
    MLogRegrRobustTransitionState<ArrayHandle<double> > stateRight = args[1];

    // We first handle the trivial case where this function is called with one
    // of the states being the initial state
    if (stateLeft.numRows == 0)
        return stateRight;
    else if (stateRight.numRows == 0)
        return stateLeft;

    // Merge states together and return
    stateLeft += stateRight;
    return stateLeft;
}

AnyType MLrobuststateToResult(
    const Allocator &inAllocator,
    int ref_category,
    const HandleMap<const ColumnVector, TransparentHandle<double> >& inCoef,
    const ColumnVector &diagonal_of_varianceMat) {

	MutableNativeColumnVector variance(
        inAllocator.allocateArray<double>(inCoef.size()));
    MutableNativeColumnVector stdErr(
        inAllocator.allocateArray<double>(inCoef.size()));
    MutableNativeColumnVector waldZStats(
        inAllocator.allocateArray<double>(inCoef.size()));
    MutableNativeColumnVector waldPValues(
        inAllocator.allocateArray<double>(inCoef.size()));

    for (Index i = 0; i < inCoef.size(); ++i) {
        variance(i) = diagonal_of_varianceMat(i);
        stdErr(i) = std::sqrt(diagonal_of_varianceMat(i));
        waldZStats(i) = inCoef(i) / stdErr(i);
        waldPValues(i) = 2. * prob::cdf(
            prob::normal(), -std::abs(waldZStats(i)));
    }

    // Return all coefficients, standard errors, etc. in a tuple
    AnyType tuple;
    tuple <<  ref_category << inCoef << stdErr << waldZStats << waldPValues;
    return tuple;
}

/**
 * @brief Perform the logistic-regression final step
 */
AnyType
mlogregr_robust_step_final::run(AnyType &args) {
    // We request a mutable object. Depending on the backend, this might perform
    // a deep copy.
    MLogRegrRobustTransitionState<ArrayHandle<double> > state = args[0];
    // Aggregates that haven't seen any data just return Null.
    if (state.numRows == 0)
        return Null();

    // See MADLIB-138. At least on certain platforms and with certain versions,
    // LAPACK will run into an infinite loop if pinv() is called for non-finite
    // matrices. We extend the check also to the dependent variables.
    if (!state.X_transp_AX.is_finite())
        throw NoSolutionFoundException("Over- or underflow in intermediate "
            "calulation. Input data is likely of poor numerical condition.");

    SymmetricPositiveDefiniteEigenDecomposition<Matrix> decomposition(
        -1 * state.X_transp_AX, EigenvaluesOnly, ComputePseudoInverse);

	for(int i = 0; i < state.X_transp_AX.rows(); i++)
	{
		for(int j = 0; j < state.X_transp_AX.cols(); j++)
		{
			elog(INFO, "Bread %i, %i, %f",  i,j, static_cast<float>(state.X_transp_AX(i,j)));
		}
	}

    // Precompute (X^T * A * X)^-1
    Matrix bread = decomposition.pseudoInverse();
	Matrix varianceMat;
    varianceMat = bread * state.meat * bread;

    if(!state.coef.is_finite())
        throw NoSolutionFoundException("Over- or underflow in Newton step, "
            "while updating coefficients. Input data is likely of poor "
            "numerical condition.");

    return MLrobuststateToResult(*this, state.ref_category, state.coef, varianceMat.diagonal());
}

/**
 * @brief Return the difference in log-likelihood between two states
 */
AnyType
__internal_mlogregr_irls_step_distance::run(AnyType &args) {
    MLogRegrIRLSTransitionState<ArrayHandle<double> > stateLeft = args[0];
    MLogRegrIRLSTransitionState<ArrayHandle<double> > stateRight = args[1];

    return std::abs(stateLeft.logLikelihood - stateRight.logLikelihood);
}

/**
 * @brief Return the coefficients and diagnostic statistics of the state
 */
AnyType
__internal_mlogregr_irls_result::run(AnyType &args) {
    MLogRegrIRLSTransitionState<ArrayHandle<double> > state = args[0];

    return mLogstateToResult(*this, state.ref_category, state.coef,
        state.gradient, state.logLikelihood, state.X_transp_AX(0,0));
}

/**
 * @brief Compute the diagnostic statistics
 *
 * This function wraps the common parts of computing the results for IRLS.
 */
AnyType mLogstateToResult(
    const Allocator &inAllocator,
    int ref_category,
    const HandleMap<const ColumnVector, TransparentHandle<double> > &inCoef,
    const ColumnVector &diagonal_of_heissian,
    double logLikelihood,
    double conditionNo) {

    MutableNativeColumnVector stdErr(
        inAllocator.allocateArray<double>(inCoef.size()));
    MutableNativeColumnVector waldZStats(
        inAllocator.allocateArray<double>(inCoef.size()));
    MutableNativeColumnVector waldPValues(
        inAllocator.allocateArray<double>(inCoef.size()));
    MutableNativeColumnVector oddsRatios(
        inAllocator.allocateArray<double>(inCoef.size()));

    for (Index i = 0; i < inCoef.size(); ++i) {
        stdErr(i) = std::sqrt(diagonal_of_heissian(i));
        waldZStats(i) = inCoef(i) / stdErr(i);
        waldPValues(i) = 2. * prob::cdf( prob::normal(),
            -std::abs(waldZStats(i)));
        oddsRatios(i) = std::exp( inCoef(i) );
    }

    // Return all coefficients, standard errors, etc. in a tuple
    AnyType tuple;
    tuple << ref_category << inCoef << logLikelihood << stdErr << waldZStats << waldPValues
        << oddsRatios << conditionNo;
    return tuple;
}

// ---------------------------------------------------------------------------
//             Marginal Effects Multi-Logistic Regression States
// ---------------------------------------------------------------------------
/**
 * @brief State for marginal effects calculation for logistic regression
 *
 * TransitionState encapsualtes the transition state during the
 * marginal effects calculation for the logistic-regression aggregate function.
 * To the database, the state is exposed as a single DOUBLE PRECISION array,
 * to the C++ code it is a proper object containing scalars and vectors.
 *
 * Note: We assume that the DOUBLE PRECISION array is initialized by the
 * database with length at least 5, and all elemenets are 0.
 *
 */

template <class Handle>
class mlogregrMarginalTransitionState {

    // By §14.5.3/9: "Friend declarations shall not declare partial
    // specializations." We do access protected members in operator+=().
    template <class OtherHandle>
    friend class mlogregrMarginalTransitionState;

public:
    mlogregrMarginalTransitionState(const AnyType &inArray)
        : mStorage(inArray.getAs<Handle>()) {
        rebind(static_cast<uint16_t>(mStorage[0]),static_cast<uint16_t>(mStorage[1]));
    }

    /**
     * @brief Convert to backend representation
     *
     * We define this function so that we can use State in the
     * argument list and as a return type.
     */
    inline operator AnyType() const {
        return mStorage;
    }

    /**
     * @brief Initialize the iteratively-reweighted-least-squares state.
     *
     * This function is only called for the first iteration, for the first row.
     */
    inline void initialize(const Allocator &inAllocator,
        uint16_t inWidthOfX, uint16_t inNumCategories, uint16_t inRefCategory) {
        mStorage = inAllocator.allocateArray<double, dbal::AggregateContext,
            dbal::DoZero, dbal::ThrowBadAlloc>(arraySize(inWidthOfX, inNumCategories));
        rebind(inWidthOfX, inNumCategories);
        widthOfX = inWidthOfX;
        numCategories = inNumCategories;
        ref_category = inRefCategory;
    }

    /**
     * @brief We need to support assigning the previous state
     */
    template <class OtherHandle> mlogregrMarginalTransitionState &operator=(
        const mlogregrMarginalTransitionState<OtherHandle> &inOtherState) {
        for (size_t i = 0; i < mStorage.size(); i++)
            mStorage[i] = inOtherState.mStorage[i];
        return *this;
    }

    /**
     * @brief Merge with another State object by copying the intra-iteration
     *     fields
     */
    template <class OtherHandle> mlogregrMarginalTransitionState &operator+=(
        const mlogregrMarginalTransitionState<OtherHandle> &inOtherState) {
        if (mStorage.size() != inOtherState.mStorage.size() ||
            widthOfX != inOtherState.widthOfX)
            throw std::logic_error("Internal error: Incompatible transition "
                "states");

        numRows += inOtherState.numRows;
        margins_matrix += inOtherState.margins_matrix;
        X_bar += inOtherState.X_bar;
        X_transp_AX += inOtherState.X_transp_AX;
        reference_margins+= inOtherState.reference_margins;
        return *this;
    }

    /**
     * @brief Reset the inter-iteration fields.
     */
    inline void reset() {
        numRows = 0;
        margins_matrix.fill(0);
        X_bar.fill(0);
        X_transp_AX.fill(0);
        reference_margins.fill(0);
    }

private:
    static inline uint32_t arraySize(const uint16_t inWidthOfX,
        const uint16_t inNumCategories) {
        return 4 + 3*inWidthOfX * inNumCategories + 1*inWidthOfX +
           inNumCategories * inWidthOfX * inWidthOfX * inNumCategories;
    }

    /**
     * @brief Rebind to a new storage array
     *
     * @param inWidthOfX             The number of independent variables.
     * @param inNumCategories The number of categories of the dependant var

     * Array layout (iteration refers to one aggregate-function call):
     * Inter-iteration components (updated in final function):
     * - 0: widthOfX (number of independant variables)
     * - 1: numCategories (number of categories)
     * - 2: ref_category
     * - 3: coef (vector of coefficients)
     *
     * Intra-iteration components (updated in transition step):
     * - 3 + widthOfX*numCategories: numRows (number of rows already processed in this iteration)
     * - 4 + widthOfX * inNumCategories: margins_matrix
     */
    void rebind(uint16_t inWidthOfX = 0, uint16_t inNumCategories = 0) {
        widthOfX.rebind(&mStorage[0]);
        numCategories.rebind(&mStorage[1]);
        ref_category.rebind(&mStorage[2]);
        coef.rebind(&mStorage[3], inWidthOfX * inNumCategories);
        numRows.rebind(&mStorage[3 + inWidthOfX * inNumCategories]);
        margins_matrix.rebind(&mStorage[4 + inWidthOfX * inNumCategories],
            inNumCategories , inWidthOfX );
        X_bar.rebind(&mStorage[4 + 2*inWidthOfX * inNumCategories], inWidthOfX);
        reference_margins.rebind(&mStorage[4 + 2*inWidthOfX + 2*inWidthOfX*inNumCategories],
            inWidthOfX);
        X_transp_AX.rebind(&mStorage[4 + 3*inWidthOfX * inNumCategories + 1*inWidthOfX],
            inNumCategories*inWidthOfX, inWidthOfX*inNumCategories);
    }

    Handle mStorage;

public:
    typename HandleTraits<Handle>::ReferenceToUInt16 widthOfX;
    typename HandleTraits<Handle>::ReferenceToUInt16 numCategories;

    typename HandleTraits<Handle>::ColumnVectorTransparentHandleMap coef;
    typename HandleTraits<Handle>::ReferenceToUInt64 numRows;
    typename HandleTraits<Handle>::MatrixTransparentHandleMap margins_matrix;
    typename HandleTraits<Handle>::MatrixTransparentHandleMap X_transp_AX;
    typename HandleTraits<Handle>::ColumnVectorTransparentHandleMap reference_margins;
    typename HandleTraits<Handle>::ReferenceToUInt16 ref_category;
    typename HandleTraits<Handle>::ColumnVectorTransparentHandleMap X_bar;
};

AnyType
mlogregr_marginal_step_transition::run(AnyType &args) {
    using std::endl;

	  mlogregrMarginalTransitionState<MutableArrayHandle<double> > state = args[0];
    // Get x as a vector of double
    MappedColumnVector x = args[4].getAs<MappedColumnVector>();
    // Get the category & numCategories as integer
    int16_t category = args[1].getAs<int>();
    // Number of categories after pivoting (We pivot around the first category)
    int16_t numCategories = (args[2].getAs<int>() - 1);
    int32_t ref_category = args[3].getAs<int32_t>();
	  MappedColumnVector coefVec = args[5].getAs<MappedColumnVector>();

    // The following check was added with MADLIB-138.
    if (!x.is_finite())
            throw std::domain_error("Design matrix is not finite.");

    if (state.numRows == 0) {
        if (x.size() > std::numeric_limits<uint16_t>::max())
                throw std::domain_error("Number of independent variables cannot be "
                        "larger than 65535.");
        if (numCategories < 1)
                throw std::domain_error("Number of cateogires must be at least 2");
        if (category > numCategories)
                throw std::domain_error("You have entered a category > numCategories"
                    "Categories must be of values {0,1... numCategories-1}");

        // Init the state (requires x.size() and category.size())
        state.initialize(*this,
            static_cast<uint16_t>(x.size()) ,
            static_cast<uint16_t>(numCategories),
            static_cast<uint16_t>(ref_category));

        state.coef = coefVec;
        state.numCategories = numCategories;
        state.ref_category = ref_category;
    }

    // Now do the transition step
    state.numRows++;
    /*
        Get y: Convert to 1/0 boolean vector
        Example: Category 4 : 0 0 0 1 0 0
        Storing it in this forms helps us get a nice closed form expression
    */

    //To pivot around the specified reference category
    ColumnVector y(numCategories);
    y.fill(0);
    if (category > ref_category) {
        y(category - 1) = 1;
    } else if (category < ref_category) {
        y(category) = 1;
    }

    //    Marginal Effect calculations
    // ----------------------------------------------------------------------
    Matrix coef = state.coef;
    coef.resize(numCategories, state.widthOfX);

    // prob is vector of size # categories
    /*
        Note: The above 2 lines could have been written as:
        ColumnVector prob = -coef*x; but this creates warnings.
        See multilog for details.
    */
    ColumnVector prob(numCategories);
    prob = coef*x;

    // Calculate the odds ratio
    prob = prob.array().exp();
    double prob_sum = prob.sum();

    prob = prob / (1 + prob_sum);

    // Reference category computations. They have been taken out of
    // the output but left in the infrastructure
    double ref_prob = 1 / (1 + prob_sum);

    Matrix probDiag = prob.asDiagonal();

    // Marginal effects (reference calculated separately)
    ColumnVector coef_trans_prob;
    coef_trans_prob = coef.transpose() * prob;
    Matrix margins_matrix = coef;
    margins_matrix.rowwise() -= coef_trans_prob.transpose();
    margins_matrix = probDiag * margins_matrix;

    //    Variance Calculations
    // ----------------------------------------------------------------------
    /*
         a is a matrix of size JxJ where J is the number of categories
         a_j1j2 = -pi(j1)*(1-pi(j2))if j1 == j2
         a_j1j2 =  pi(j1)*pi(j2) if j1 != j2
    */
    Matrix a(numCategories,numCategories);
    a = prob * prob.transpose() - probDiag;

    // Start the Hessian calculations
    Matrix X_transp_AX(numCategories * state.widthOfX, numCategories * state.widthOfX);

    /*
        Again: The following 3 lines could have been written as
        Matrix XXTrans = x * x.transpose();
        but it creates warnings related to the type of x. Here is an easy fix
    */
    Matrix cv_x = x;
    Matrix XXTrans = trans(cv_x);
    XXTrans = cv_x * XXTrans;

    //Eigen doesn't supported outer-products for matrices, so we have to do our own.
    //This operation is also known as a tensor-product.
    for (int i1 = 0; i1 < state.widthOfX; i1++){
         for (int i2 = 0; i2 <state.widthOfX; i2++){
            int rowOffset = numCategories * i1;
            int colOffset = numCategories * i2;

            X_transp_AX.block(rowOffset, colOffset, numCategories,  numCategories) = XXTrans(i1, i2) * a;
        }
    }

    triangularView<Lower>(state.X_transp_AX) += X_transp_AX;
    state.margins_matrix += margins_matrix;
    state.reference_margins += -coef_trans_prob * ref_prob;
    state.X_bar += x; // It is called X_bar but it really is the sum

    return state;

}

/**
 * @brief Perform the perliminary aggregation function: Merge transition states
     This merge step is the same as that of logistic regression.
 */
AnyType
mlogregr_marginal_step_merge_states::run(AnyType &args) {
    mlogregrMarginalTransitionState<MutableArrayHandle<double> > stateLeft = args[0];
    mlogregrMarginalTransitionState<ArrayHandle<double> > stateRight = args[1];

    // We first handle the trivial case where this function is called with one
    // of the states being the initial state
    if (stateLeft.numRows == 0)
        return stateRight;
    else if (stateRight.numRows == 0)
        return stateLeft;

    // Merge states together and return
    stateLeft += stateRight;
    return stateLeft;
}

AnyType mlogregr_marginalstateToResult(
    const Allocator &inAllocator,
    const int numRows,
    const ColumnVector &inCoef,
    const ColumnVector &inMargins,
    const ColumnVector &inVariance
    ) {

    MutableNativeColumnVector margins(
        inAllocator.allocateArray<double>(inMargins.size()));
    MutableNativeColumnVector coef(
        inAllocator.allocateArray<double>(inMargins.size()));
    MutableNativeColumnVector stdErr(
        inAllocator.allocateArray<double>(inMargins.size()));
    MutableNativeColumnVector tStats(
        inAllocator.allocateArray<double>(inMargins.size()));
    MutableNativeColumnVector pValues(
        inAllocator.allocateArray<double>(inMargins.size()));

    for (Index i = 0; i < inMargins.size(); ++i) {
        margins(i) = inMargins(i);
        coef(i) = inCoef(i);
        stdErr(i) = std::sqrt(inVariance(i));
        tStats(i) = margins(i) / stdErr(i);

        // P-values only make sense if numRows > coef.size()
        if (numRows > inCoef.size())
          pValues(i) = 2. * prob::cdf(
              boost::math::complement(
                  prob::students_t(
                      static_cast<double>(numRows - inCoef.size())
                  ),
                  std::fabs(tStats(i))
              ));
    }

    // Return all coefficients, standard errors, etc. in a tuple
    // Note: PValues will return NULL if numRows <= coef.size
    AnyType tuple;
    tuple << margins
          << coef
          << stdErr
          << tStats
        	<< (numRows > inCoef.size()? pValues: Null());
    return tuple;
}

/**
 * @brief Perform the logistic-regression final step
 */
AnyType
mlogregr_marginal_step_final::run(AnyType &args) {
    // We request a mutable object. Depending on the backend, this might perform
    // a deep copy.
    mlogregrMarginalTransitionState<ArrayHandle<double> > state = args[0];
    // Aggregates that haven't seen any data just return Null.
    if (state.numRows == 0)
        return Null();

    if(!state.coef.is_finite())
        throw NoSolutionFoundException("Over- or underflow in Newton step, "
            "while updating coefficients. Input data is likely of poor "
            "numerical condition.");

    // See MADLIB-138. At least on certain platforms and with certain versions,
    // LAPACK will run into an infinite loop if pinv() is called for non-finite
    // matrices. We extend the check also to the dependent variables.
    if (!state.X_transp_AX.is_finite())
        throw NoSolutionFoundException("Over- or underflow in intermediate "
            "calulation. Input data is likely of poor numerical condition.");

    // Include marginal effects of reference variable:
    // FIXME: They have been taken out of the output for now
    //const int size = state.coef.size() + numIndepVars;
    const int size = state.coef.size();

    // Variance-covariance calculation
    // ----------------------------------------------------------
    SymmetricPositiveDefiniteEigenDecomposition<Matrix> decomposition(
        -1 * state.X_transp_AX, EigenvaluesOnly, ComputePseudoInverse);

    // Precompute -(X^T * A * X)^-1
    Matrix V = decomposition.pseudoInverse();

    // Marginal Effect calculation
    // ---------------------------------------------------------
    Matrix margins_matrix;
    margins_matrix = state.margins_matrix / state.numRows;

    // Standard error calculation
    // ----------------------------------------------------------
    ColumnVector x_bar = state.X_bar / state.numRows;
    int numIndepVars = state.coef.size() / state.numCategories;
    int numCategories = state.numCategories;

    Matrix coef = state.coef;
    coef.resize(state.numCategories, state.widthOfX);

    // Variance & Marginal gradient
    ColumnVector marginal_gradient(size);
    ColumnVector variance(size);
    variance.setOnes();

    // Probibility vector at the mean
    ColumnVector p_bar(state.numCategories);
    ColumnVector coef_x_bar(state.numCategories);
    ColumnVector coef_trans_p_bar(state.widthOfX);

    coef_x_bar = coef*x_bar;
    coef_trans_p_bar = coef.transpose() * p_bar;

    p_bar = coef_x_bar;
    p_bar = p_bar.array().exp();
    p_bar = p_bar / (1 + p_bar.sum());

    // Marginal effects at the mean
    Matrix margins_mean_matrix(numCategories, numIndepVars);
    margins_mean_matrix.rowwise() -= coef_trans_p_bar.transpose();
    margins_mean_matrix = p_bar.asDiagonal() * margins_matrix;

    // Compute the variance for each marginal effect
    int index, e_j_J, e_k_K, e_k_K_j_J;
    for (int K=0; K < numIndepVars; K++){
      for (int J=0; J < numCategories; J++){

        for (int k=0; k < numIndepVars; k++){
          for (int j=0; j < numCategories; j++){

            e_j_J = (j==J) ? 1: 0;
            e_k_K = (k==K) ? 1: 0;
            e_k_K_j_J = (j==J && k==K) ? 1: 0;

            index = k*numCategories + j;

            marginal_gradient(index) =
                x_bar(k) * (e_j_J  - p_bar(j)) * margins_mean_matrix(J,K);

            marginal_gradient(index) += p_bar(J) *
                ( e_k_K_j_J - p_bar(j) * e_k_K - x_bar(k) * margins_mean_matrix(j,K));
          }
        }

        // NOTE: Since the earlier Variance calculations are being done by
        // stacking up the indepdent variables for each category separtely
        variance(K + numIndepVars * J) =
                          marginal_gradient.transpose() * V * marginal_gradient;

      }
    }

    // Add in reference variables to all the calculations
    // ----------------------------------------------------------
    ColumnVector coef_with_ref(size);
    ColumnVector margins_with_ref(size);

    // Vectorize the margins_matrix and add the reference variable
    for (int j=0; j < numCategories; j++){
      for (int k=0; k < numIndepVars; k++){

        index = k + numIndepVars *j;
        coef_with_ref(index) = coef(j,k);

        index = k + numIndepVars*j;
        margins_with_ref(index) = margins_matrix(j,k);
      }
    }

    return mlogregr_marginalstateToResult(*this,
                                          state.numRows,
                                          coef_with_ref,
                                          margins_with_ref,
                                          variance);
}

// ------------------------ End of Marginal ------------------------------------

} // namespace regress

} // namespace modules

} // namespace madlib


/**
 * @author     : Zhao Chonyyao (cyzhao@zju.edu.cn)
 * @date       : 2021-04-30
 * @description: elasticity finite element method problem
 * @version    : 1.0
 */
#include <memory>
#include <iomanip>
#include <boost/property_tree/ptree.hpp>

#include "Common/DEFINE_TYPE.h"
#include "Common/error.h"

// TODO: possible bad idea of having dependence to model in problem module
#include "Model/fem/elas_energy.h"
#include "Model/fem/mass_matrix.h"

#include "Problem/energy/basic_energy.h"
#include "Io/io.h"
#include "Geometry/extract_surface.imp"

#include "elas_fem_problem.h"

namespace PhysIKA {
using namespace std;
using namespace Eigen;

template <typename T>
int read_elas_mtr(const char* file, VEC<T>& Young, VEC<T>& Poi, const size_t num_cells)
{
    cout << "read_elas_mtr" << endl;
    Young.resize(num_cells);
    Poi.resize(num_cells);
    ifstream ifs(file);
    if (ifs.fail())
    {
        std::cerr << "[info] "
                  << "can not open file" << file << std::endl;
        return __LINE__;
    }
    size_t cell_id = 0;
    while (!ifs.eof())
    {
        if (cell_id == num_cells)
            break;
        ifs >> Young(cell_id) >> Poi(cell_id);
        ++cell_id;
    }
    ifs.close();
    return 0;
}

template <typename T>
elas_problem_builder<T>::elas_problem_builder(const T* x, const boost::property_tree::ptree& pt)
    : pt_(pt)
{
    if (pt.get<string>("solver_type") == "explicit")
    {
#define SEMI_IMPLICIT
    }

    //TODO: need to check exception
    const string filename = pt.get<string>("filename");
    MAT<T>       nods(1, 1);
    MatrixXi     cells(1, 1);

    const string type = pt.get<string>("type", "tet");

    if (type == "tet")
    {
        IF_ERR(exit, mesh_read_from_vtk<T, 4>(filename.c_str(), nods, cells));
    }
    else if (type == "hex")
        exit_if(mesh_read_from_vtk<T, 8>(filename.c_str(), nods, cells));
    else
    {
        // error_msg("type:<%s> is not supported.", type.c_str());
    }

    const size_t num_nods = nods.cols(), num_cells = cells.cols();
    cout << "V" << nods.rows() << " " << nods.cols() << endl
         << "T " << cells.rows() << " " << cells.cols() << endl;
    if (x != nullptr)
        nods = Map<const MAT<T>>(x, nods.rows(), nods.cols());
    cout << "Boundary Box :\n"
         << nods.rowwise().minCoeff() << endl
         << nods.rowwise().maxCoeff() << endl;

    REST_  = nods;
    cells_ = cells;

    Matrix<T, 3, 3> rot;
    rot << 0, 0, 1, 0, 1, 0, -1, 0, 0;
    cout << "rotatoin matrix is " << rot << endl;
    MAT<T> rotated_nods = rot * nods;

    // const string outdir = argv[3];
    auto phy_paras = pt.get_child("physics");
    //set mtr
    const T      rho       = phy_paras.get<T>("rho", 20);
    const T      Young     = phy_paras.get<T>("Young", 2000.0);
    const T      poi       = phy_paras.get<T>("poi", 0.3);
    const T      gravity   = phy_paras.get<T>("gravity", 9.8);
    const T      dt        = phy_paras.get<T>("dt", 0.01);
    const T      w_pos     = phy_paras.get<T>("w_pos", 1e6);
    const size_t num_frame = phy_paras.get<size_t>("num_frames", 100);

    //set mtr
    VEC<T>       Young_vec(num_cells), Poi_vec(num_cells);
    const string mtr_file = pt.get<string>("mtr_file", "");
    if (mtr_file != "")
        read_elas_mtr(mtr_file.c_str(), Young_vec, Poi_vec, num_cells);

    //read fixed points
    vector<size_t> cons(0);
    const string   cons_file_path = pt.get<string>("cons", "");
    if (cons_file_path != "")
        IF_ERR(exit, read_fixed_verts_from_csv(cons_file_path.c_str(), cons));
    cout << "constrint " << cons.size() << " points" << endl;

    //calc mass vector
    Matrix<T, -1, 1> mass_vec(num_nods);
    // calc_mass_vector<T>(nods, cells, rho, mass_vec);
    if (type == "tet")
        mass_calculator<T, 3, 4, 1, 1, basis_func, quadrature>(nods, cells, rho, mass_vec);
    else if (type == "hex")
        mass_calculator<T, 3, 8, 1, 2, basis_func, quadrature>(nods, cells, rho, mass_vec);

    cout << "build energy" << endl;
    int ELAS = 0;
    int GRAV = 1;
    int KIN  = 2;
    int POS  = 3;
    if (pt_.get<string>("solver_type") == "explicit")
        POS = 2;

    ebf_.resize(POS + 1);
    {
        const string csttt_type = phy_paras.get<string>("csttt", "linear");
        if (pt.get<bool>("rotate", false))
            nods = rotated_nods;
        if (mtr_file != "")
            gen_elas_energy_intf<T>(type, csttt_type, nods, cells, Young_vec, Poi_vec, ebf_[ELAS], &elas_intf_);
        else
            gen_elas_energy_intf<T>(type, csttt_type, nods, cells, Young, poi, ebf_[ELAS], &elas_intf_);
        nods = REST_;
        // to lowercase.
        char axis = pt.get<char>("grav_axis", 'y') | 0x20;
        if (axis > 'z' || axis < 'x')
        {
            // error_msg("grav_axis should be one of x(X), y(Y) or z(Z).");
        }
        ebf_[GRAV] = make_shared<gravity_energy<T, 3>>(num_nods, 1, gravity, mass_vec, axis);

        kinetic_ = pt.get<bool>("dynamics", true) ? make_shared<momentum<T, 3>>(nods.data(), num_nods, mass_vec, dt)
                                                  : nullptr;

        if (pt_.get<string>("solver_type") == "implicit")
            ebf_[KIN] = kinetic_;

        ebf_[POS] = make_shared<position_constraint<T, 3>>(nods.data(), num_nods, w_pos, cons);
    }

    //set constraint

    enum constraint_type
    {
        COLL
    };
    cbf_.resize(COLL + 1);
    collider_  = nullptr;
    cbf_[COLL] = collider_;

    if (pt_.get<string>("solver_type") == "explicit")
    {
        Map<Matrix<T, -1, 1>> position(REST_.data(), REST_.size());
        semi_implicit_ = make_shared<semi_implicit<T>>(dt, mass_vec, position);
    }
}

template <typename T>
std::shared_ptr<Problem<T, 3>> elas_problem_builder<T>::build_problem() const
{
    cout << "assemble energy" << endl;
    shared_ptr<Functional<T, 3>> energy;
    try
    {
        energy = build_energy_t<T, 3>(ebf_);
    }
    catch (std::exception& e)
    {
        cerr << e.what() << endl;
        exit(EXIT_FAILURE);
    }

    shared_ptr<Constraint<T>> constraint;
    cout << "assemble constraint" << endl;
    bool all_null = true;
    for (auto& c : cbf_)
        if (c != nullptr)
            all_null = false;
    if (all_null)
    {
        constraint = nullptr;
        cout << "WARNGING: No hard constraints." << endl;
    }
    else
    {
        try
        {
            constraint = build_constraint_t<T>(cbf_);
        }
        catch (std::exception& e)
        {
            cerr << e.what() << endl;
            exit(EXIT_FAILURE);
        }
    }
    exit_if(constraint != nullptr && energy->Nx() != constraint->Nx(), "energy and constraint has different dimension.");
    return make_shared<Problem<T, 3>>(energy, constraint);
}

template <typename T>
int elas_problem_builder<T>::update_problem(const T* x, const T* v)
{
    if (kinetic_ != nullptr)
        IF_ERR(return, kinetic_->update_location_and_velocity(x, v));
    if (collider_ != nullptr)
        IF_ERR(return, collider_->update(x));
    return 0;
}

template class elas_problem_builder<double>;

template class elas_problem_builder<float>;

}  // namespace PhysIKA