/* -*- 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 <mia/core/cmeansinit/kmeans.hh>

#include <iomanip>

NS_MIA_BEGIN

CKMeansInitializer::CKMeansInitializer(size_t nclasses):m_nclasses(nclasses)
{
}

bool kmeans_step(const CMeans::NormalizedHistogram& nh, std::vector<int>& classification,
		 CMeans::DVector& classes, size_t l, int& biggest_class )
{
	cvdebug()<<  "kmeans enter: ";
	for (size_t i = 0; i <= l; ++i )
		cverb << std::setw(8) << classes[i]<< " ";  
	cverb << "\n"; 
	
	biggest_class = -1; 
	const double convLimit = 0.005;	// currently fixed
	
	std::vector<double> sums(classes.size()); 
	std::vector<double> count(classes.size()); 
	
	bool conv = false;
	int iter = 50; 
	
	while( iter-- && !conv) {

		sort(classes.begin(), classes.end()); 
		
		// assign closest cluster center
		auto ob = classification.begin(); 
		for (auto b = nh.begin(); b != nh.end(); ++b, ++ob) {
			const double val = b->first;
			double dmin = std::numeric_limits<double>::max();
			int c = 0;
			for (size_t i = 0; i <= l; i++) {
				double d = fabs (val - classes[i]);
				if (d < dmin) {
					dmin = d;
					c = i;
				};
			};
			*ob = c; 
			
			count[c] += b->second;
			sums[c] += val * b->second;
		};

		
		// recompute cluster centers
		conv = true;
		double max_count = 0; 
		for (size_t i = 0; i <= l; i++) {
			if (count[i]) {
				double a = sums[i] / count[i];
				if (a  && fabs ((a - classes[i]) / a) > convLimit)
					conv = false;
				classes[i] = a;
				
				if (max_count < count[i]) {
					max_count  = count[i]; 
					biggest_class = i; 
				}
			} else { // if a class is empty move it closer to neighbour 
				if (i == 0)
					classes[i] = (classes[i] + classes[i + 1]) / 2.0; 
				else
					classes[i] = (classes[i] + classes[i - 1]) / 2.0; 
				conv = false;
			}
			sums[i] = 0;
			count[i] = 0;
		};
	};

	cvinfo()<<  "kmeans: " << l + 1 << " classes " << 50 - iter << "  iterations";
	for (size_t i = 0; i <= l; ++i )
		cverb << std::setw(8) << classes[i]<< " ";  
	cverb << "\n"; 
	
	return conv; 

}
	

CMeans::DVector CKMeansInitializer::run(const CMeans::NormalizedHistogram& nh) const
{
	if (m_nclasses < 2)
		throw create_exception<std::invalid_argument>("cmeans: requested ", m_nclasses, 
						    "class(es), required are at least two");
	
	if ( nh.size() < m_nclasses ) 
		throw create_exception<std::invalid_argument>("kmeans: insufficient input: want ", m_nclasses , 
						    " classes, but git only ",  nh.size(), " distinct input values"); 

	CMeans::DVector classes(m_nclasses);

	std::vector<int> classification(nh.size(), 0);
	
	double sum = 0.0;
	for(auto h : nh) {
		sum += h.first * h.second;
	}; 

	
	// simple initialization splitting at the mean 
	classes[0] = sum / 1.99;  
	classes[1] = sum / 2.01;

	cvinfo() << "kmeans: initial classes: " << classes << "\n"; 
	
	// first run calles directly 
	int biggest_class = 0; 
	kmeans_step(nh, classification, classes, 1, biggest_class);
	
	// further clustering always splits biggest class 
	for (size_t  l = 2; l < m_nclasses; l++) {
		const size_t pos = biggest_class > 0 ? biggest_class - 1 : biggest_class + 1; 
		classes[l] = 0.5 * (classes[biggest_class] + classes[pos]);
		kmeans_step(nh, classification, classes, l, biggest_class); 
	};		
	
	// some post iteration until centers no longer change 
	for (size_t  l = 1; l < 3; l++) {
		if (kmeans_step(nh, classification, classes, m_nclasses - 1, biggest_class)) 
			break; 
	}
	return classes; 
        
}

CKMeansInitializerPlugin::CKMeansInitializerPlugin():
	CMeansInitializerSizedPlugin("kmeans")
{
}


CMeansInitializerPlugin::Product * CKMeansInitializerPlugin::do_create() const
{
        return new CKMeansInitializer(get_size_param());
}

const std::string CKMeansInitializerPlugin::do_get_descr() const
{
        return "C-Means initializer that sets the initial class centers by using a k-means classification";
}

extern "C" EXPORT CPluginBase  *get_plugin_interface()
{

	return new CKMeansInitializerPlugin(); 
}


NS_MIA_END
