/* -*- mia-c++  -*-
 *
 * This file is part of MIA - a toolbox for medical image analysis 
 * Copyright (c) Leipzig, Madrid 1999-2017 Gert Wollny
 *
 * MIA is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with MIA; if not, see <http://www.gnu.org/licenses/>.
 *
 */

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include <gsl/gsl_cblas.h>
#include <iomanip>
#include <limits>

#include <boost/cast.hpp>

#include <mia/3d/fifof/mlv.hh>

NS_BEGIN(mlv_2dstack_filter)

NS_MIA_USE
using namespace std;
using namespace boost;


template <class InputIterator, class OutputIterator>
static void convert(InputIterator a, InputIterator e, OutputIterator o)
{
	while (a != e) {
		*o = numeric_cast<typename iterator_traits<OutputIterator>::value_type>(*a);
		++o;
		++a;
	}
}

C2DMLVnFifoFilter::C2DMLVnFifoFilter(size_t hwidth):
	C2DImageFifoFilter(4 * hwidth + 1, 2 * hwidth + 1, 2 * hwidth),
	m_hw(hwidth),
	m_w(2 * hwidth + 1), 
	m_buf_slice_size(0)
{
}

// This operator reads out the result

template <typename T>
C2DImage *C2DMLVnFifoFilter::operator()(const T3DImage<T>& /*dummy*/) const
{
	TRACE("C2DMLVnFifoFilter::operator() (pull)");
	
	static int slice = 0; 
	
	size_t start = 2 * m_w - 2;
	size_t end = min(get_end() + m_w, get_buffer_size()); 

	cvdebug() << "Evaluate "<< slice++ << " output " << start << " - " << end << "\n"; 
	C2DFImage mu_result(m_slice_size);
	C2DFImage sigma_buffer(m_slice_size);

	fill(sigma_buffer.begin(), sigma_buffer.end(), numeric_limits<float>::max());

	for (size_t z = start; z <  end; ++z) {
		for (size_t y = 0; y < m_slice_size.y; ++y) {
			for (size_t iy = 0; iy < m_w; ++iy)
				for (size_t ix = 0; ix < m_w; ++ix) {
					
					C3DFImage::const_iterator mu_i = m_mu[z].begin_at(ix, y + iy);
					C3DFImage::const_iterator si_i = m_sigma[z].begin_at(ix, y + iy);
					C2DFImage::iterator omu_i = mu_result.begin_at(0, y);
					C2DFImage::iterator osi_i = sigma_buffer.begin_at(0, y);

					for (size_t x = 0; x < m_slice_size.x; ++x) {
						if (osi_i[x] > si_i[x]) {
							osi_i[x] = si_i[x];
							omu_i[x] = mu_i[x];
						}
					}
				}

			}

	}
	T2DImage<T> *result = new T2DImage<T>(m_slice_size);
	convert(mu_result.begin(), mu_result.end(), result->begin());

	return result;
}

static void do_evaluate(C2DFImage::const_iterator ni, C2DFImage::const_iterator ne,
		     C2DFImage::iterator mu, C2DFImage::iterator sigma)
{
	while (ni != ne) {
		const float n = *ni;
		const float muq = *mu / *ni;

		*sigma = (n > 1.0) ? (*sigma - muq * *mu) / (n - 1.0f) : 0.0f;

		*mu++ = muq;
		++sigma; 
		++ni;
	}
}

template <typename Iterator>
void rotate(Iterator begin, Iterator end)
{
	typename Iterator::value_type help = *(end - 1);
	copy_backward(begin, end - 1, end);
	*begin = help;

	fill(begin->begin(), begin->end(), 0.0); 
}

void C2DMLVnFifoFilter::shift_buffer()
{
	TRACE_FUNCTION; 
	cvdebug() << "Evaluate shift\n"; 
	rotate(m_mu.begin(), m_mu.end());
	rotate(m_sigma.begin(),m_sigma.end());
	rotate(m_n.begin(), m_n.end());
	copy_backward(m_evaluated.begin(), m_evaluated.end() - 1, m_evaluated.end());
	m_evaluated[0] = false; 
}

void C2DMLVnFifoFilter::evaluate(size_t slice)
{
	TRACE_FUNCTION; 
	if (m_evaluated[slice]) 
		return; 

	cvdebug() << "Evaluate mu(sigma) for  " << slice << "\n"; 
	do_evaluate(m_n[slice].begin(),
		    m_n[slice].end(),
		    m_mu[slice].begin(),
		    m_sigma[slice].begin());
	m_evaluated[slice] = true; 
}

// this operator add the new slice

template <typename In, typename Out>
bool copy_was_zero(In b, In e, Out o)
{
	bool result = true; 
	while (b != e) {
		result &= (*b == 0); 
		*o = *b; 
		++o; 
		++b; 
	}
	return result; 
}


