/* -*- mia-c++  -*-
 *
 * This file is part of MIA - a toolbox for medical image analysis 
 * Copyright (c) Leipzig, Madrid 1999-2014 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 <cmath>
#include <functional>
#include <stdexcept>

#include <boost/type_traits.hpp>
#include <mia/2d/filter.hh>
#include <libmia/fftkernel.hh>

// to avoid the indention
#define NSBEGIN namespace lnfft_2dimage_filter {
#define NSEND }

NSBEGIN

NS_MIA_USE;
using namespace std; 


static const char *plugin_name = "lnfft"; 
static const CStringOption param_kernel("kernel", "filter kernel", NULL); 


class C2DLnFFT: public C2DFilter {
public: 
	C2DLnFFT(C2DFFTKernelFactory::ProductPtr kernel); 
	~C2DLnFFT();
	
	void init(const C2DBounds& size); 
	
	template <typename T>
	C2DLnFFT::result_type operator()(const T2DImage<T>& image)const; 
		
	
private: 
	C2DFFTKernelFactory::ProductPtr m_kernel; 
	C2DBounds m_size; 
	union {
		fftw_complex *m_cbuffer; 
		double   *m_fbuffer; 
	}; 
	fftw_plan m_forward_plan; 
	fftw_plan m_backward_plan; 
	size_t m_realsize_x; 
};

class C2DImageLnFFTFilter: public C2DImageFilterBase {
	mutable C2DLnFFT m_filter; 
public:
	C2DImageLnFFTFilter(C2DFFTKernelFactory::ProductPtr kernel);

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

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


C2DLnFFT::C2DLnFFT(C2DFFTKernelFactory::ProductPtr kernel):
	m_kernel(kernel), 
	m_size(0,0), 
	m_cbuffer(NULL), 
	m_realsize_x(0)
{
}

C2DLnFFT::~C2DLnFFT()
{
	if (m_cbuffer) {
		fftw_free(m_cbuffer); 
		fftw_destroy_plan( m_forward_plan); 
		fftw_destroy_plan( m_backward_plan); 		
	}
}

void C2DLnFFT::init(const C2DBounds& size)
{
	if (m_size == size)
		return; 
	cvdebug() << "C2DLnFFT::init ...\n"; 
	
	m_size = size; 
	m_realsize_x = 2 * (m_size.x /2 + 1); 

	cvdebug() << "size = " << m_size.x << ", " << m_size.y << "\n"; 
	if (m_cbuffer) {
		delete[] m_cbuffer; 
		fftw_destroy_plan( m_forward_plan); 
		fftw_destroy_plan( m_backward_plan); 		
	}
	
	m_cbuffer = (fftw_complex*)fftw_malloc(sizeof(fftw_complex) * m_size.y * m_realsize_x);
	// create the fftw plans
	
	cvdebug() << "buffer at " << m_cbuffer << ":" << m_fbuffer << "\n"; 
	
	if (!m_cbuffer) {
		throw runtime_error("unable to allocate fftw buffers"); 
	}
	
	 m_forward_plan = fftw_plan_dft_r2c_2d(m_size.y, m_size.x,
						m_fbuffer, m_cbuffer, FFTW_ESTIMATE);

	 cvdebug() << "forward plan at " <<  m_forward_plan << "\n"; 
	 if (!m_forward_plan) {
		 //fftw_free(m_cbuffer); 
		 throw runtime_error("unable to create forward plans ..."); 
	 }
	 
	 m_backward_plan = fftw_plan_dft_c2r_2d(m_size.y, m_size.x,
						 m_cbuffer, m_fbuffer, FFTW_ESTIMATE);
	 
	 if (!m_backward_plan) {
		 //fftw_free(m_cbuffer); 
		 //fftw_destroy_plan( m_forward_plan); 		 
		 throw runtime_error("unable to create backward plans ..."); 
	 }
	 
	 cverb << "done\n"; 
}

template <typename T, bool is_integral>
struct FBackConvert {
	FBackConvert(double scale):
		m_scale(scale) 
		{
			cvdebug() << "scale = " << m_scale <<"\n"; 
		}
	
	T operator ()(double x) {
		return T(exp(x *m_scale) - 1.0); 
	}
private: 
	double m_scale; 
};

template <typename T>
struct FBackConvert<T, true> {
	FBackConvert(double scale):
		m_scale(scale) 
		{
			cvdebug() << "scale = " << m_scale <<"\n"; 
		}
	
	T operator ()(double x) {
		double xc = exp(x *m_scale) - 1.0; 
		return xc < numeric_limits<T>::min() ? numeric_limits<T>::min() : 
			( xc < numeric_limits<T>::max() ?  T(xc) : numeric_limits<T>::max()); 
	}
private: 
	double m_scale; 
};

struct ForwardLogify {
	ForwardLogify(double f):m_f(f){}
	double operator ()(double x) {
		return m_f * log(x + 1.0); 
	}
private: 
	double m_f; 
};

template <typename T>
C2DLnFFT::result_type C2DLnFFT::operator()(const T2DImage<T>& image)const
{
	ForwardLogify flp(1.0); 
	ForwardLogify flm(-1.0); 
	
	cvdebug() << "C2DLnFFT::operator() begin\n";
	for (size_t y = 0; y < image.get_size().y; ++y) 
		transform(&image(0,y), &image(0,y) + image.get_size().x, &m_fbuffer[y * m_realsize_x], y & 1 ? flm : flp); 
		
	cvdebug() << "C2DLnFFT::operator() forward transform\n";
	fftw_execute( m_forward_plan); 

	cvdebug() << "C2DLnFFT::operator() call kernel\n";
	m_kernel->apply(m_size, m_cbuffer); 
	
	cvdebug() << "C2DLnFFT::operator() backward transform\n";
	fftw_execute( m_backward_plan); 

	cvdebug() << "C2DLnFFT::operator() back copy\n";
	T2DImage<T> *result = new T2DImage<T>(image.get_size(), image); 
	
	
	const bool is_integral = ::boost::is_integral<T>::value; 
	
	FBackConvert<T, is_integral> convertp(1.0/ image.size());
	FBackConvert<T, is_integral> convertm(-1.0/ image.size());
	// copy back and scale filtered output
	for (size_t y = 0; y < image.get_size().y; ++y)
		transform(&m_fbuffer[y * m_realsize_x], &m_fbuffer[y * m_realsize_x + m_size.x] , 
			  &(*result)(0,y), (y & 1) ? convertm : convertp);
	
	cvdebug() << "C2DLnFFT::operator() end\n";
	return P2DImage(result); 
	
}

C2DImageLnFFTFilter::C2DImageLnFFTFilter(C2DFFTKernelFactory::ProductPtr kernel):
	m_filter(kernel)
{
}
	

P2DImage C2DImageLnFFTFilter::do_filter(const C2DImage& image) const
{
	m_filter.init(image.get_size()); 
	return wrap_filter(m_filter, image); 
}


C2DImageLnFFTFilterFactory::C2DImageLnFFTFilterFactory():
	C2DFilterPlugin(plugin_name)
{
}

C2DFilter *C2DImageLnFFTFilterFactory::create(const CParsedOptions& options) const
{
	string kernel_descr = param_kernel.get_value(options); 
	if (kernel_descr.empty())
	    throw invalid_argument("fft: no kernel given"); 
	
	C2DFFTKernelPluginHandler fftkh; 
	C2DFFTKernelFactory::ProductPtr kernel(C2DFFTKernelFactory::produce(kernel_descr.c_str(), fftkh)); 
	if (!kernel) 
		throw invalid_argument(string("fft: unable to create cernel from '") + kernel_descr + "'");
	
	return ProductPtr(new C2DImageLnFFTFilter(kernel)); 
}

const string C2DImageLnFFTFilterFactory::do_get_descr()const
{
	return "2D image frequency domain filters based on LnFFT";
}


struct FCompare {
	bool operator () (float a, float b){
		return ::fabs(a-b) < 0.0001; 
	}
};


int C2DImageLnFFTFilterFactory::do_test() const 
{
	srand48(time(NULL)); 
	class CFFTKernelDummy : public CFFTKernel {
		virtual void do_apply(const C2DBounds& size, fftw_complex *cbuffer) const{};
	};
	
	CFFTKernel *k = new CFFTKernelDummy(); 
	C2DFFTKernelFactory::ProductPtr kernel(k); 
	
	C2DImageLnFFTFilter f(kernel); 
	
	C2DBounds size(16, 16); 
	C2DFImage *src = new C2DFImage(size); 
	P2DImage src_wrap(src); 

	for (C2DFImage::iterator i = src->begin(); i != src->end(); ++i)
		*i = 1; //drand48(); 
	
	P2DImage res_wrap = f.filter(src_wrap); 
	
	C2DFImage *res = res_wrap.getC2DFImage(); 
	if (!res) {
		cvfail() << get_name() << "image pixel type\n"; 
		return -1; 
	}
	
	if (res->get_size() != src->get_size()) {
		cvfail() << get_name() << "image size\n"; 
		return -1; 
	}
	
	if (equal(src->begin(), src->end(), res->begin(), FCompare()))
		return 1; 
	
	if (cverb.show_debug()) {
		for(C2DFImage::const_iterator s = src->begin(), r = res->begin(); s != src->end(); ++s, ++r)
			cvdebug() << *s << " : " << *r << "\n"; 
		    
	}
	
	return -1; 
}


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



NSEND
