/* -*- 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/>.
 *
 */

#include <algorithm>
#include <functional>
#include <cmath>
#include <boost/type_traits.hpp>

#include <mia/2d/filter.hh>
#include <libmia/histogram.hh>
#include <libmia/monaAlgorithms.hh>

namespace ianiso_2dimage_filter {
NS_MIA_USE;
using namespace std;

// scale factor for MAD to zero-mean normal distribution  
static const float zmn_weight = 1.4826; 


static char const *plugin_name = "ianiso";
static const CIntOption param_iterations("iter", "number of iterations", 100, 1, 10000); 
static const CFloatOption param_epsilon("epsilon", "delta value to finish iteration", 1, 0.001, 100); 

inline  float psi(float x, float sigma) {
	if (x > sigma || -x > sigma)
		return 0.0; 
	float val = x / sigma; 
	val *= val; 
	val = 1.0 - val; 
	return x * val * val; 
}



class C2DAnisoDiff: public TUnaryImageFilter<int> {
public:
	C2DAnisoDiff(int maxiter, float epsilon); 

	template <class Data2D>
	typename C2DAnisoDiff::result_type operator () (Data2D& data) const ;

	void work(C2DSLImage& work_copy)const; 
private: 
	float estimate_MAD(const C2DSLImage& data, vector<C2DSLImage>& buffer)const; 
	
	float diffuse(C2DSLImage& data, float sigma, float gamma, const vector<C2DSLImage>& buffer)const; 


	int m_maxiter; 
	float m_epsilon; 
};


class CAnisoDiff2DImageFilter: public C2DImageFilterBase {
	C2DAnisoDiff m_filter; 
public:
	CAnisoDiff2DImageFilter(int maxiter, float epsilon);

