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

#define VSTREAM_DOMAIN "rigidreg"

#include <mia/2d/rigidregister.hh>
#include <mia/2d/filter.hh>
#include <mia/2d/transformfactory.hh>
#include <mia/core/filter.hh>


NS_MIA_BEGIN

using namespace std;

struct C2DRigidRegisterImpl {

       C2DRigidRegisterImpl(P2DImageCost cost, PMinimizer minimizer,
                            P2DTransformationFactory transform_creator,  size_t mg_levels);

       P2DTransformation run(P2DImage src, P2DImage ref) const;
private:

       P2DImageCost m_cost;
       PMinimizer m_minimizer;
       P2DTransformationFactory m_transform_creator;
       size_t m_mg_levels;
};

class C2DRegGradientProblem: public CMinimizer::Problem
{
public:
       C2DRegGradientProblem(const C2DImage& model, C2DTransformation& transf, const C2DImageCost& m_cost);
private:
       void    do_df(const CDoubleVector& x, CDoubleVector&  g);
       double  do_fdf(const CDoubleVector& x, CDoubleVector&  g);
protected:
       P2DImage apply(const CDoubleVector& x);
       double  do_f(const CDoubleVector& x);
       size_t do_size() const;

       const C2DImage& m_model;
       C2DTransformation& m_transf;
       const C2DImageCost& m_cost;
};
typedef shared_ptr<C2DRegGradientProblem> P2DGradientProblem;


class C2DRegFakeGradientProblem: public C2DRegGradientProblem
{
public:
       C2DRegFakeGradientProblem(const C2DImage& model, C2DTransformation& transf,
                                 const C2DImageCost& m_cost);
private:
       void    do_df(const CDoubleVector& x, CDoubleVector&  g);
       double  do_fdf(const CDoubleVector& x, CDoubleVector&  g);

};

class C2DRegProblem: public CMinimizer::Problem
{
public:
       C2DRegProblem(const C2DImage& model, C2DTransformation& transf,
                     const C2DImageCost& m_cost);
private:
       double  do_f(const CDoubleVector& x);
       void    do_df(const CDoubleVector& x, CDoubleVector&  g);
       double  do_fdf(const CDoubleVector& x, CDoubleVector&  g);
       size_t do_size() const;

       const C2DImage& m_model;
       C2DTransformation& m_transf;
       const C2DImageCost& m_cost;
};
typedef shared_ptr<C2DRegProblem> P2DRegProblem;


C2DRigidRegister::C2DRigidRegister(P2DImageCost cost, PMinimizer minimizer,
                                   P2DTransformationFactory transform_creator, size_t mg_levels):
       impl(new C2DRigidRegisterImpl( cost, minimizer, transform_creator, mg_levels))
{
}


C2DRigidRegister::~C2DRigidRegister()
{
       delete impl;
}

P2DTransformation C2DRigidRegister::run(P2DImage src, P2DImage ref) const
{
       return impl->run(src, ref);
}

C2DRigidRegisterImpl::C2DRigidRegisterImpl(P2DImageCost cost, PMinimizer minimizer,
              P2DTransformationFactory transform_creator,
              size_t mg_levels):
       m_cost(cost),
       m_minimizer(minimizer),
       m_transform_creator(transform_creator),
       m_mg_levels(mg_levels)
{
}

P2DTransformation C2DRigidRegisterImpl::run(P2DImage src, P2DImage ref) const
{
       assert(src);
       assert(ref);
       assert(src->get_size() == ref->get_size());
       P2DTransformation transform;
       int x_shift = m_mg_levels + 1;
       int y_shift = m_mg_levels + 1;

       while (x_shift && y_shift) {
              if (x_shift)
                     x_shift--;

              if (y_shift)
                     y_shift--;

              C2DBounds BlockSize(1 << x_shift, 1 << y_shift);
              cvinfo() << "Blocksize = " << BlockSize.x << "x" << BlockSize.y << "\n";
              stringstream downscale_descr;
              downscale_descr << "downscale:bx=" << BlockSize.x << ",by=" << BlockSize.y;
              auto downscaler = C2DFilterPluginHandler::instance().produce(downscale_descr.str().c_str());
              P2DImage src_scaled = x_shift && y_shift ? downscaler->filter(*src) : src;
              P2DImage ref_scaled = x_shift && y_shift ? downscaler->filter(*ref) : ref;
              m_cost->set_reference(*ref_scaled);

              if (transform)
                     transform = transform->upscale(src_scaled->get_size());
              else
                     transform = m_transform_creator->create(src_scaled->get_size());

              cvmsg() << "register at " << src_scaled->get_size() << "\n";
              CMinimizer::PProblem gp = m_minimizer->has(property_gradient) ?
                                        CMinimizer::PProblem(new C2DRegFakeGradientProblem(*src_scaled, *transform, *m_cost)) :
                                        CMinimizer::PProblem(new C2DRegProblem(*src_scaled, *transform, *m_cost));
              m_minimizer->set_problem(gp);
              auto x = transform->get_parameters();
              m_minimizer->run(x);
              transform->set_parameters(x);
              auto params = transform->get_parameters();
              cvinfo() << "\nParams:";

              for (auto i = params.begin(); i != params.end(); ++i)
                     cverb << *i << " ";

              cverb << "\n";
       }

       return transform;
}

