/* -*- 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 <functional>
#include <stdexcept>
#include <boost/type_traits.hpp>
#include <mia/2d/filter.hh>
#include <libmia/fftkernel.hh>

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

NSBEGIN

NS_MIA_USE;
using namespace std; 


static const char *plugin_name = "frequency"; 

class C2DFFT: public C2DFilter {
public: 
	C2DFFT(); 
	~C2DFFT();
	
	void init(const C2DBounds& size); 
	
	template <typename T>
	C2DFFT::result_type operator()(const T2DImage<T>& image)const; 
		
	
private: 
	C2DBounds m_size; 
	fftw_complex *m_cbuffer; 
	fftw_plan m_forward_plan; 
};

class C2DImageFFTFilter: public C2DImageFilterBase {
	mutable C2DFFT m_filter; 
public:
	C2DImageFFTFilter();

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

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


C2DFFT::C2DFFT():
	m_size(0,0), 
	m_cbuffer(NULL)
{
}

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

void C2DFFT::init(const C2DBounds& size)
{
	if (m_size == size)
		return; 
	cvdebug() << "C2DFFT::init ...\n"; 
	
	m_size = size; 

	cvdebug() << "size = " << m_size.x << ", " << m_size.y << "\n"; 
	if (m_cbuffer) {
		delete[] m_cbuffer; 
		fftw_destroy_plan( m_forward_plan); 
	}
	
	m_cbuffer = (fftw_complex*)fftw_malloc(sizeof(fftw_complex) * m_size.y * m_size.x);
	// create the fftw plans
	
	cvdebug() << "buffer at " << m_cbuffer << "\n"; 
	
	if (!m_cbuffer) {
		throw runtime_error("unable to allocate fftw buffers"); 
	}
	
	 m_forward_plan = fftw_plan_dft_2d(m_size.y, m_size.x,
					    m_cbuffer, m_cbuffer, FFTW_FORWARD, 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 ..."); 
	 }
	 
	 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(x *m_scale); 
	}
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 = x *m_scale; 
		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; 
};

template <typename T>
C2DFFT::result_type C2DFFT::operator()(const T2DImage<T>& image)const
{
	cvdebug() << "C2DFFT::operator() begin\n";
	typename T2DImage<T>::const_iterator i = image.begin(); 
	fftw_complex *t = m_cbuffer; 
	for (size_t y = 0; y < image.get_size().y; ++y) {
		float sign = y & 1 ? -1 : 1; 
		for (size_t x = 0; x < image.get_size().x; ++x, sign *=-1, ++i, ++t) {
			(*t)[0] = sign * *i; 
			(*t)[1] = 0.0; 
		}
	}
	
	cvdebug() << "C2DFFT::operator() forward transform\n";
	fftw_execute( m_forward_plan); 
	
	fftw_complex *cbuffer = m_cbuffer;
	
	for (size_t i = 0; i < image.size(); ++i, ++cbuffer) {
		(*cbuffer)[0] = sqrt((*cbuffer)[0] * (*cbuffer)[0] + (*cbuffer)[1] * (*cbuffer)[1]); 
	}

	double factor =1.0/sqrt(image.size()); 
	cvdebug() << "factor = " << factor << "\n"; 
	
	T2DImage<T> *result = new T2DImage<T>(image.get_size(), image);
	typename T2DImage<T>::iterator r = result->begin(); 
	cbuffer = m_cbuffer;
	
	for (size_t y = 0; y < image.get_size().y; ++y)
		for (size_t x = 0; x < image.get_size().x; ++x, ++cbuffer, ++r) {
			double val = factor * (*cbuffer)[0]; 
			*r = val > numeric_limits<T>::max() ? numeric_limits<T>::max(): 
				val < numeric_limits<T>::min() ? val > numeric_limits<T>::min() : T(val); 
		}

	return P2DImage(result); 
}

C2DImageFFTFilter::C2DImageFFTFilter()
{
}
	

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


C2DImageFFTFilterFactory::C2DImageFFTFilterFactory():
	C2DFilterPlugin(plugin_name)
{
}

C2DFilter *C2DImageFFTFilterFactory::create(const CParsedOptions& options) const
{
	return ProductPtr(new C2DImageFFTFilter()); 
}

const string C2DImageFFTFilterFactory::do_get_descr()const
{
	return "2D image convert-to-frequency domain filters based on FFT";
}


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

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



NSEND