	virtual void do_filter(P2DImage& image) const;
};

C2DAnisoDiff::C2DAnisoDiff(int maxiter, float epsilon):
	m_maxiter(maxiter),
	m_epsilon(epsilon)
{
}

class CAnisoDiff2DImageFilterFactory: public C2DFilterPlugin {
public: 
	CAnisoDiff2DImageFilterFactory();
	virtual C2DFilter *create(const CParsedOptions& options) const;
	virtual const string do_get_descr()const; 
};

typedef pair<float, float> HScale; 

inline bool operator < (const HScale& a, const HScale& b)
{
	return a.first < b.first; 
}


/* estimate the MAD */
float C2DAnisoDiff::estimate_MAD(const C2DSLImage& data, vector<C2DSLImage>& buffer)const
{
	
	long img_min, img_max; 
	get_min_max(data.begin(), data.end(), img_min, img_max);

	// this is to avoid some warnings 
	float dist = img_max - img_min; 

	cvdebug() << "Histogram spread = " << dist << " in ["<< (float)img_min <<", "<< (float)img_max  <<"]\n"; 

	vector<unsigned long> histogram(img_max); 

	cvdebug() << "evaluate gradient\n"; 
	C2DSLImage::const_iterator id = data.begin(); 
	advance(id, data.get_size().x); 
	// evaluate the gradient values 
	unsigned int nelements = 0; 

	for (size_t y = 1; y < data.get_size().y - 1; ++y) {
		++id; 

		for (size_t x = 1; x < data.get_size().x - 1; ++x, ++id) {
			if (*id > 0) {
				long idd = *id; 
				
				long d1 = id[-data.get_size().x]  - idd; 
				buffer[0](x,y) = d1; 
				
				long d2 = id[-1]  - idd; 
				buffer[1](x,y) = d2; 
				
				long d3 = id[+1] - idd;
				buffer[2](x,y) = d3; 
				
				long d4 = id[data.get_size().x]  - idd; 
				buffer[3](x,y) = d4; 
			
				++histogram[ d1 > 0 ? d1 : -d1 ];
				++histogram[ d2 > 0 ? d2 : -d2 ];	
				++histogram[ d3 > 0 ? d3 : -d3 ];
				++histogram[ d4 > 0 ? d4 : -d4 ];
				nelements += 4; 
			}
		}
		++id; 
	}
	cvdebug() << "used: " <<nelements << "\n"; 
	cvdebug() << "evaluate median\n"; 
	
	// need to know how many elements we are
	float nelemnts_median = nelements / 2.0;
	
	long median = img_min; 
	float thresh = 0; 
	for (vector<unsigned long>::const_iterator i = histogram.begin(); 
	     (i != histogram.end() && thresh < nelemnts_median); ++i, ++median) {
		thresh += *i;
	}
	
	cvdebug() << "median = " << median << "\n"; 

	// evaluate the distance between median and actual value
	vector<HScale> new_hist;
	float val = img_min;
	for (vector<unsigned long>::const_iterator i = histogram.begin(); 
	     i != histogram.end(); ++i, val += 1.0 ) {
		float v = val - median; 
		new_hist.push_back(HScale( v > 0 ? v : -v, *i)); 
	}

	sort(new_hist.begin(), new_hist.end()); 
	
#if 0
	for (vector<HScale>::const_iterator k = new_hist.begin(); k != new_hist.end(); ++k)
		cvdebug() << k->first << " : " << k->second << "\n"; 
#endif

	
	thresh = 0; 
	vector<HScale>::const_iterator melement = new_hist.begin(); 
	vector<HScale>::const_iterator old_element = new_hist.end(); 
	while (melement != new_hist.end() && thresh < nelemnts_median) {
		thresh += melement->second;
		old_element = melement; 
		++melement; 
	}
	
	// weight the enclosing neighbours
	if (old_element != new_hist.end()) {
		const float f1 = (thresh - nelemnts_median) / old_element->second; 
		return (1.0 - f1)  * melement->first + f1 *  old_element->first; 
	}else
		return melement->first; 
}
float C2DAnisoDiff::diffuse(C2DSLImage& data, float sigma, float gamma, const vector<C2DSLImage>& buffer)const
{
	vector<C2DSLImage::const_iterator> ib(4); 
	
	for (int i = 0; i < 4; ++i)
		ib[i] = buffer[i].begin(); 

	float sum = 0.0; 
	for (C2DSLImage::iterator id = data.begin(); id != data.end();  ++id) {
		if (*id != 0) { // skip all the black pixels, they should be f
			float delta = 0.0; 
			for (int i = 0; i < 4; ++i) {
				delta += psi(*(ib[i]), sigma);
			}
			delta *= gamma; 
			
			const long delta2 = (long)(delta + 0.5); 
			sum += delta2 * delta2; 
		
			*id = *id + delta2;
		}
		for (int i = 0; i < 4; ++i) 
			++ib[i]; 
	}
	return sum; 
}

void C2DAnisoDiff::work(C2DSLImage& work_copy)const 
{
	int iter = 0; 

	transform(work_copy.begin(), work_copy.end(), work_copy.begin(), bind2nd(multiplies<long>(), 16)); 

	vector<C2DSLImage> buffer; 
	for (int i = 0; i < 4; ++i)
		buffer.push_back(C2DSLImage(work_copy.get_size())); 
	
	char endline = cverb.show_debug() ? '\n' : '\r'; 

	float delta; 
	do {
		++iter; 
		cvmsg() <<iter <<": "; 
		float sigma_e = zmn_weight *  estimate_MAD(work_copy, buffer); 
		cvmsg() << " sigma_e = " << sigma_e <<" "; 

		if (sigma_e == 0.0) { // image contains only one colour
			break; 
		}
		float sigma = sqrt(5.0) *  sigma_e; 
		

		
		float gamma = 1.0 / (psi(sigma_e, sigma) * 4.0); 
				     
		cvmsg() <<" gamma = "<<gamma <<" "; 
		delta = diffuse(work_copy, sigma, gamma, buffer); 
		
		cvmsg() <<" delta " << delta << "       " << endline; 
		
		
	} while (delta > m_epsilon && iter < m_maxiter); 
	cvmsg() << '\n'; 
	
}

template <typename Data2D, bool supported>
struct __dispatch {
	static void apply(Data2D& image, const C2DAnisoDiff& diff) {
		throw invalid_argument("input image type not supported by this plugin, please use 'aniso'");
	}
};


template <typename Data2D>
struct __dispatch<Data2D, true> {
	static void apply(Data2D& image, const C2DAnisoDiff& diff) {

		C2DSLImage work_copy(image);
		
		diff.work(work_copy); 
		
		auto di = image.begin();
		for (auto wi = work_copy.begin(); wi != work_copy.end(); ++wi, ++di) {
			*di = static_cast<typename Data2D::value_type>((*wi + 8) / 16); 
		}
	}
};


template <class Data2D>
typename C2DAnisoDiff::result_type C2DAnisoDiff::operator () (Data2D& image) const
{
	const bool is_supported = ::boost::is_integral<typename Data2D::value_type>::value; 
	__dispatch<Data2D, is_supported>::apply(image, *this); 
	return 0; 
}

CAnisoDiff2DImageFilter::CAnisoDiff2DImageFilter(int maxiter, float epsilon):
	m_filter(maxiter, epsilon)
{
}

void CAnisoDiff2DImageFilter::do_filter(P2DImage& image) const
{
	wrap_filter(m_filter, image); 
}

CAnisoDiff2DImageFilterFactory::CAnisoDiff2DImageFilterFactory():
	C2DFilterPlugin(plugin_name)
{
	add_help(param_iterations); 
	add_help(param_epsilon); 
}

C2DFilter *CAnisoDiff2DImageFilterFactory::create(const CParsedOptions& options) const
{
	int max_iter = param_iterations.get_value(options); 
	float epsilon = param_epsilon.get_value(options); 
	return C2DFilter *new CAnisoDiff2DImageFilter(max_iter, epsilon); 
}

const string CAnisoDiff2DImageFilterFactory::do_get_descr()const
{
	return "2D image anisotropic filter (with errors!)"; 
}

extern "C" EXPORT CPluginBase *get_plugin_interface()
{
	return new CAnisoDiff2DImageFilterFactory(); 
}






} // end namespace ianiso_2dimage_filter
