/* -*- 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 <mia/core/histogram.hh>
#include <mia/core/cmdlineparser.hh>
#include <mia/2d.hh>

#include <memory>
#include <vector>
#include <numeric>

using namespace mia;
using namespace std; 

const SProgramDescription g_description = {
        {pdi_group, "Analysis, filtering, combining, and segmentation of 2D images"}, 
	{pdi_short, "Run a fuzzy c-means segmentation of a 2D image."}, 	
	{pdi_description, "This program is a  implementation of a fuzzy c-means segmentation algorithm"},
	{pdi_example_descr, "Run a 5-class segmentation over inpt image input.v and store the class "
	 "probability images in cls.v."}, 
	{pdi_example_code, "-i input.v -a 5 -o cls.v"}
}; 


class CSegment2DFuzzy: public TFilter<P2DImage> {
public: 
	struct Params {
		Params();
		vector<double> class_centres; 
		P2DFilter neighbourhood_filter; 
		float p; 
		float alpha; 
		float epsilon; 
	}; 

	CSegment2DFuzzy(const Params& p); 

	template <typename T> 
	P2DImage operator() (const T2DImage<T>& image); 

	const C2DImageVector& get_probs() const; 
	P2DImage get_gain() const; 
private:
	Params m_params; 
	C2DImageVector m_probs;
	P2DImage m_gain; 
}; 

CSegment2DFuzzy::Params::Params():
	p(2), 
	alpha(0.7), 
	epsilon(0.01)
{
}

CSegment2DFuzzy::CSegment2DFuzzy(const Params& p):
	m_params(p)
{
}

const C2DImageVector& CSegment2DFuzzy::get_probs() const
{
	return m_probs; 
}

P2DImage CSegment2DFuzzy::get_gain() const
{
	return m_gain; 
}

class MyFCM {
public: 
	template <typename T> 
	MyFCM(const T2DImage<T>& image, const CSegment2DFuzzy::Params& params);

	void run();



private: 
	friend class CSegment2DFuzzy; 
	void filter_neighborhood(C2DFImage& deltasq); 
	void estimate_prob(const vector<C2DFImage>& deltasq); 
	P2DImage get_gain()const; 
	C2DImageVector get_probs() const; 


	template <typename T> 
	P2DImage get_corrected_image() const; 
	
	C2DBounds m_size; 
	C2DFImage m_image;
	C2DFImage m_gain; 
	P2DFilter m_neighbourhood; 
	vector<double> m_class_centres; 
	vector<C2DFImage> m_probs; 
	float m_p; 
	float m_alpha; 
	float m_epsilon; 
}; 

/*
  Initialization of needed fields and transformation to the logarithmic space 
*/
template <typename T> 
MyFCM::MyFCM(const T2DImage<T>& image, const CSegment2DFuzzy::Params& params):
	m_size(image.get_size()), 
	m_image(m_size), 
	m_gain(m_size), 
	m_neighbourhood(params.neighbourhood_filter), 
	m_class_centres(params.class_centres.size()), 
	m_probs(m_class_centres.size(), C2DFImage(m_size)), 
	m_p(params.p), 
	m_alpha(params.alpha),
	m_epsilon(params.epsilon)
{
	fill(m_gain.begin(), m_gain.end(), 0.0); 
	transform(image.begin(), image.end(), m_image.begin(), [](float x){return log(x + 1);}); 
	transform(params.class_centres.begin(), params.class_centres.end(), 
		  m_class_centres.begin(), [](float x){return log(x + 1);}); 
}

P2DImage MyFCM::get_gain()const
{
	return P2DImage(m_gain.clone()); 
	
}

C2DImageVector MyFCM::get_probs() const
{
	C2DImageVector result; 
	
	transform(m_probs.begin(), m_probs.end(), back_inserter(result), 
		  [](const C2DFImage& image){ return P2DImage(image.clone());}); 
	return result; 
}

template <typename T> 
P2DImage MyFCM::get_corrected_image()const
{
	T2DImage<T> *result = new T2DImage<T>(m_image.get_size(), m_image); 
	
	transform(m_image.begin(), m_image.end(), m_gain.begin(), result->begin(), 
		  [](float y, float b){return mia_round_clamped<T>(exp(y-b)); }); 
	
	return P2DImage(result); 
}

void MyFCM::filter_neighborhood(C2DFImage& deltasq) 
{
	auto h = m_neighbourhood->filter(deltasq);
	const C2DFImage& h2 = dynamic_cast<const C2DFImage&>(*h); 
	
	transform(h2.begin(), h2.end(), deltasq.begin(), deltasq.begin(),
		  [this](float m, float D){return (D + m_alpha * m);});
}

