/* -*- 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 <mia/internal/plugintester.hh>
#include <mia/3d/vfregularizer/sor.hh>
#include <mia/3d/vfregularizerkernel.hh>
#include <mia/core/parallel.hh>

using namespace std; 
using namespace mia; 


class Mock3DVectorfieldRegularizerKernel: public C3DFVectorfieldRegularizerKernel {
public:
	Mock3DVectorfieldRegularizerKernel(); 
private: 
        virtual float do_evaluate_row(unsigned y, unsigned z, CBuffers& buffers); 
        virtual float do_evaluate_row_sparse(unsigned y, unsigned z, CBuffers& buffers); 
        virtual unsigned do_get_boundary_padding() const; 
        float solve_at(C3DFVector *v, const C3DFVector& b); 
	
	void post_set_data_fields(); 
	int m_dx; 
	int m_dxy;
	vector<CMutex> m_slice_mutex;

	

}; 

BOOST_AUTO_TEST_CASE(test_sor_solver) 
{
	C3DBounds size(5, 6, 7); 
	C3DBounds rbe = size - 2 * C3DBounds::_1;

	C3DFVectorfield v(size); 
	C3DFVectorfield b(size); 


	auto useful_range = v.get_range(2 * C3DBounds::_1, rbe);
	auto irv = useful_range.begin(); 
	auto erv = useful_range.end(); 
	
	while ( irv != erv ) {
		C3DFVector x(irv.pos());
		
		float fz = x.z - 3.5; 
		fz = 1.0 / (1 + fz * fz); 
		float fy = fz * cos((x.y - 3.0)/6.0 * M_PI);
		float fx = x.x - 2.5f;

		*irv = C3DFVector( fy *  1.0 / (1 + fx * fx), 
				   fy * fy *  1.0 / (1 + fx * fx), 
				   fz * fy *  2.0 / (1 + fx * fx)); 

		cvdebug() << "input[" << irv.pos() << "]=" << *irv << "\n"; 
		
		++irv; 
	}

	
	C3DBounds rbe1 = size - C3DBounds::_1;
	auto b_range = b.get_range(C3DBounds::_1, rbe1);
	
	auto irb = b_range.begin(); 
	auto erb = b_range.end(); 
	
	while (irb != erb) {
		C3DBounds p = irb.pos(); 
		
		*irb = (6.0f * v(p) + 
			v(C3DBounds(p.x + 1, p.y, p.z)) + 
			v(C3DBounds(p.x - 1, p.y, p.z)) + 
			v(C3DBounds(p.x, p.y + 1, p.z)) + 
			v(C3DBounds(p.x, p.y - 1, p.z)) + 
			v(C3DBounds(p.x, p.y, p.z + 1)) + 
			v(C3DBounds(p.x, p.y, p.z - 1))) * 1.0f/12.0f; 
		
		++irb; 
	}
	
	C3DSORVectorfieldRegularizer sor_solver(1e-7, 1e-9, 1000, 
						make_shared<Mock3DVectorfieldRegularizerKernel>()); 
	C3DFVectorfield vs(size); 
	C3DFVectorfield d(size); 

	sor_solver.set_size(size); 
	sor_solver.run(vs, b, d); 
	
	auto vs_full_range = vs.get_range(C3DBounds::_0, vs.get_size()); 
	auto v_full_range = v.get_range(C3DBounds::_0, vs.get_size()); 
	
	for (auto ivs = vs_full_range.begin(), iv = v_full_range.begin(); ivs != vs_full_range.end(); ++ivs, ++iv) {
		
		cvdebug() << "[" << ivs.pos() << "]=" << *ivs << " expect " << *iv << "\n"; 

		if (fabs(iv->x) > 1e-5) 
			BOOST_CHECK_CLOSE(ivs->x, iv->x, 1e-4);
		else 
			BOOST_CHECK_SMALL(ivs->x, 1e-5f);

		if (fabs(iv->y) > 1e-5) 
			BOOST_CHECK_CLOSE(ivs->y, iv->y, 1e-4);
		else 
			BOOST_CHECK_SMALL(ivs->y, 1e-5f);

		if (fabs(iv->z) > 1e-5) 
			BOOST_CHECK_CLOSE(ivs->z, iv->z, 1e-4);
		else 
			BOOST_CHECK_SMALL(ivs->z, 1e-5f);


	}
	
}

Mock3DVectorfieldRegularizerKernel::Mock3DVectorfieldRegularizerKernel():
	C3DFVectorfieldRegularizerKernel(false),
	m_dx(0),
	m_dxy(0)
{
}


float Mock3DVectorfieldRegularizerKernel::do_evaluate_row(unsigned y, unsigned z, 
                                                          CBuffers& MIA_PARAM_UNUSED(buffers))
{
        float residuum = 0.0f; 
        
        auto& v = get_output_field(); 
        auto& b = get_input_field(); 

	unsigned linear_index = 1 + y * m_dx + z * m_dxy; 

	// There should be a more fine-grained method, but for the test it's okay
	
	CScopedLock lock_zm1(m_slice_mutex[z-1]);
	CScopedLock lock_z(m_slice_mutex[z]);
	CScopedLock lock_zp1(m_slice_mutex[z+1]); 
	
        for (int x = 1; x < m_dx - 1; ++x, ++linear_index) {
                residuum += solve_at(&v[linear_index], b[linear_index]); 
        }
        return residuum; 

}

float Mock3DVectorfieldRegularizerKernel::do_evaluate_row_sparse(unsigned MIA_PARAM_UNUSED(y),
								 unsigned MIA_PARAM_UNUSED(z), 
                                                                 CBuffers& MIA_PARAM_UNUSED(buffers))
{
        return numeric_limits<float>::max(); 
}

float Mock3DVectorfieldRegularizerKernel::solve_at(C3DFVector *v, const C3DFVector& b)
{
        C3DFVector s = 6.0f * v[0]; 
	s += v[1]; 
	s += v[-1]; 

	s += v[m_dx]; 
	s += v[m_dxy]; 
	
	s += v[-m_dx]; 

	s += v[-m_dxy]; 
		
	
	s /= 12.0f; 
        
        C3DFVector r = 1.2 * ( b - s); 

        v[0] += r; 
        
        return r.norm(); 
}

unsigned Mock3DVectorfieldRegularizerKernel::do_get_boundary_padding() const
{
        return 1; 
}

void Mock3DVectorfieldRegularizerKernel::post_set_data_fields()
{
	auto s = get_input_field().get_size();
	m_dx = s.x; 
	m_dxy = s.x * s.y;
	m_slice_mutex = vector<CMutex>(s.z); 
}
