/* -*- 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 <limits>
#include <mia/core/msgstream.hh>
#include <mia/3d/transform/vectorfield.hh>
//#include <mia/3d/vfio.hh>

NS_MIA_BEGIN
using namespace std;

C3DGridTransformation::C3DGridTransformation(const C3DBounds& size, const C3DInterpolatorFactory& ipf):
	C3DTransformation(ipf),
	m_field(size), 
	m_upscale_interpolator_factory("bspline:d=1", "zero")
{
}

C3DFVector C3DGridTransformation::get_displacement_at(const C3DFVector& x) const
{
	return m_field.get_interpol_val_at(x);
}

P3DTransformation C3DGridTransformation::do_upscale(const C3DBounds& size) const
{
	/* This implementation could be improved by using something like the spline interpolator 
	   that applies the scaling in a separable way. */

	TRACE("C3DGridTransformation::upscale");
	DEBUG_ASSERT_RELEASE_THROW(m_field.get_size().x != 0 && m_field.get_size().y != 0, 
				   "C3DGridTransformation::do_upscale: input field has a zero dimension"); 


	C3DGridTransformation *result = new C3DGridTransformation(size, get_interpolator_factory());

	unique_ptr<T3DInterpolator<C3DFVector> >
		interp(m_upscale_interpolator_factory.create(m_field)); 
	

	C3DFVector mult(float(size.x) / (float)m_field.get_size().x, 
			float(size.y) / (float)m_field.get_size().y,
			float(size.z) / (float)m_field.get_size().z); 
	
	C3DFVector imult = C3DFVector::_1 / mult; 
	
	auto i = result->m_field.begin_range(C3DBounds::_0, m_field.get_size());
	auto e = result->m_field.end_range(C3DBounds::_0, m_field.get_size());
	

	while (i != e)
		*i++ = (*interp)(imult * C3DFVector(i.pos())) * mult;
	
	return P3DTransformation(result);
}

const C3DBounds& C3DGridTransformation::get_size() const
{
	return m_field.get_size();
}

void C3DGridTransformation::update(float step, const C3DFVectorfield& a)
{
	C3DFVectorfield::const_iterator inf = a.begin();
	C3DFVectorfield::const_iterator enf = a.end();
	C3DFVectorfield::iterator onf = m_field.begin();

	while (inf != enf)
		*onf++ += step * *inf++;

}

C3DTransformation *C3DGridTransformation::do_clone() const
{
	return new C3DGridTransformation(*this);
}

C3DTransformation *C3DGridTransformation::invert() const
{
	assert(0 && "not implemented"); 
	return new C3DGridTransformation(*this);
}


size_t C3DGridTransformation::degrees_of_freedom() const
{
	return 3 * m_field.size();
}

C3DFMatrix C3DGridTransformation::derivative_at(const  C3DFVector& x) const
{
	
	C3DFMatrix result = C3DFMatrix::_1; 

	const int ix = static_cast<int>(floor(x.x)); 
	const int iy = static_cast<int>(floor(x.y)); 
	const int iz = static_cast<int>(floor(x.z)); 

	const int sx = m_field.get_size().x;
	const int sy = m_field.get_size().y;
	const int sz = m_field.get_size().z;
	
	if (ix > 0 && ix < sx - 2 && 
	    iy > 0 && iy < sy - 2 && 
	    iz > 0 && iz < sz - 2) {
		
		const float dx = x.x - ix; 
		const float dy = x.y - iy; 
		const float dz = x.z - iz; 
		
		const float fx = 1.0 - dx; 
		const float fy = 1.0 - dy; 
		const float fz = 1.0 - dz; 
		
		C3DFMatrix m[8]; 
		m[0] = field_derivative_at(ix    , iy    , iz    );
		m[1] = field_derivative_at(ix + 1, iy    , iz    );
		m[2] = field_derivative_at(ix    , iy + 1, iz    );
		m[3] = field_derivative_at(ix + 1, iy + 1, iz    );
		m[4] = field_derivative_at(ix    , iy    , iz + 1);
		m[5] = field_derivative_at(ix + 1, iy    , iz + 1);
		m[6] = field_derivative_at(ix    , iy + 1, iz + 1);
		m[7] = field_derivative_at(ix + 1, iy + 1, iz + 1);

		result -= 
			fz * (
				fy * (fx * m[0] + dx * m[1]) + 
				dy * (fx * m[2] + dx * m[3])) + 
			dz * (
				fy * (fx * m[4] + dx * m[5]) + 
				dy * (fx * m[6] + dx * m[7])); 
	}
	return result; 
}

C3DFMatrix C3DGridTransformation::derivative_at(int x, int y, int z) const
{
	const int sx = m_field.get_size().x;
	const int sy = m_field.get_size().y;
	const int sz = m_field.get_size().z;

	C3DFMatrix result = C3DFMatrix::_1; 
	if (z >= 0 && z < sz && 
	    y >= 0 && y < sy && 
	    x >= 0 && x < sx) {
		result -= field_derivative_at(x,y,z);
	}
	return result; 
}