void MyFCM::estimate_prob(const vector<C2DFImage>& deltasq)
{
	const float SMALL = 1e-8; 
	const int nclasses = m_probs.size(); 

	vector<C2DFImage::const_iterator> idsq(nclasses); 
	vector<C2DFImage::iterator> iprob(nclasses); 
	
	transform(deltasq.begin(), deltasq.end(), idsq.begin(), [](const C2DFImage& img){return img.begin();}); 
	transform(m_probs.begin(), m_probs.end(), iprob.begin(), [](C2DFImage& img){return img.begin();}); 

	auto ii = m_image.begin();  

	while (idsq[0] != deltasq[0].end()) {
		float sum = 0.0;
		bool defined_prob = false; 
		cvdebug() << "pixel=" << *ii++ << ", v = "; 
		for (int i = 0; i < nclasses && !defined_prob; ++i) {
			auto v = *idsq[i];
			cverb << v << " "; 
			// if this value is very small, then the pixel has a very high probability to be 
			// of the current class, 
			if (fabs(v) >= SMALL) 
				sum += 1.0f / v;
			else {
				for (int j = 0; j < nclasses; ++j) {
					if (j != i) 
						*iprob[j] = 0.0; 
					else 
						*iprob[j] = 1.0; 
				}
				defined_prob = true;
			}
		}
		if (!defined_prob) {
			cvdebug() << ", probs = (sum="<<sum<<") "; 
			for (int i = 0; i < nclasses; ++i) {
				*iprob[i] =  1.0f / (*idsq[i] * sum); 
				cverb << *iprob[i] << ", ";
			}
			cverb << "\n"; 
		} else {
			cvdebug() << "probs = permutation of (1 0 0 ...) \n"; 
		}
		for (int i = 0; i < nclasses; ++i) {
			++iprob[i]; 
			++idsq[i]; 
		}
	
	}
}

void MyFCM::run()
{
	// helper for differences
	vector<C2DFImage> deltasq(m_probs.size(), C2DFImage(m_size)); 
	C2DFImage corrected(m_size); 
	const int nclasses = m_class_centres.size(); 

	cvmsg() << "Initial class centers = " << m_class_centres << "\n"; 

	auto old_class_centres = m_class_centres; 
	float delta_cls = 0.0; 
	int maxiter = 30; 
	do  {
		
		// evaluate the b-field corrected image \| \y_k - \beta_k \|^2  in log-space 
		transform(m_image.begin(), m_image.end(), m_gain.begin(), corrected.begin(), 
			  [](float y, float b) {return y - b;}); 
			
		for (int i = 0; i < nclasses; ++i) {
			float vi = m_class_centres[i]; 
			
			// evaluate D_{ik} = \| \y_k - \beta_k - v_i \|^2  -> deltasq
			transform(corrected.begin(), corrected.end(), deltasq[i].begin(), 
				  [&vi](float x) -> float {
					  const float value = x - vi; 
					  return value * value; 
				  }); 
			
			// apply the neighbourhood filter
			filter_neighborhood(deltasq[i]);

			// if p = 2, the exponent becomes 1.0 and we can save this step 
			if (m_p != 2.0) {
				const float p_help = 1 / (1.0 - m_p); 
				transform(deltasq[i].begin(), deltasq[i].end(), deltasq[i].begin(), 
					  [&p_help](float x) { return pow(x, p_help); }); 
			}
		}

		
		
		// evaluate the probabilites according to eq. 13
		estimate_prob(deltasq);
		
		// now evaluate the power to p of the probabilities 
		if (m_p != 2.0) {
			for (int i = 0; i < nclasses; ++i) {
				transform(m_probs[i].begin(),  m_probs[i].end(), deltasq[i].begin(), 
					  [this](float x) { return pow(x, m_p);}); 
			}
		} else {
			for (int i = 0; i < nclasses; ++i) {
				transform(m_probs[i].begin(),  m_probs[i].end(), deltasq[i].begin(), 
					  [](float x){return x*x;});
			}
		}
		
                // apply the neighbourhood filter to the corrected image 
		filter_neighborhood(corrected);
		
		// evaluate new class centres according to eq. 15
		for (int i = 0; i < nclasses; ++i) {
			float sum = 0.0f; 
			float n = 0.0f; 
			auto ic = corrected.begin(); 
			for (auto up = deltasq[i].begin(); up != deltasq[i].end(); ++ic, ++up) {
				sum += *up * *ic; 
				n += *up; 
			}
			m_class_centres[i] = sum / (( 1.0 + m_alpha) * n); 
		}
		
		// now update the gain field according to eq. 19  
		
		vector<C2DFImage::const_iterator> iprobs(nclasses); 
		transform(deltasq.begin(), deltasq.end(), iprobs.begin(), [](const C2DFImage& img){return img.begin();});

		auto ii = m_image.begin(); 
		for (auto ig = m_gain.begin(); ig != m_gain.end(); ++ig, ++ii) {
			float sum = 0.0; 
			float n = 0.0; 
			for (int i = 0; i < nclasses; ++i) {
				sum += m_class_centres[i] * *iprobs[i]; 
				n += *iprobs[i]; 
				++iprobs[i]; 
			}
			*ig = *ii - sum / n ; 
		}
		delta_cls = inner_product(old_class_centres.begin(), old_class_centres.end(), 
					  m_class_centres.begin(), 0.0f, 
					  [](float x, float y) {return x + y;}, 
					  [](float x, float y) -> float {float d = x-y; return d * d; });
		cvmsg() << "Class centers = " << m_class_centres << ", delta = " << delta_cls << "\n"; 
		old_class_centres = m_class_centres; 
		--maxiter; 
	} while (delta_cls > m_epsilon && maxiter); 
	if (!maxiter) 
		throw runtime_error("no convergence"); 
}