C2DRegGradientProblem::C2DRegGradientProblem(const C2DImage& model, C2DTransformation& transf,
              const C2DImageCost& cost):
       m_model(model),
       m_transf(transf),
       m_cost(cost)
{
       add(property_gradient);
}

size_t C2DRegGradientProblem::do_size() const
{
       return m_transf.degrees_of_freedom();
}

P2DImage C2DRegGradientProblem::apply(const CDoubleVector& x)
{
       cvinfo() << "\nParams:";

       for (auto i = x.begin(); i != x.end(); ++i)
              cverb << *i << " ";

       cverb << "\n";
       m_transf.set_parameters(x);
       return m_transf(m_model);
}

double  C2DRegGradientProblem::do_f(const CDoubleVector& x)
{
       P2DImage temp = apply(x);
       const double value = m_cost.value(*temp);
       cvmsg() << "\rCost = " << value;
       return value;
}

void    C2DRegGradientProblem::do_df(const CDoubleVector& x, CDoubleVector&  g)
{
       P2DImage temp = apply(x);
       C2DFVectorfield gradient(m_model.get_size());
       m_cost.evaluate_force(*temp, gradient);
       transform(gradient.begin(), gradient.end(), gradient.begin(),
       [](const C2DFVector & x) {
              return -1.0f * x;
       });
       m_transf.translate(gradient, g);
}

double  C2DRegGradientProblem::do_fdf(const CDoubleVector& x, CDoubleVector&  g)
{
       P2DImage temp = apply(x);
       C2DFVectorfield gradient(m_model.get_size());
       double result = m_cost.evaluate_force(*temp, gradient);
       transform(gradient.begin(), gradient.end(), gradient.begin(),
       [](const C2DFVector & x) {
              return -1.0f * x;
       });
       m_transf.translate(gradient, g);
       return result;
}

C2DRegFakeGradientProblem::C2DRegFakeGradientProblem(const C2DImage& model,
              C2DTransformation& transf, const C2DImageCost& m_cost):
       C2DRegGradientProblem(model, transf,  m_cost)
{
       add(property_gradient);
}

void C2DRegFakeGradientProblem::do_df(const CDoubleVector& x, CDoubleVector&  g)
{
       CDoubleVector x_tmp(x.size());
       copy(x.begin(), x.end(), x_tmp.begin());

       for (size_t i = 0; i < g.size(); ++i) {
              x_tmp[i] += 0.01;
              double cost_p = do_f(x_tmp);
              x_tmp[i] -= 0.02;
              double cost_m = do_f(x_tmp);
              x_tmp[i] += 0.01;
              g[i] = (cost_p - cost_m) * 50.0;
              cvdebug() << "g[" << i << "] = " << g[i] << "\n";
       }
}

double  C2DRegFakeGradientProblem::do_fdf(const CDoubleVector& x, CDoubleVector&  g)
{
       double cost_value = do_f(x);
       do_df(x, g);
       return cost_value;
}

C2DRegProblem::C2DRegProblem(const C2DImage& model, C2DTransformation& transf,
                             const C2DImageCost& cost):
       m_model(model),
       m_transf(transf),
       m_cost(cost)
{
}

double  C2DRegProblem::do_f(const CDoubleVector& x)
{
       cvinfo() << "\nParams:";

       for (auto i = x.begin(); i != x.end(); ++i)
              cverb << *i << " ";

       cverb << "\n";
       m_transf.set_parameters(x);
       P2DImage test =  m_transf(m_model);
       const double value = m_cost.value(*test);
       cvmsg() << "Cost = " << value << "\n";
       return value;
}

void   C2DRegProblem::do_df(const CDoubleVector&, CDoubleVector&  )
{
       DEBUG_ASSERT_RELEASE_THROW(false, "C2DRegProblem doesn't provide a gradient.");
}

double  C2DRegProblem::do_fdf(const CDoubleVector&, CDoubleVector&  )
{
       DEBUG_ASSERT_RELEASE_THROW(false, "C2DRegProblem doesn't provide a gradient.");
}

size_t C2DRegProblem::do_size() const
{
       return m_transf.degrees_of_freedom();
}


NS_MIA_END
