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

#include <mia/core/filter.hh>
#include <mia/3d/filter.hh>
#include <mia/3d/fullcost/taggedssd.hh>

NS_BEGIN(taggedssd_3d)
NS_MIA_USE
using namespace std; 

C3DTaggedSSDCost::C3DTaggedSSDCost(const std::string& src_x, 
				   const std::string& ref_x, 
				   const std::string& src_y, 
				   const std::string& ref_y, 
				   const std::string& src_z, 
				   const std::string& ref_z,
				   double weight):
	C3DFullCost(weight)
	
{
	auto& imgio = C3DImageIOPluginHandler::instance();  

	m_src_key[0] = imgio.load_to_pool(src_x);
	m_src_key[1] = imgio.load_to_pool(src_y);
	m_src_key[2] = imgio.load_to_pool(src_z);
	
	m_ref_key[0] = imgio.load_to_pool(ref_x);
	m_ref_key[1] = imgio.load_to_pool(ref_y);
	m_ref_key[2] = imgio.load_to_pool(ref_z);

	add(property_gradient);
}

inline double sqd(double x, double y) {
	double h = x - y; 
	return h * h; 
}

struct FTaggedSSDAccumulatorX : public TFilter<double> {
	FTaggedSSDAccumulatorX(C3DFVectorfield& force, float scale):
		m_force(force), m_scale(scale) {} 
	
	template <typename T, typename S> 
	double operator () ( const T3DImage<T>& src, const T3DImage<S>& ref) {
		assert(src.size() == ref.size()); 
		
		double result = 0.0; 
		auto is = src.begin(); 
		auto ir = ref.begin(); 
		auto iforce = m_force.begin() + 1; 
		for (size_t z = 0; z < src.get_size().z; ++z) 
			for (size_t y = 0; y < src.get_size().y; ++y) {
				result += sqd(*is++, *ir++); 
				for (size_t x = 1; x < src.get_size().x - 1; ++x, is++, ir++, ++iforce) {
					double delta = *is - *ir; 
					iforce->x = m_scale * delta * (is[1] - is[-1]); 
					result += delta * delta; 
				}
				result += sqd(*is++, *ir++); 
				iforce += 2; 
			} 
		return result; 
	}; 
private: 
	C3DFVectorfield& m_force; 
	float m_scale; 
}; 


struct FTaggedSSDAccumulatorY : public TFilter<double> {
	FTaggedSSDAccumulatorY(C3DFVectorfield& force, float scale):
		m_force(force), m_scale(scale) {} 
	
	template <typename T, typename S> 
	double operator () ( const T3DImage<T>& src, const T3DImage<S>& ref) {
		assert(src.size() == ref.size()); 
		double result = 0.0; 
		
		int dx = src.get_size().x; 
		auto is = src.begin(); 
		auto ir = ref.begin(); 
		auto iforce = m_force.begin(); 
		for (size_t z = 0; z < src.get_size().z; ++z) {
			for (size_t x = 0; x < src.get_size().x; ++x, ++is, ++ir, ++iforce) {
				result += sqd(*is, *ir); 
			}
			for (size_t y = 1; y < src.get_size().y - 1; ++y) {
				for (size_t x = 0; x < src.get_size().x; ++x, is++, ir++, ++iforce) {
					double delta = *is - *ir; 
					iforce->y = m_scale * delta * (is[dx] - is[-dx]); 
					result += delta * delta; 
				}
			}
			for (size_t x = 0; x < src.get_size().x; ++x, ++is, ++ir, ++iforce) {
				result += sqd(*is, *ir); 
			}
		}
		return result; 
	}; 
private: 
	C3DFVectorfield& m_force;
	float m_scale; 
}; 


struct FTaggedSSDAccumulatorZ : public TFilter<double> {
	FTaggedSSDAccumulatorZ(C3DFVectorfield& force, float scale):
		m_force(force), m_scale(scale) {} 
	