C3DFMatrix C3DGridTransformation::field_derivative_at(int x, int y, int z) const
{
	C3DFMatrix result;
	
	const int sx = m_field.get_size().x;
	const int sy = m_field.get_size().y;
	const int sz = m_field.get_size().z;
	const int sxy = sx * sy; 

	C3DFVectorfield::const_pointer center = &m_field[(z * sy + y) * sx  + x];
	if (x > 0 && x < sx - 1) {
		result.x = (center[ 1] - center[-1]) * 0.5f;
	}
	if (y > 0 && y < sy - 1 ) {
		result.y = (center[ sx] - center[-sx]) * 0.5f;
	}
	if (z > 0 && z < sz - 1 ) {
		result.z = (center[ sxy] - center[-sxy]) * 0.5f;
	}
	return result;
}

void C3DGridTransformation::set_identity()
{
	fill(m_field.begin(),m_field.end(), C3DFVector::_0);
}

float C3DGridTransformation::get_max_transform() const
{
	C3DFVectorfield::const_iterator i = m_field.begin();
	C3DFVectorfield::const_iterator e = m_field.end();

	assert(i != e);
	float value = i->norm2();

	++i;
	while (i != e) {
		float v = i->norm2();
		if (value < v)
			value = v;
		++i;
	}
	return sqrt(value);
}


CDoubleVector C3DGridTransformation::get_parameters() const
{
	CDoubleVector result(m_field.size() * 3);
	auto r = result.begin();
	for(auto f = m_field.begin(); f != m_field.end(); ++f) {
		*r++ = f->x;
		*r++ = f->y;
		*r++ = f->z;
	}
	return result;
}

void C3DGridTransformation::set_parameters(const CDoubleVector& params)
{
	assert(3 * m_field.size() == params.size());
	auto r = params.begin();
	for(auto f = m_field.begin(); f != m_field.end(); ++f) {
		f->x = *r++;
		f->y = *r++;
		f->z = *r++;
	}
}


C3DGridTransformation::iterator_impl::iterator_impl(const C3DBounds& pos, const C3DBounds& size, 
						    C3DFVectorfield::const_iterator start):
	C3DTransformation::iterator_impl(pos, size), 
	m_current(start)
{
	if (pos.y < size.y && pos.x < size.x && pos.z < size.z ) 
		m_value = C3DFVector(get_pos()) - *m_current; 
}

C3DGridTransformation::iterator_impl::iterator_impl(const C3DBounds& pos, const C3DBounds& begin, const C3DBounds& end, 
						    const C3DBounds& size, C3DFVectorfield::const_iterator start):
C3DTransformation::iterator_impl(pos, begin, end, size), 
		    m_current(start)
{
	if (pos.y < size.y && pos.x < size.x && pos.z < size.z ) 
		m_value = C3DFVector(get_pos()) - *m_current; 
}
		    

C3DTransformation::const_iterator C3DGridTransformation::begin_range(const C3DBounds& begin, const C3DBounds& end) const
{
	TRACE_FUNCTION;
	return C3DTransformation::const_iterator(new iterator_impl(begin, begin, end, get_size(), 
								   m_field.begin_at(begin.x, begin.y, begin.z)));
}


C3DTransformation::const_iterator C3DGridTransformation::end_range(const C3DBounds& begin, const C3DBounds& end) const
{
	TRACE_FUNCTION;
	return C3DTransformation::const_iterator(new iterator_impl(end, begin, end, get_size(), 
								   m_field.begin_at(end.x, end.y, end.z)));

}


C3DTransformation::iterator_impl * C3DGridTransformation::iterator_impl::clone() const
{
	return new C3DGridTransformation::iterator_impl(get_pos(), get_size(), m_current); 
}

const C3DFVector&  C3DGridTransformation::iterator_impl::do_get_value()const
{
	return m_value; 
}

void C3DGridTransformation::iterator_impl::do_x_increment()
{
	++m_current; 
	m_value = C3DFVector(get_pos()) - *m_current; 
}

void C3DGridTransformation::iterator_impl::do_y_increment()
{
	++m_current; 
	m_value = C3DFVector(get_pos()) - *m_current; 
}

void C3DGridTransformation::iterator_impl::do_z_increment()
{
	++m_current; 
	m_value = C3DFVector(get_pos()) - *m_current; 
}


C3DFVector C3DGridTransformation::operator ()(const  C3DFVector& x) const
{
	return x - get_displacement_at(x);
}

C3DGridTransformation::const_iterator C3DGridTransformation::begin() const
{
	return const_iterator(new iterator_impl(C3DBounds::_0, m_field.get_size(), m_field.begin()));
}

