/* -*- 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/core/cmdlineparser.hh>
#include <mia/3d/transformio.hh>
#include <mia/3d/transformfactory.hh>
#include <mia/3d/vfio.hh>
#include <mia/3d/cost.hh>
#include <mia/3d/imageio.hh>
#include <mia/internal/main.hh>


NS_MIA_USE
using namespace std;

const SProgramDescription g_description = {
       {pdi_group, "Registration, Comparison, and Transformation of 3D images"},
       {pdi_short, "Evaluate the cost gradient between two images and convert it to a spline representation."},
       {
              pdi_description, "Evaluate the cost gradient between two images and evaluate the "
              "transformation related gradient for it based on the given transformation model."
       },
       {
              pdi_example_descr, "Evaluate the SSD cost between src.v and ref.v and store the "
              "gradient gradient corresponding to a spline transformation in grad.v3dt."
       },
       {pdi_example_code, "-i src.v -o grad.v3dt -f spline:rate=8 -c ssd"}
};

int do_main(int argc, char **argv)
{
       CCmdOptionList options(g_description);
       string src_filename;
       string ref_filename;
       string out_filename;
       string grad_image_filename;
       string cost_grad_filename;
       P3DImageCost cost;
       P3DTransformationFactory transform_creator;
       const auto& imageio  = C3DImageIOPluginHandler::instance();
       options.add(make_opt( src_filename, "in-file", 'i', "input image ", CCmdOptionFlags::required_input, &imageio));
       options.add(make_opt( ref_filename, "ref-file", 'r', "reference image ", CCmdOptionFlags::required_input, &imageio));
       options.add(make_opt( out_filename, "out-file", 'o', "output vector field ",
                             CCmdOptionFlags::required_output, &C3DTransformationIOPluginHandler::instance()));
       options.add(make_opt( grad_image_filename, "gradimg-file", 'g', "norm image of the spline transformed gradient",
                             CCmdOptionFlags::output, &imageio));
       options.add(make_opt( cost_grad_filename, "cost-gradimg-file", 'C', "norm image of the cost gradient",
                             CCmdOptionFlags::output, &imageio));
       options.add(make_opt( transform_creator, "spline:rate=5", "transForm", 'f', "Transformation the gradient relates to"));
       options.add(make_opt( cost, "ssd", "cost", 'c', "cost function to use"));

       if (options.parse(argc, argv) != CCmdOptionList::hr_no)
              return EXIT_SUCCESS;

       auto source = load_image<P3DImage>(src_filename);
       auto ref    = load_image<P3DImage>(ref_filename);

       if (!source) {
              throw create_exception<runtime_error>("No image found in '", src_filename, "'");
       }

       if (!ref)
              throw create_exception<runtime_error>("No image found in '", src_filename, "'");

       C3DFVectorfield forcefield(source->get_size());
       cost->set_reference(*ref);
       cost->evaluate_force(*source, forcefield);
       auto t = transform_creator->create(forcefield.get_size());
       auto grad = t->get_parameters();
       t->translate(forcefield, grad);
       t->set_parameters(grad);

       if (!C3DTransformationIOPluginHandler::instance().save(out_filename, *t))
              throw create_exception<runtime_error>("Grad can not be saved to  '", out_filename, "'");

       if (!cost_grad_filename.empty()) {
              C3DFImage image(forcefield.get_size());
              float maxnorm = 0.0;
              transform(forcefield.begin(), forcefield.end(), image.begin(),
              [&maxnorm](const C3DFVector & x)->float{
                     float n = x.norm();

                     if (maxnorm < n)
                            maxnorm = n;
                     return n;
              });
//		float imn = 1.0/ maxnorm;
//		transform(image.begin(), image.end(), image.begin(), [imn](float x) {return imn * x;});
              save_image(cost_grad_filename, image);
              cvmsg() << "max gradient norm before translation = " << maxnorm << "\n";
       }

       if (!grad_image_filename.empty()) {
              C3DFImage image(forcefield.get_size());
              float maxnorm = 0.0;
              auto ti = t->begin();
              auto te = t->end();
              auto ii = image.begin_range(C3DBounds::_0, forcefield.get_size());

              while ( ti != te )  {
                     const C3DFVector d = *ti - C3DFVector(ii.pos());
                     *ii = d.norm();

                     if (maxnorm < *ii)
                            maxnorm = *ii;

                     ++ii;
                     ++ti;
              }

//		float imn = 1.0/ maxnorm;
//		transform(image.begin(), image.end(), image.begin(), [imn](float x) {return imn * x;});
              save_image(grad_image_filename, image);
              cvmsg() << "max gradient norm after translation = " << maxnorm << "\n";
       }

       return EXIT_SUCCESS;
}

MIA_MAIN(do_main);