	template <typename T, typename S> 
	double operator () ( const T3DImage<T>& src, const T3DImage<S>& ref) {
		assert(src.size() == ref.size()); 
		
		double result = 0.0; 
		int dxy = src.get_size().x * src.get_size().y; 
		auto is = src.begin(); 
		auto ir = ref.begin(); 
		auto iforce = m_force.begin(); 
		
		for (size_t y = 0; y < src.get_size().y; ++y) {
			for (size_t x = 0; x < src.get_size().x; ++x, ++is, ++ir, ++iforce) {
				result += sqd(*is, *ir); 
			}
		}

		for (size_t z = 1; z < src.get_size().z - 1; ++z) {
			for (size_t y = 0; y < src.get_size().y; ++y) {
				for (size_t x = 0; x < src.get_size().x; ++x, is++, ir++, ++iforce) {
					double delta = *is - *ir; 
					iforce->z = m_scale * delta * (is[dxy] - is[-dxy]); 
					result += delta * delta; 
				}
			}
		}
		for (size_t y = 0; y < src.get_size().y; ++y) {
			for (size_t x = 0; x < src.get_size().x; ++x, ++is, ++ir, ++iforce) {
				result += sqd(*is, *ir); 
			}
		}
		return result; 
	}; 
private: 
	C3DFVectorfield& m_force; 
	float m_scale; 
}; 


double C3DTaggedSSDCost::do_evaluate(const C3DTransformation& t, CDoubleVector& gradient) const
{
	C3DFVectorfield force(get_current_size()); 
	double value = 0.0; 
	float scale = 0.5f / m_ref_scaled[0]->size(); 
	{
		auto temp = t(*m_src_scaled[0]); 
		FTaggedSSDAccumulatorX acc(force, scale);
		double r = mia::accumulate(acc, *temp, *m_ref_scaled[0]);
		value += r; 
	}
	{
		auto temp = t(*m_src_scaled[1]); 
		FTaggedSSDAccumulatorY acc(force, scale);
		double r = mia::accumulate(acc, *temp, *m_ref_scaled[1]);
		value += r; 
	}
	{
		auto temp = t(*m_src_scaled[2]); 
		FTaggedSSDAccumulatorZ acc(force, scale);
		double r = mia::accumulate(acc, *temp, *m_ref_scaled[2]);
		value += r; 
	}
	
	t.translate(force, gradient); 
	return value / (6.0 * m_ref_scaled[0]->size()); 	
}

void C3DTaggedSSDCost::do_set_size()
{
	TRACE_FUNCTION; 
	assert(m_src[0]); 
	assert(m_src[1]); 
	assert(m_src[2]); 
	
	assert(m_ref[0]); 
	assert(m_ref[1]); 
	assert(m_ref[2]); 

	if (!m_src_scaled[0] || m_src_scaled[0]->get_size() != get_current_size() ||
	    !m_ref_scaled[0] || m_ref_scaled[0]->get_size() != get_current_size() ||
	    !m_src_scaled[1] || m_src_scaled[1]->get_size() != get_current_size() ||
	    !m_ref_scaled[1] || m_ref_scaled[1]->get_size() != get_current_size() ||
	    !m_src_scaled[2] || m_src_scaled[2]->get_size() != get_current_size() ||
	    !m_ref_scaled[2] || m_ref_scaled[2]->get_size() != get_current_size()) {
		
		if (m_src[0]->get_size() == get_current_size()) {
			for (int i = 0; i < 3; ++i) {
				m_src_scaled[i] = m_src[i]; 
				m_ref_scaled[i] = m_ref[i];
			}
		}else{
			stringstream filter_descr; 
			filter_descr << "scale:s=[" << get_current_size()<<"]"; 
			auto scaler = C3DFilterPluginHandler::instance().produce(filter_descr.str()); 
			assert(scaler); 
			cvdebug() << "C3DImageFullCost:scale images to " << get_current_size() << 
				" using '" << filter_descr.str() << "'\n"; 
			for (int i = 0; i < 3; ++i) {
				m_src_scaled[i] = scaler->filter(*m_src[i]); 
				m_ref_scaled[i] = scaler->filter(*m_ref[i]);
			}
		}
	}
}

