#include "tt.h"

#include <vector>
#include "boost/array.hpp"
#include "boost/multi_array.hpp"

struct SortPredicate
{
	bool operator() (const float& start, const float& end) {
		return start < end;
	}
};

bool tt (boost::multi_array< float, 2> &s_in, float& v_out, std::vector<float>& v_dist_out)
{
	// Jake -- note that I assume the input does not contain all 0's, in which Claudia's
	// original code returns NaN's.

	// Add a row of 0's
	int size1, size2;
	size1 = s_in.shape()[0];
	size2 = s_in.shape()[1];
	boost::multi_array< float, 2>::extent_gen extents;
	s_in.resize(extents[size1+1][size2]);
	for (int i = 0; i < size2; i++) {
		s_in[size1][i] = 0.0f;
	}

	// Convert matrix into vector; remember column-major order!
	int lengthX = (size1+1) * size2;
//	float *x = new float[lengthX];
    std::vector<float> x(lengthX, 0.0);
	for (int i = 0; i < size1 + 1; i++) {
		for (int j = 0; j < size2; j++) {
			x[j * (size1 + 1) + i] = s_in[i][j];
		}
	}

	// Compute z1 and z0
	std::vector<float> z0, z1;
	for (int i = 0; i < lengthX - 1; i++) {
		float diff = x[i+1] - x[i];

		if (diff == 1) {
			z0.push_back(i);
		} else if (diff == -1) {
			z1.push_back(i);
		}
	}

//	assert(z0.size() > 0 && z1.size() > 0);
    if( (z0.size() <= 0) || (z1.size() <= 0) )
        return false;

	// Measure the length of black lines
	if (z0[0] > z1[0]) {
		// Note - have to change "0" in Matlab code to "-1" in C code due to C-style indexing!
		z0.insert(z0.begin(), -1);
		if (z0.size() > z1.size()) {
			// Remove last element
			z0.pop_back();
		}
	}

//	assert(z1.size() == z0.size());
    if( z1.size() != z0.size() )
        return false;

	// Compute t
	std::vector<float> t;
	for ( unsigned int i = 0; i < z1.size(); i++) {
		t.push_back(z1[i] - z0[i]);
	}
	// Sort t
	std::sort(t.begin(), t.end(), SortPredicate());

	// Compute t1	
	std::vector<float> t1;
	for ( unsigned int i = 0; i < t.size(); i++) {
		if (t[i] - 1 != 0) {
			t1.push_back(t[i]);
		}
	}

	// Compute return values
	if (t1.size() == 0) {
		v_out = 0;
	} else {
		v_out = 0;
		// Compute mean(t1)
		for ( unsigned int i = 0; i < t1.size(); i++) {
			v_out += t1[i];
		}
		v_out /= t1.size();
	}

	// v_dist_out := t
	for ( unsigned int i = 0; i < t.size(); i++) {
		v_dist_out.push_back(t[i]);
	}

//	delete [] x;

    return true;
}