template <typename T>
C2DImage *C2DMLVnFifoFilter::operator()(const T2DImage<T>& src)
{
	TRACE("C2DMLVnFifoFilter::operator()(const T2DImage<T>& input)(push)");

	cvdebug() << "Evaluate input\n"; 

	
	C2DBounds temp_size = m_mu[0].get_size(); 

	vector<float> val(src.get_size().x); 
	vector<float> val2(src.get_size().x);
			
	vector<float> sum_mu_l1(temp_size.x); 
	vector<float> sum_sigma_l1(temp_size.x);

	vector<float> sum_mu_l2(m_mu[0].size(), 0.0); 
	vector<float> sum_sigma_l2(m_mu[0].size(), 0.0);

	
	for (size_t y = 0; y < src.get_size().y; ++y){
		
		// copy one input row and check whether it is all zero, if yes, shortcut
		if (copy_was_zero(src.begin_at(0,y), src.begin_at(0,y) + src.get_size().x, val.begin()))
			continue; 

		fill(sum_mu_l1.begin(), sum_mu_l1.end(), 0.0); 
		fill(sum_sigma_l1.begin(), sum_sigma_l1.end(), 0.0); 
		
		transform(val.begin(), val.end(), val2.begin(),[](float x) { return  x * x;});
		
		for(size_t x = 0; x < m_w; ++x) {
			cblas_saxpy(src.get_size().x, 1.0f, &val[0],  1, &sum_mu_l1[x], 1); 
			cblas_saxpy(src.get_size().x, 1.0f, &val2[0], 1, &sum_sigma_l1[x], 1);
		}
		
		for(size_t iy = 0; iy < m_w; ++iy) {
			cblas_saxpy(sum_mu_l1.size(), 1.0f, &sum_mu_l1[0],  1, 
				    &sum_mu_l2[(y + iy) * temp_size.x], 1); 
			
			cblas_saxpy(sum_sigma_l1.size(), 1.0f, &sum_sigma_l1[0],  1, 
				    &sum_sigma_l2[(y + iy) * temp_size.x], 1);
		}
	}

	// update the numbers
	for (size_t z = 0; z < m_w; ++z) {
		cblas_saxpy(sum_mu_l2.size(), 1.0f, &sum_mu_l2[0],  1, &m_mu[z](0,0) , 1); 
		cblas_saxpy(sum_sigma_l2.size(), 1.0f, &sum_sigma_l2[0],  1,&m_sigma[z](0,0) , 1);
		cblas_saxpy(m_n_template.size(), 1.0f, &m_n_template(0,0), 1, &m_n[z](0,0), 1); 
	}
	// the next shift will push this slice out if the input range 
	// ence we can evaluate the real values here. 
	evaluate(m_w-1); 
	

	return NULL;
}

void C2DMLVnFifoFilter::post_finalize()
{
	TRACE_FUNCTION; 
	m_mu.resize(0);
	m_sigma.resize(0);
	m_n.resize(0);
}

void C2DMLVnFifoFilter::do_initialize(::boost::call_traits<P2DImage>::param_type x)
{
	TRACE("C2DMLVnFifoFilter::do_initialize");

	m_slice_size = x->get_size();
	cvdebug() << "Input slices are of size " << m_slice_size << "\n"; 
	
	C2DBounds size(m_slice_size.x + 2 * m_hw,
		       m_slice_size.y + 2 * m_hw);

	size_t n_slices = get_buffer_size();

	m_mu = vector<C2DFImage>(n_slices);
	m_sigma = vector<C2DFImage>(n_slices);
	m_n = vector<C2DFImage>(n_slices);
	m_evaluated = vector<bool>(n_slices, false);
		
	for (size_t i = 0; i < n_slices; ++i) {
		m_mu[i] = C2DFImage(size);
		m_sigma[i] = C2DFImage(size);
		m_n[i] = C2DFImage(size);
	}

	m_buf_slice_size = size.x * size.y;

	m_buf1.resize(m_slice_size.x);
	m_buf2.resize(m_slice_size.x);
	m_prototype.reset(create_buffer(C2DBounds(1,1), 1, x->get_pixel_type()));


	m_n_template = C2DFImage(C2DBounds(m_slice_size.x + 2 * m_hw, m_slice_size.y + 2 * m_hw));
	// create the pattern that is used to update the count
	// this is the slow version, but we have to do it only once
	for (size_t y = 0; y < m_slice_size.y; ++y)
		for (size_t x = 0; x < m_slice_size.x; ++x)
			for (size_t iy = 0; iy < 2 * m_hw + 1; ++iy)
				for (size_t ix = 0; ix < 2 * m_hw + 1; ++ix)
					m_n_template(x + ix, y + iy) += 1.0;
}

void C2DMLVnFifoFilter::do_push(::boost::call_traits<P2DImage>::param_type x)
{
	TRACE("C2DMLVnFifoFilter::do_push");
	mia::accumulate(*this, *x);
}

P2DImage C2DMLVnFifoFilter::do_filter()
{
	TRACE("C2DMLVnFifoFilter::do_filter");
	return P2DImage(mia::filter(*this, *m_prototype));
}


class C2DMLVnFifoFilterPlugin : public C2DFifoFilterPlugin {
public:
	C2DMLVnFifoFilterPlugin();
private:

	virtual const string do_get_descr() const;
	virtual C2DImageFifoFilter *do_create()const;

	mutable int m_hw;
};

C2DMLVnFifoFilterPlugin::C2DMLVnFifoFilterPlugin():
	C2DFifoFilterPlugin("mlv"),
	m_hw(1)
{
	add_parameter("w", make_lc_param(m_hw, 1, false, "filter width parameter"));
}

const string C2DMLVnFifoFilterPlugin::do_get_descr() const
{
	return "Mean of Least Variance 2D stack image filter";
}

C2DImageFifoFilter *C2DMLVnFifoFilterPlugin::do_create()const
{
	return new C2DMLVnFifoFilter(m_hw);
}

extern "C" EXPORT CPluginBase *get_plugin_interface()
{

	return new C2DMLVnFifoFilterPlugin();
}

NS_END

