/* -*- mia-c++  -*-
 *
 * This file is part of MIA - a toolbox for medical image analysis 
 * Copyright (c) Leipzig, Madrid 1999-2016 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 <iomanip>
#include <ostream>
#include <fstream>
#include <map>
#include <mia/core.hh>
#include <queue>

#include <mia/internal/main.hh>
#include <mia/2d/imageio.hh>
#include <mia/2d/filter.hh>
#include <mia/2d/ica.hh>


NS_MIA_USE;
using namespace std;

const SProgramDescription g_description = {
	{pdi_group, "Tools for Myocardial Perfusion Analysis"}, 

	{pdi_short, "Evaluate the time-intensity correlation in a series of images."}, 
	
	{pdi_description, "Given a set of images of temporal sucession, evaluates images that represent "
	 "the time-intensity correlation in horizontal and vertical direction as "
	 "well as average correlation of each pixel with its neighbors. "
	 "All input images must be of the same pixel type and size."}, 

	
	{pdi_example_descr,"Evaluate the time-intensity correaltions for an image series "
	 "imageXXXX.png starting at image 2 and stop at image 30. "
	 "Store the results in horizontal.exr, and vertical.exr."}, 
	
	{pdi_example_code,	"-i image0000.png -k 2 -e 30 -z horizontal.exr -t vertical.exr"}
}; 

struct FCorrelationAccumulator : public TFilter<bool> {

	FCorrelationAccumulator(const C2DBounds & size);

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

	P2DImage get_horizontal_corr() const;
	P2DImage get_vertical_corr() const;
	P2DImage get_avg_corr() const;
private:
	void evaluate_ver()const; 
	void evaluate_hor()const; 
	C2DDImage sx2;

	C2DDImage sxy_horizontal;
	C2DDImage sxy_vertical;

	C2DDImage sx;
	C2DDImage sy;
	C2DBounds size;
	
	mutable C2DFImage *corr_hor; 
	mutable P2DImage  pcorr_hor; 
	mutable C2DFImage *corr_ver;
	mutable P2DImage  pcorr_ver;

	size_t len;
};

int do_main( int argc, char *argv[] )
{
	string src_name("data0000.exr");
	string out_name("output.v");
	string out_hor_name("horizontal.v");
	string out_ver_name("vertical.v");
	size_t first =  2;
	size_t last  = 60;

	const auto& image2dio = C2DImageIOPluginHandler::instance();

	CCmdOptionList options(g_description);
	options.add(make_opt( src_name, "in-base", 'i', "input file name base", CCmdOptionFlags::required_input, &image2dio));
	options.add(make_opt( src_name, "outname", 'o', "output file name to save the avarage per-pixel correlation", 
			      CCmdOptionFlags::required_output, &image2dio));
	options.add(make_opt( out_hor_name, "horizontal", 'z', "horiZontal correlation output file name", 
			      CCmdOptionFlags::output, &image2dio));
	options.add(make_opt( out_ver_name, "vertical", 't', "verTical  correlation output file name", 
			      CCmdOptionFlags::output, &image2dio));
	options.add(make_opt( first, "skip", 'k', "skip images at beginning of series"));
	options.add(make_opt( last, "end", 'e', "last image in series"));

	if (options.parse(argc, argv) != CCmdOptionList::hr_no)
		return EXIT_SUCCESS; 


	size_t start_filenum = 0;
	size_t end_filenum  = 0;
	size_t format_width = 0;

	string src_basename = get_filename_pattern_and_range(src_name, start_filenum, end_filenum, format_width);

	if (start_filenum < first)
		start_filenum = first;
	if (end_filenum > last)
		end_filenum = last;

	// load images
	vector<P2DImage> series;
	for (size_t i = start_filenum; i < end_filenum; ++i) {
		string src_name = create_filename(src_basename.c_str(), i);
		P2DImage image = load_image<P2DImage>(src_name);
		series.push_back(image );
	}

	// evaluate all series correlation coefficients
	cvmsg()<< "Got series of " << series.size() << " images\n";
	FCorrelationAccumulator acc(series[0]->get_size());
	for (auto i = series.begin(); i != series.end(); ++i)
		::mia::accumulate(acc, **i);

	P2DImage hor = acc.get_horizontal_corr();
	P2DImage ver = acc.get_vertical_corr();
	P2DImage avgcorr = acc.get_avg_corr();

	if (!save_image(out_hor_name, hor))
		throw create_exception<runtime_error>( "unable to save horizontal correlation to '", out_hor_name,"'");

	if (!save_image(out_ver_name, ver))
		throw create_exception<runtime_error>( "unable to save vertical correlation to '", out_ver_name, "'");

	if (!save_image(out_name, avgcorr))
		throw create_exception<runtime_error>( "unable to save average correlation to '", out_name, "'");


	

	return EXIT_SUCCESS;

};



FCorrelationAccumulator::FCorrelationAccumulator(const C2DBounds & _size):
	sx2(_size),
	sxy_horizontal(_size),
	sxy_vertical(_size),
	sx(_size),
	size(_size),
	corr_hor(nullptr), 
	corr_ver(nullptr), 
	len(0)
{
}

template <typename T>
bool FCorrelationAccumulator::operator ()(const T2DImage<T>& image)
{
	if (image.get_size() != size)
		throw create_exception<invalid_argument>( "Input image size ",  size,  " expected, but got ", image.get_size());
	// sum x
	transform(image.begin(), image.end(), sx.begin(), sx.begin(), 
		  [](T x, double y){return x + y;}); 

	// sum x^2
	transform(image.begin(), image.end(), sx2.begin(), sx2.begin(), 
		  [](double x, double y){return x*x + y;}); 

	// sum horizontal
	for (size_t y = 0; y < size.y; ++y) {
		auto irow = image.begin_at(0,y);
		auto orow = sxy_horizontal.begin_at(0,y);
		for (size_t x = 0; x < size.x-1; ++x, ++irow, ++orow) {
			*orow += irow[0] * irow[1];
		}
	}

	// sum vertical
	for (size_t y = 1; y < size.y; ++y) {
		auto irow0 = image.begin_at(0,y-1);
		auto irow1 = image.begin_at(0,y);
		auto orow = sxy_vertical.begin_at(0,y-1);
		for (size_t x = 0; x < size.x; ++x, ++irow0, ++irow1,++orow) {
			*orow += *irow0 * *irow1;
		}
	}
	++len;
	return true;
}

P2DImage FCorrelationAccumulator::get_horizontal_corr() const
{
	if (!pcorr_hor)
		evaluate_hor(); 
	
	return pcorr_hor;
}
void FCorrelationAccumulator::evaluate_hor()const
{
	if (!len)
		throw create_exception<invalid_argument>( "No input images");

	
	corr_hor = new C2DFImage(C2DBounds(size.x-1, size.y));
	pcorr_hor.reset(corr_hor); 
	for (size_t y = 0; y < size.y; ++y) {
		auto irow_xy = sxy_horizontal.begin_at(0,y);
		auto irow_xx = sx2.begin_at(0,y);
		auto irow_yy = sx2.begin_at(1,y);
		auto irow_x  = sx.begin_at(0,y);
		auto irow_y  = sx.begin_at(1,y);
		auto orow    = corr_hor->begin_at(0,y);
		
		for (size_t x = 1; x < size.x;
		     ++x, ++irow_xy, ++irow_xx, ++irow_yy, ++irow_x, ++irow_y, ++orow) {
			
			const float ssxy = *irow_xy - *irow_x * *irow_y / len;
			const float ssxx = *irow_xx - *irow_x * *irow_x / len;
			const float ssyy = *irow_yy - *irow_y * *irow_y / len;
			
			if (fabs(ssxx) < 1e-10 && fabs(ssyy) < 1e-10)
				*orow =  1.0;
			else if (fabs(ssxx) < 1e-10 || fabs(ssyy) < 1e-10)
				*orow =  0.0;
			else
				*orow = (ssxy * ssxy) /  (ssxx * ssyy);
			
		}
		++irow_xy; ++irow_xx; ++irow_yy; ++irow_x; ++irow_y;
	}
}

P2DImage FCorrelationAccumulator::get_vertical_corr() const
{
	if (!pcorr_ver)
		evaluate_ver(); 
	return pcorr_ver;
}

void FCorrelationAccumulator::evaluate_ver()const
{
	if (!len)
		throw create_exception<invalid_argument>( "No input images");
	
	corr_ver = new C2DFImage(C2DBounds(size.x, size.y-1));
	pcorr_ver.reset(corr_ver); 
	
	for (size_t y = 0; y < size.y-1; ++y) {
		auto irow_xy = sxy_vertical.begin_at(0,y);
		auto irow_xx = sx2.begin_at(0,y);
		auto irow_yy = sx2.begin_at(0,y+1);
		auto irow_x  = sx.begin_at(0,y);
		auto irow_y  = sx.begin_at(0,y+1);
		auto orow    = corr_ver->begin_at(0,y);
		
		for (size_t x = 0; x < size.x;
		     ++x, ++irow_xy, ++irow_xx, ++irow_yy, ++irow_x, ++irow_y, ++orow) {
			
			const float ssxy = *irow_xy - *irow_x * *irow_y / len;
			const float ssxx = *irow_xx - *irow_x * *irow_x / len;
			const float ssyy = *irow_yy - *irow_y * *irow_y / len;
			
			if (ssxx == 0 && ssyy == 0)
				*orow =  1.0;
			else if (ssxx == 0 || ssyy == 0)
				*orow = 0.0;
			else
				*orow = (ssxy * ssxy) /  (ssxx * ssyy);
		}
	}
}

P2DImage FCorrelationAccumulator::get_avg_corr() const
{
	if (!pcorr_ver) 
		evaluate_ver(); 
	if (!pcorr_hor) 
		evaluate_hor();

	C2DFImage *result= new C2DFImage(C2DBounds(size.x, size.y));
	P2DImage presult(result);

	const C2DFImage& h = *corr_hor;
	const C2DFImage& v = *corr_ver;
	
	auto r = result->begin(); 
	auto ch = h.begin(); 
	auto cv = v.begin(); 
	
	*r++ = (*ch + *cv++) * 0.5f; 
	
	for (size_t x = 1; x < size.x-1; ++x, ++r, ++ch, ++cv) {
		*r = (*ch + ch[1] + *cv) * 1.0f/3.0f;
	}
	*r++ = (*ch++ + *cv++) * 0.5f;

	auto cvm = corr_ver->begin(); 
	for (size_t y = 1; y < size.y-1; ++y) {
		*r++ = (*ch + *cv++ + *cvm++) * 1.0f/3.0f;
		for (size_t x = 1; x < size.x-1; ++x, ++r, ++ch, ++cv, ++cvm)
			*r = (*ch + ch[1] + *cv + *cvm) * 0.25f;
		*r++ = (*ch++ + *cv++ + *cvm++) * 1.0f/3.0f;
	}

	assert(cv == v.end()); 	
	
	*r++ = (*ch + *cvm++) * 0.5f; 
	
	for (size_t x = 1; x < size.x-1; ++x, ++r, ++ch, ++cvm) {
		*r = (*ch + ch[1] + *cvm) * 1.0f/3.0f;
	}
	*r++ = (*ch++ + *cvm++) * 0.5f;

	assert(ch == h.end()); 
	assert(cvm == v.end()); 


	return presult; 
}

MIA_MAIN(do_main); 
