/* -*- 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 <iomanip>

#include "vfluid.hh"
#include "elast.hh"
#include <memory>
#include <mia/core/msgstream.hh>
#include <mia/2d/deformer.hh>


NS_MIA_BEGIN
using namespace std;

TElastReg::TElastReg(const C2DFImage& _Ref, const C2DFImage& _Model, float __regrid_tresh,
                     int level, float __epsilon, float __lambda, float __mu):
       ut(_Ref.get_size()),
       B(_Ref.get_size()),
       u(_Ref.get_size()),
       r(_Ref.get_size()),
       Ref(_Ref),
       Model(_Model),
       min_stepsize(MIN_FORCE_MULT / level),
       final_level(level == 1),
       epsilon(__epsilon),
       lambda(__lambda),
       mu(__mu),
       regrid_thresh(__regrid_tresh)
{
       stepSize = .00001;
       Start = C2DBounds(1, 1);
       End = Ref.get_size() - Start;
       // note: grid is assumed to be unity in all directions
       a = mu;
       b = lambda + mu;
       c = 1 / (4 * a + 2 * b);
       b_4 = 0.25 * b * c;
       omega = 1.0;
       a_b = (a + b) * c;
       a_ = a * c;
       delta = 0;
}
TElastReg::~TElastReg()
{
}

float  TElastReg::solveAt(int x, int y)
{
       const C2DFVector& B_ = B(x, y);
       C2DFVector *vp = &u(x, y);
       C2DFVector *vpp = &vp[ u.get_size().x];
       C2DFVector *vpm = &vp[-u.get_size().x];
       C2DFVector p(B_.x + (a_b) * ( vp[-1].x + vp[+1].x ) + a_ * ( vpp->x + vpm->x ),
                    B_.y + (a_b) * ( vpm->y + vpp->y ) + a_ * ( vp[-1].y + vp[1].y ));
       C2DFVector q;
       q.y = ( ( vpm[-1].x + vpp[1].x ) - ( vpm[1].x + vpp[-1].x ) ) * b_4;
       q.x = ( ( vpm[-1].y + vpp[1].y ) - ( vpm[1].y + vpp[-1].y ) ) * b_4;
       C2DFVector hmm((( p + q ) - *vp) * omega);
       *vp += hmm;
       return hmm.norm();
}

void	TElastReg::solvePDE(unsigned int nit)
{
       unsigned int i, x, y;
       double start_residuum = 0;
       double residuum;

       // init velocity fields
       //v.clear();

       for (y = Start.y + 1; y < End.y - 1; y++)
              for (x = Start.x + 1; x < End.x - 1; x++) {
                     start_residuum += solveAt(x, y);
              }

       residuum = start_residuum;

       for (i = 0; i < nit && residuum > epsilon; i++)  {
              residuum = 0;

              for (y = Start.y + 1; y < End.y - 1; y++)
                     for (x = Start.x + 1; x < End.x - 1; x++) {
                            residuum += solveAt(x, y);
                     }
       }

       cvmsg() << " res = " <<  residuum;
}


C2DFVector  TElastReg::forceAt(const C2DFVector& p, float s)
{
       float d = (*target_interp)(p) - s;
       C2DFVector gradt =  target_interp->derivative_at(p);
       return gradt * d;
}

void	TElastReg::calculateForces()
{
       float sum_force = 0.0;

       for (unsigned int y = Start.y; y < End.y; y++)  {
              for (unsigned int x = Start.x; x < End.x; x++)  {
                     // compute force at (xi, yi):
                     C2DFVector f  = stepSize * forceAt(C2DFVector(x, y) - u(x, y), Ref(x, y)) * c;
                     sum_force += f.norm();
                     B(x, y) = f;
              }
       }

       cvdebug() << "sum_force = " << sqrt(sum_force) << "\n";
}


float  TElastReg::calculateStep()
{
       float gamma = 0;

       for (unsigned int y = Start.y; y < End.y; y++)  {
              for (unsigned int x = Start.x; x < End.x; x++)  {
                     float g = ut(x, y).norm();

                     if (g > gamma) gamma = g;
              }
       }

       if (gamma > 0) {
              return stepSize / gamma;
       } else
              return 0;
}


float  TElastReg::calculateMismatch(bool apply)
{
       //	int n = (End.y  - Start.y) * (End.x - Start.x);
       float diff = 0;

       if (apply) {
              for (unsigned int y = Start.y; y < End.y; y++)  {
                     for (unsigned int x = Start.x; x < End.x; x++)  {
                            C2DFVector xi = C2DFVector(x, y) - u(x, y);
                            float p = (*target_interp)(xi);
                            diff += fabs(p - Ref(x, y));
                     }
              }
       } else {
              C2DFImage::const_iterator ra = Ref.begin();
              C2DFImage::const_iterator re = Ref.end();
              C2DFImage::const_iterator t = Template.begin();

              do {
                     diff += fabs(*ra - *t);
                     ++t;
                     ++ra;
              } while (ra != re);
       }

       return diff ;
}

void TElastReg::work(C2DFVectorfield *Shift, const C2DInterpolatorFactory& ipfac)
{
       int trys = 0;
       C2DFVectorfield usave(*Shift);
       // init displacement field
       u = *Shift;
       Template = C2DFImage(Model.get_size());
       {
              FDeformer2D deformer(*Shift, ipfac);
              deformer(Model, Template);
       }
       target_interp.reset(ipfac.create(Template.data()));
       double dmin = calculateMismatch(false);
       double diff = dmin;
       int maxiter = (final_level) ? 500 : 20;

       for (int i = 0; i < maxiter && trys < 5 && dmin > 0; i++)  {
              cvmsg() << "[" << setw(3) << "]:" << setw(10) << dmin << ", " << setw(10) << stepSize  << "\n";
              // Algorithm 2, step 2+3:
              calculateForces();
              //calculateForces();
              // Algorithm 2, step 4:
              solvePDE(1000);
              // Algorithm 2, step 5+6:
              // Algorithm 2, step 7:
              // apply weighted perturbation to displacement field
              // see algorithm 2, step 7, second part
              // calculate difference between source and template
              diff = calculateMismatch(true);

              if (diff < dmin) {
#if 0

                     if (diff < dmin * 0.9)
                            increaseStep();

#endif
                     trys = 0;
                     float delta = dmin - diff;
                     cvdebug() << " delta cost " <<  delta << "\n";
                     usave = u;

                     if (delta < 100)
                            break;

                     dmin = diff;
              } else {
#if 0

                     if (!decreaseStep()) {
                            trys = 5;
                            fprintf(stderr, "try %d", trys);
                     }

#endif
                     trys++;
                     cvdebug() << " delta cost " <<  delta << "\n";
              }

              if (dmin < 1000) break;

              cvmsg() << "\n";
       }

       cvmsg() << "final diff = " << dmin << "\n";
       *Shift = usave;
}


NS_MIA_END