template <typename T> 
P2DImage CSegment2DFuzzy::operator() (const T2DImage<T>& image)
{
	MyFCM myfmc(image,  m_params); 
	myfmc.run();
	m_probs = myfmc.get_probs();
	m_gain = myfmc.get_gain();
	
	return myfmc.get_corrected_image<T>(); 
}


int do_main(int argc, char *argv[])
{

	string in_filename; 
	string cls_filename; 
	string out_filename; 
	string gain_filename; 

	CSegment2DFuzzy::Params params; 
	unsigned int n_classes = 3; 
	const auto& image2dio = C2DImageIOPluginHandler::instance();
	
	CCmdOptionList opts(g_description);
	opts.set_group("File I/O"); 
	opts.add(make_opt( in_filename, "in-file", 'i', "image to be segmented", CCmdOptionFlags::required_input , &image2dio)); 
	opts.add(make_opt( cls_filename, "cls-file", 'c', "class probability images, the image type must "
			   "support multiple images and floating point values", CCmdOptionFlags::output, &image2dio )); 
	opts.add(make_opt( out_filename, "out-file", 'o', "B-field corrected image", CCmdOptionFlags::output , &image2dio)); 
	opts.add(make_opt( gain_filename, "gain-log-file", 'g', "Logarithmic gain field, the image type must "
			   "support floating point values",  CCmdOptionFlags::output , &image2dio)); 
	
	opts.set_group("Segmentation parameters");
	opts.add(make_opt( n_classes, "no-of-classes", 'n', "number of classes to segment"));
	opts.add(make_opt( params.class_centres, "class-centres", 'C', "initial class centers"));
	opts.add(make_opt( params.neighbourhood_filter, "shmean:shape=8n", "neighborhood", 'N', "neighborhood filter for B-field correction"));
	opts.add(make_opt( params.alpha, "alpha", 'a', "weight of neighborhood filter for B-field correction"));
	opts.add(make_opt( params.p, "fuzziness", 'p', 
			   " parameter describing the fuzzyness of mattar distinction")); 
	opts.add(make_opt( params.epsilon, "epsilon", 'e', 
			   "Stopping criterion for class center estimation.")); 

	
	if (opts.parse(argc, argv) != CCmdOptionList::hr_no)
		return EXIT_SUCCESS; 
	
	if (cls_filename.empty() && out_filename.empty() && gain_filename.empty()) {
		throw invalid_argument("Not a single output given, processing makes no sense");  
	}

	auto in_image = load_image2d(in_filename);	
	
	if (params.class_centres.empty()) {
		cvdebug() << "no class centers given at command line, estimate using k-means\n"; 

		ostringstream kmeans_filter; 
		kmeans_filter << "kmeans:c=" << n_classes; 
		
		P2DImage kmeans_image = run_filter(*in_image, kmeans_filter.str().c_str()); 
		params.class_centres = 
			kmeans_image->get_attribute_as<vector<double> >(ATTR_IMAGE_KMEANS_CLASSES); 
	}else{
		if (params.class_centres.size() != n_classes) {
			cvwarn() << "Specified initial class centres but also number of classes. "
				 << "The number of classes resulting from the given centres takes precedence\n"; 
		}
	}
	
	// initialize the functor
	CSegment2DFuzzy segment(params);
	auto b0_corrected = mia::accumulate (segment, *in_image);
	
	if (!out_filename.empty()) {
		if (!save_image(out_filename, b0_corrected)) {
			cvwarn() << "Unable to save B0 corrected image to '" << out_filename << "'\n"; 
		}
	}

	if (!cls_filename.empty()) {
		auto probs = segment.get_probs(); 
		if (!C2DImageIOPluginHandler::instance().save(cls_filename, probs)) {
			cvwarn() << "Unable to save class probabilities to '" << cls_filename << "'\n"; 
		}
	}

	if (!gain_filename.empty()) {
		auto gain = segment.get_gain(); 
		if (!save_image(gain_filename, gain)) {
			cvwarn() << "Unable to save gain to '" << gain_filename << "'\n"; 
		}
	}

	return EXIT_SUCCESS; 
}

#include <mia/internal/main.hh>
MIA_MAIN(do_main); 