C3DGridTransformation::const_iterator C3DGridTransformation::end() const
{
	return const_iterator(new iterator_impl( m_field.get_size(), m_field.get_size(), 
						 m_field.end()));
}


C3DGridTransformation::field_iterator C3DGridTransformation::field_begin()
{
	return m_field.begin();
}

C3DGridTransformation::field_iterator C3DGridTransformation::field_end()
{
	return m_field.end();
}

C3DGridTransformation::const_field_iterator C3DGridTransformation::field_begin() const
{
	return m_field.begin();
}

C3DGridTransformation::const_field_iterator C3DGridTransformation::field_end()const
{
	return m_field.end();
}

void C3DGridTransformation::translate(const C3DFVectorfield& gradient, CDoubleVector& params) const
{
	assert(3 * params.size() != gradient.size());

	// translating the gradient also means multiplication with -1 
	// because the transformation is I-u(x)
	auto r = params.begin();
	for(auto f = gradient.begin(); f != gradient.end(); ++f, r+=3) {
		r[0] = -f->x;
		r[1] = -f->y;
		r[2] = -f->z;
	}
}

float C3DGridTransformation::pertuberate(C3DFVectorfield& v) const
{
	C3DFVectorfield::iterator iv = v.begin();
	float max_gamma = 0.0f;
	for (size_t z = 0; z < v.get_size().z; ++z)
		for (size_t y = 0; y < v.get_size().y; ++y)
			for (size_t x = 0; x < v.get_size().x; ++x, ++iv){
				const C3DFMatrix j = field_derivative_at(x,y,z);
				const C3DFVector ue = j * *iv;
				*iv -= ue;
				float gamma = iv->norm2();
				if (gamma > max_gamma)
					max_gamma = gamma;
			}
	return sqrt(max_gamma);
}



float C3DGridTransformation::get_jacobian(const C3DFVectorfield& v, float delta) const
{
	assert(v.get_size() == get_size());
	assert(0 && "this function needs to be fixed"); 
	const int dx = v.get_size().x;
	const int dxy = v.get_size().x * v.get_size().y;
	
	float j_min = numeric_limits<float>::max();

	for(size_t z = 1; z < get_size().z - 1; ++z) {
		for(size_t y = 1; y < get_size().y - 1; ++y) {
			C3DFVectorfield::const_iterator iu = m_field.begin_at(1,y,z);
			C3DFVectorfield::const_iterator iv = v.begin_at(1,y,z);
			for(size_t x = 1; x < get_size().x - 1; ++x, ++iu, ++iv) {
				C3DFVector jx((iu[1] - iu[-1]) + (delta * (iv[1] - iv[-1])));
				C3DFVector jy((iu[dx] - iu[-dx]) + (delta * (iv[dx] - iv[-dx])));
				C3DFVector jz((iu[dxy] - iu[-dxy]) + (delta * (iv[dxy] - iv[-dxy])));
				
				const float j = (2.0 - jx.x) * (2.0 - jy.y) * (2.0 - jz.z) - jx.y * jy.x * jy.z;
				if ( j_min > j) {
					j_min = j;
				}
			}
		}
	}
	return 0.25 * j_min;
}

EXPORT_3D C3DGridTransformation operator + (const C3DGridTransformation& a, const C3DGridTransformation& b)
{
	assert( a.get_size() == b.get_size());

	C3DGridTransformation result(a.get_size(), a.get_interpolator_factory());

	C3DFVectorfield::iterator ri = result.field_begin();
	C3DFVectorfield::const_iterator bi = b.field_begin();
	
	
	for (size_t z = 0; z < a.get_size().z; ++z)  {
		for (size_t y = 0; y < a.get_size().y; ++y)  {
			for (size_t x = 0; x < a.get_size().x; ++x, ++ri, ++bi)  {
				const C3DFVector xi = C3DFVector(x,y,z) - *bi;
				*ri = a.get_displacement_at(xi) +  *bi;
			}
		}
	}
	return result;
}

/**
   Transformation creator 
 */

C3DGridTransformCreator::C3DGridTransformCreator(const C3DInterpolatorFactory& ipf):
	C3DTransformCreator(ipf)
{
}

P3DTransformation C3DGridTransformCreator::do_create(const C3DBounds& size, const C3DInterpolatorFactory& ipf) const
{
	return P3DTransformation(new C3DGridTransformation(size, ipf));
}


C3DGridTransformCreatorPlugin::C3DGridTransformCreatorPlugin():
	C3DTransformCreatorPlugin("vf")
{
}

C3DTransformCreator *C3DGridTransformCreatorPlugin::do_create(const C3DInterpolatorFactory& ipf) const
{
	return new C3DGridTransformCreator(ipf);
}

const std::string C3DGridTransformCreatorPlugin::do_get_descr() const
{
	return "This plug-in implements a transformation that defines a translation for "
		"each point of the grid defining the domain of the transformation.";
}


NS_MIA_END