P3DImage C3DTaggedSSDCost::get_from_pool(const C3DImageDataKey& key)
{
	C3DImageIOPlugin::PData in_image_list = key.get();
		
	if (!in_image_list || in_image_list->empty())
		throw invalid_argument("C3DImageFullCost: no image available in data pool");

	return (*in_image_list)[0];

}

	

struct FTaggedSSDAccumulator :public TFilter<bool> {
	FTaggedSSDAccumulator():m_value(0.0){} 
	
	template <typename T, typename S> 
	bool operator () ( const T3DImage<T>& src, const T3DImage<S>& ref) {
		assert(src.size() == ref.size()); 
		m_value += std::inner_product (src.begin(), src.end(), ref.begin(), 0.0, 
					       [](double x, double y) { return x + y;}, 
					       [](double x, double y) -> double { double d = x - y; return d*d; }); 
		return true; 
	}; 
	double get_value() const {
		return m_value; 
	}
private: 
	double m_value; 
}; 

double C3DTaggedSSDCost::do_value(const C3DTransformation& t) const
{
	P3DImage temp[3]; 
	FTaggedSSDAccumulator acc;
	
	for (int i = 0; i < 3; ++i) {
		temp[i] = t(*m_src_scaled[i]); 
		mia::accumulate(acc, *temp[i], *m_ref_scaled[i]); 
	}

	return acc.get_value() / (6.0 * m_ref_scaled[0]->size()); 

}

double C3DTaggedSSDCost::do_value() const
{
	FTaggedSSDAccumulator acc;
	for (int i = 0; i < 3; ++i)
		mia::accumulate(acc, *m_src_scaled[i], *m_ref_scaled[i]); 
	return acc.get_value() / (6.0 * m_ref_scaled[0]->size()); 
}

struct FAccumulateMeanVariance: public TFilter<int>  {
	FAccumulateMeanVariance(); 
	
	template <typename T> 
	int operator() (const T3DImage<T>& image); 

	P3DFilter get_intensity_scaling_filter() const; 
private: 
	double sum; 
	double sum2; 
	double n; 
};
	
FAccumulateMeanVariance::FAccumulateMeanVariance():
	sum(0.0), 
	sum2(0.0), 
	n(0.0)
{
}

template <typename T> 
int FAccumulateMeanVariance::operator() (const T3DImage<T>& image)
{
	n += image.size(); 
	for (auto i = image.begin(); i != image.end(); ++i) {
		sum += *i; 
		sum2 += *i * *i; 
	}
	return 0; 
}

P3DFilter FAccumulateMeanVariance::get_intensity_scaling_filter() const
{
	double mean = sum / n; 
	double sigma = sqrt((sum2 - sum * sum / n) / (n - 1));

	// both images are of the same single color 
	if (sigma == 0.0) 
		return P3DFilter(); 

	// I want a conversion filter, that makes the images together zero mean 
	// and diversion 1
	stringstream filter_descr; 
	filter_descr << "convert:repn=float,map=linear,b=" << -mean/sigma << ",a=" << 1.0/sigma; 
	cvinfo() << "Will convert using the filter:" << filter_descr.str() << "\n"; 
	
	return C3DFilterPluginHandler::instance().produce(filter_descr.str()); 
}


