/* -*- 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/2d/model/navier.hh>

NS_MIA_USE
using namespace boost;
using namespace std;

NS_BEGIN(navier_regmodel)


C2DNavierRegModel::C2DNavierRegModel(float mu, float lambda, size_t maxiter,
                                     float omega, float epsilon):
       m_mu(mu),
       m_lambda(lambda),
       m_omega(omega),
       m_epsilon(epsilon),
       m_max_iter(maxiter)
{
       float a = mu;
       float b = lambda + mu;
       m_c = 1 / (4 * a + 2 * b);
       m_b_4 = 0.25 * b * m_c;
       m_a_b = ( a + b ) * m_c;
       m_a = a * m_c;
}


float C2DNavierRegModel::do_force_scale() const
{
       return m_c;
}

void C2DNavierRegModel::do_solve(const C2DFVectorfield& b, C2DFVectorfield& v) const
{
       // init velocity fields
       float start_residuum = 0.0;
       float residuum;
       size_t i = 0;

       do {
              residuum = 0;
              C2DFVectorfield::const_iterator ib = b.begin() + b.get_size().x;

              for (size_t y = 1; y < b.get_size().y - 1; y++) {
                     ++ib;

                     for (size_t x = 1; x < b.get_size().x - 1; x++, ++ib) {
                            residuum += solve_at(x, y, *ib, v);
                     }

                     ++ib;
              }

              if (i == 0)
                     start_residuum = residuum;

              if (residuum < 0.01)
                     break;

              ++i;
       } while (i < m_max_iter && residuum / start_residuum > m_epsilon);
}

float  C2DNavierRegModel::solve_at(unsigned int x, unsigned int y,
                                   const C2DFVector& b, C2DFVectorfield& v)const
{
       const int sizex = v.get_size().x;
       C2DFVector *vp = &v(x, y);
       C2DFVector *vpp = &vp[  sizex];
       C2DFVector *vpm = &vp[ -sizex];
       C2DFVector p(b.x + m_a_b * ( vp[-1].x + vp[+1].x ) +
                    m_a * ( vpp->x   + vpm->x  ),
                    b.y + m_a_b * ( vpm->y   + vpp->y   ) +
                    m_a * ( vp[-1].y + vp[1].y ));
       C2DFVector q;
       q.y = ( ( vpm[-1].x + vpp[1].x ) - ( vpm[1].x + vpp[-1].x ) ) * m_b_4;
       q.x = ( ( vpm[-1].y + vpp[1].y ) - ( vpm[1].y + vpp[-1].y ) ) * m_b_4;
       C2DFVector hmm((( p + q ) - *vp) * m_omega);
       *vp += hmm;
       return hmm.norm();
}

bool C2DNavierRegModel::test_kernel()
{
       return true;
}

class C2DNavierRegModelPlugin: public C2DRegModelPlugin
{
public:
       C2DNavierRegModelPlugin();
       C2DRegModel *do_create()const;

private:
       const string do_get_descr()const;

       float m_mu;
       float m_lambda;
       float m_omega;
       float m_epsilon;
       int m_maxiter;
};

C2DNavierRegModelPlugin::C2DNavierRegModelPlugin():
       C2DRegModelPlugin("navier"),
       m_mu(1.0),
       m_lambda(1.0),
       m_omega(1.0),
       m_epsilon(0.0001),
       m_maxiter(100)
{
       add_parameter("mu", make_lc_param(m_mu, 0.0, false, "isotropic compliance"));
       add_parameter("lambda", make_lc_param(m_lambda, 0.0, false, "isotropic compression"));
       add_parameter("omega", make_ci_param(m_omega, 0.1, 10, false, "relexation parameter"));
       add_parameter("epsilon", make_oci_param(m_epsilon, 0.0, 0.1, false, "stopping parameter"));
       add_parameter("iter", make_ci_param(m_maxiter, 1, 10000, false, "maximum number of iterations"));
}

C2DRegModel *C2DNavierRegModelPlugin::do_create()const
{
       return new C2DNavierRegModel(m_mu, m_lambda,
                                    m_maxiter, m_omega, m_epsilon);
}

const string C2DNavierRegModelPlugin::do_get_descr()const
{
       return "navier-stokes based registration model";
}

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

NS_END

