/* -*- 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/naviera.hh>

NS_BEGIN(naviera_regmodel)
NS_MIA_USE
using namespace boost;
using namespace std;


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;
       float c = 1 / (4 * a + 2 * b);
       m_b_4 = 0.25 * b * c;
       m_a_b = ( a + b ) * c;
       m_a = a * c;
}

void set_update(C2DUBImage::iterator isetupt, size_t dx)
{
       isetupt[-1 - dx] = true;
       isetupt[0 - dx] = true;
       isetupt[1 - dx] = true;
       isetupt[-1] = true;
       *isetupt = true;
       isetupt[1] = true;
       isetupt[dx - 1] = true;
       isetupt[dx] = true;
       isetupt[dx + 1] = true;
}

void C2DNavierRegModel::do_solve(const C2DFVectorfield& b, C2DFVectorfield& v) const
{
       // init velocity fields
       v.clear();
       float start_residuum = 0.0;
       float residuum;
       const size_t end_x = b.get_size().x - 1;
       const size_t dx = b.get_size().x;
       size_t i = 0;
       C2DFImage residua(b.get_size());
       C2DUBImage needupdate1(b.get_size());
       C2DUBImage needupdate2(b.get_size());
       C2DUBImage *needupdate_get = &needupdate1;
       C2DUBImage *needupdate_set = &needupdate2;
       fill(needupdate_get->begin(), needupdate_get->end(), true);
       fill(residua.begin(), residua.end(), 0.0f);
       {
              C2DFVectorfield::const_iterator ib = b.begin() + end_x + 1;
              C2DFVectorfield::iterator iv = v.begin() + dx;
              C2DFImage::iterator ires = residua.begin() + dx;

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

                     for (size_t x = 1; x < end_x; x++, ++ib, ++iv, ++ires) {
                            *ires = solve_at(x, y, *ib, iv, dx);
                            start_residuum += *ires;
                     }

                     ++ib;
                     ++ires;
              }
       }
       residuum = start_residuum;

       do {
              fill(needupdate_set->begin(), needupdate_set->end(), false);
              float rthresh = residuum / (b.size() + i);
              ++i;
              residuum = 0;
              C2DFVectorfield::iterator iv = v.begin() + dx;
              C2DFVectorfield::const_iterator ib = b.begin() + dx;
              C2DFImage::iterator ires = residua.begin() + dx;
              C2DUBImage::const_iterator igetupt = needupdate_get->begin()  + dx;
              C2DUBImage::iterator isetupt = needupdate_set->begin()  + dx;

// don't use get_size (or inline it)
// don't use a bit image for the update, since it needs shifts
              for (size_t y = 1; y < b.get_size().y - 1; y++) {
                     ++ib;
                     ++ires;
                     ++igetupt;
                     ++isetupt;
                     ++iv;

                     for (size_t x = 1; x < end_x; x++, ++iv, ++ib, ++ires, ++igetupt, ++isetupt) {
                            if (*igetupt)
                                   *ires = solve_at(x, y, *ib, iv, dx);

                            if (*ires > rthresh)
                                   set_update(isetupt, dx);

                            residuum += *ires;
                     }

                     ++ib;
                     ++ires;
                     ++igetupt;
                     ++isetupt;
                     ++iv;
              }

              cvdebug() << i << ":" << residuum << "\r";

              if (residuum < 1)
                     break;

              swap(needupdate_get, needupdate_set);
       } while (i < m_max_iter && residuum / start_residuum > m_epsilon);

       cverb << "\n";
}

float  C2DNavierRegModel::solve_at(unsigned int /*x*/, unsigned int /*y*/,
                                   const C2DFVector& b, C2DFVectorfield::iterator vp, int dx)const
{
       C2DFVectorfield::iterator vpp = vp + dx;
       C2DFVectorfield::iterator vpm = vp - dx;
       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();
}

C2DNavierRegModelPlugin::C2DNavierRegModelPlugin():
       C2DRegModelPlugin("naviera"),
       m_mu(1.0),
       m_lambda(1.0),
       m_omega(1.0),
       m_epsilon(0.0001),
       m_maxiter(100)
{
       add_parameter("mu", make_nonnegative_param(m_mu, false, "isotropic compliance"));
       add_parameter("lambda", make_nonnegative_param(m_lambda, 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 using a Gauss-Southwell relaxation as solver";
}

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

NS_END