void C3DTaggedSSDCost::do_reinit()
{
	TRACE_FUNCTION; 

	FAccumulateMeanVariance filter_creator; 
	//cvmsg() << "C3DImageFullCost: read " << m_src_key << " and " << m_ref_key << "\n"; 
	for (int i = 0; i < 3; ++i) {
		m_src[i] = get_from_pool(m_src_key[i]);
		m_ref[i] = get_from_pool(m_ref_key[i]);
		
		mia::accumulate(filter_creator, *m_src[i]); 
		mia::accumulate(filter_creator, *m_ref[i]); 

		m_src_scaled[i].reset(); 
		m_ref_scaled[i].reset(); 

		// the loading code should chack for proper loading 
		assert(m_src[i]); 
		assert(m_ref[i]); 
	
		if (m_src[i]->get_size() != m_ref[i]->get_size()) 
			throw runtime_error("C3DImageFullCost only works with images of equal size"); 
		
	
		if (m_src[i]->get_voxel_size() != m_ref[i]->get_voxel_size()) {
			cverr() << "C3DImageFullCost: src and reference image are of different pixel dimensions."
				<< "This code doesn't honour this and a proper scaling should be applied first."; 
		}
	}
	if (m_ref[0]->get_size() != m_ref[1]->get_size() ||
	    m_ref[0]->get_size() != m_ref[2]->get_size()) {
		throw invalid_argument("Input reference images are of different sizes"); 
	}

	// run intensity normalization to zero mean, variation one. 
	auto filter = filter_creator.get_intensity_scaling_filter(); 
	if (!filter)
		throw invalid_argument("Input images have constant intensity"); 

	for (int i = 0; i < 3; ++i) {
		m_src[i] = filter->filter(*m_src[i]); 
		m_ref[i] = filter->filter(*m_ref[i]); 
	}
}

bool C3DTaggedSSDCost::do_get_full_size(C3DBounds& size) const
{
	TRACE_FUNCTION; 
	assert(m_ref[0]); 
	if (size == C3DBounds::_0) {
		size = m_ref[0]->get_size(); 
		return true; 
	}else
		return 	size == m_ref[0]->get_size(); 

}


// plugin implementation 
class C3DTaggedSSDCostPlugin: public C3DFullCostPlugin {
public: 
	C3DTaggedSSDCostPlugin(); 
private: 
	C3DFullCost *do_create(float weight) const;
	const std::string do_get_descr() const;
	std::string m_srcx_name;
	std::string m_refx_name;
	std::string m_srcy_name;
	std::string m_refy_name;
	std::string m_srcz_name;
	std::string m_refz_name;
}; 

C3DTaggedSSDCostPlugin::C3DTaggedSSDCostPlugin():
C3DFullCostPlugin("taggedssd"), 
      m_srcx_name("srcx.@"), 
      m_refx_name("refx.@"), 
      m_srcy_name("srcy.@"), 
      m_refy_name("refy.@"),
      m_srcz_name("srcz.@"),
      m_refz_name("refz.@")
										      
{
	const auto& io = C3DImageIOPluginHandler::instance(); 
	add_parameter("srcx", new CStringParameter(m_srcx_name, CCmdOptionFlags::input, "Study image X-tag", &io));
	add_parameter("refx", new CStringParameter(m_refx_name, CCmdOptionFlags::input, "Reference image  X-tag", &io));
	add_parameter("srcy", new CStringParameter(m_srcy_name, CCmdOptionFlags::input, "Study image Y-tag", &io));
	add_parameter("refy", new CStringParameter(m_refy_name, CCmdOptionFlags::input, "Reference image  Y-tag", &io));
	add_parameter("srcz", new CStringParameter(m_srcz_name, CCmdOptionFlags::input, "Study image Z-tag", &io));
	add_parameter("refz", new CStringParameter(m_refz_name, CCmdOptionFlags::input, "Reference image  Z-tag", &io));
}

C3DFullCost *C3DTaggedSSDCostPlugin::do_create(float weight) const
{
	return new C3DTaggedSSDCost(m_srcx_name, m_refx_name, 
				    m_srcy_name, m_refy_name, 
				    m_srcz_name, m_refz_name, 
				    weight); 
}

const std::string C3DTaggedSSDCostPlugin::do_get_descr() const
{
	return "Evaluates the Sum of Squared Differences similarity measure by using three tagged image pairs. "
		"The cost function value is evaluated based on all image pairs, but the gradient is composed by composing "
		"its component based on the tag direction."; 
}

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



NS_END


