/* -*- 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 <fstream>
#include <sstream>
#include <iomanip>
#include <queue>
#include <boost/filesystem.hpp>

#include <mia/core/msgstream.hh>
#include <mia/core/fullstats.hh>
#include <mia/core/cmdlineparser.hh>
#include <mia/core/errormacro.hh>
#include <mia/core/filetools.hh>
#include <mia/3d/imageio.hh>
#include <mia/3d/nonrigidregister.hh>
#include <mia/3d/transformfactory.hh>
#include <mia/3d/fullcost.hh>
#include <mia/3d/similarity_profile.hh>

using namespace std;
using namespace mia;

namespace bfs = boost::filesystem;

const SProgramDescription g_description = {
       {pdi_group, "Registration of series of 3D images"},
       {pdi_short, "Registration of a series of 3D images."},
       {
              pdi_description, "This program runs the non-rigid registration of an image series "
              "by first registering an already aligned subset of the images to one reference, "
              "and then by registering the remaining images by using synthetic references. "
              "This is a 3D version of [G. Wollny, M-J Ledesma-Cabryo, P.Kellman, and A.Santos, "
              "\"Exploiting Quasiperiodicity in Motion Correction of Free-Breathing,\" "
              "IEEE Transactions on Medical Imaging, 29(8), 2010.]"
              "(http://dx.doi.org/10.1109/TMI.2010.2049270)"
       },
       {
              pdi_example_descr, "Register the image series given by images imageXXXX.v by "
              "optimizing a spline based transformation with a coefficient rate of 16 pixel, "
              "skipping two ""images at the beginning and using normalized gradient fields "
              "as initial cost measure and SSD as final measure. Penalize the transformation "
              "by using divcurl with aweight of 2.0. "
              "As optimizer an nlopt based newton method is used."
       },
       {
              pdi_example_code, "mia-3dprealign-nonrigid  -i imageXXXX.v -o registered -t vista -k 2"
              "-F spline:rate=16,penalty=[divcurl:weight=2] -1 image:cost=[ngf:eval=ds] -2 image:cost=ssd "
              "-O nlopt:opt=ld-var1,xtola=0.001,ftolr=0.001,maxiter=300"
       }
};

class C3DFImage2PImage
{
public:
       P3DImage operator () (const C3DFImage& image) const
       {
              return P3DImage(new C3DFImage(image));
       }
};


struct FAddWeighted: public TFilter<P3DImage> {
       FAddWeighted(float w):
              m_w(w)
       {
       }

       template <typename T, typename S>
       P3DImage operator() (const T3DImage<T>& a, const T3DImage<S>& b) const
       {
              if (a.get_size() != b.get_size()) {
                     throw invalid_argument("input images cann not be "
                                            "combined because they differ in size");
              }

              T3DImage<T> *result = new T3DImage<T>(a.get_size(), a);
              auto r = result->begin();
              auto e = result->end();
              auto ia = a.begin();
              auto ib = b.begin();
              float w2 = 1.0 - m_w;

              // this should be properly clamped
              while ( r != e ) {
                     *r = (T)(w2 * (float) * ia + m_w * (float) * ib);
                     ++r;
                     ++ia;
                     ++ib;
              }

              return P3DImage(result);
       }

private:
       float m_w;
};



class C3DMyocardPeriodicRegistration
{
public:
       class RegistrationParams
       {
       public:
              RegistrationParams();
              PMinimizer minimizer;
              P3DFullCost pass1_cost;
              P3DFullCost pass2_cost;
              P3DFullCost series_select_cost;
              P3DTransformationFactory transform_creator;
              size_t mg_levels;
              size_t max_candidates;
              bool save_ref;
              int global_reference;
              size_t max_delta;
       };

       C3DMyocardPeriodicRegistration(const RegistrationParams& params);
       vector<P3DTransformation> run(C3DImageSeries& images, size_t preskip, size_t postskip);
       size_t get_ref_idx()const;

private:
       vector<size_t>  get_high_contrast_candidates(const C3DImageSeries& images,
                     size_t startidx, size_t endidx);
       vector<size_t> get_prealigned_subset(const C3DImageSeries& images,
                                            size_t preskip, size_t postskip);
       void run_initial_pass(C3DImageSeries& images,
                             vector<P3DTransformation>& transforms,
                             const vector<size_t>& subset);
       void run_final_pass(C3DImageSeries& images,
                           vector<P3DTransformation>& transforms,
                           const vector<size_t>& subset);

       RegistrationParams m_params;
       size_t m_ref;
};

struct IdxVariation {
       size_t index;
       double variation;
};

bool operator < (const IdxVariation& a, const IdxVariation& b)
{
       return a.variation < b.variation;
}

class CStatsEvaluator : public TFilter<CFullStats>
{
public:
       template <typename T>
       CFullStats operator () (const T3DImage<T>& image) const
       {
              return CFullStats(image.begin(), image.end());
       }
};

vector<size_t>
C3DMyocardPeriodicRegistration::get_high_contrast_candidates(const C3DImageSeries& images,
              size_t startidx, size_t endidx)
{
       CStatsEvaluator sev;
       priority_queue<IdxVariation> q;

       while (startidx < endidx)  {
              IdxVariation v;
              v.index = startidx;
              v.variation = mia::filter(sev, *images[startidx]).sigma();
              q.push(v);
              ++startidx;
       }

       vector<size_t> result;

       while (result.size() < m_params.max_candidates && !q.empty()) {
              IdxVariation v = q.top();
              result.push_back(v.index);
              q.pop();
       }

       return result;
}

vector<size_t> C3DMyocardPeriodicRegistration::get_prealigned_subset(const C3DImageSeries& images,
              size_t preskip, size_t postskip)
{
       assert(postskip < images.size());
       assert(preskip < images.size() - postskip);
       vector<size_t> candidates;

       if (m_params.global_reference < 0) {
              cvmsg() << "estimate prealigned subset ...\n";
              candidates = get_high_contrast_candidates(images, preskip, images.size() - 2);
              assert(!candidates.empty());
       } else
              candidates.push_back(m_params.global_reference);

       C3DSimilarityProfile best_series(m_params.series_select_cost, images, candidates[0], m_params.max_delta);
       m_ref = candidates[0];

       // the skip values should be parameters
       for (size_t i = 1; i < candidates.size(); ++i) {
              C3DSimilarityProfile sp(m_params.series_select_cost, images, candidates[i], m_params.max_delta);

              if (sp.get_peak_frequency() > best_series.get_peak_frequency()) {
                     m_ref = candidates[i];
                     best_series = sp;
              }
       }

       return best_series.get_periodic_subset();
}


void C3DMyocardPeriodicRegistration::run_initial_pass(C3DImageSeries& images,
              vector<P3DTransformation>& transforms,
              const vector<size_t>& subset)
{
       cvmsg() << "run initial registration pass on " << subset << "\n";
       C3DFullCostList costs;
       costs.push(m_params.pass1_cost);
       C3DNonrigidRegister nr(costs,
                              m_params.minimizer,
                              m_params.transform_creator,
                              m_params.mg_levels);
       P3DImage ref = images[m_ref];

       for (auto i = subset.begin(); i != subset.end(); ++i) {
              if (m_ref == *i)
                     continue;

              cvmsg() << "Register " << *i << " to " << m_ref << "\n";
              P3DImage src = images[*i];
              P3DTransformation transform = nr.run(src, ref);
              images[*i] = (*transform)(*images[*i]);
              transforms[*i] = transform;
       }
}

void C3DMyocardPeriodicRegistration::run_final_pass(C3DImageSeries& images,
              vector<P3DTransformation>& transforms,
              const vector<size_t>& subset)
{
       assert(!subset.empty());
       cvmsg() << "run final registration pass ...\n";
       C3DFullCostList costs;
       costs.push(m_params.pass2_cost);
       C3DNonrigidRegister nr(costs,
                              m_params.minimizer,
                              m_params.transform_creator,
                              m_params.mg_levels);
       auto low_index = subset.begin();
       auto high_index = low_index + 1;

       for (size_t i = 0; i < images.size(); ++i) {
              if (i == *low_index)
                     continue;

              // the last images may be registered using SSD without interpolating references
              P3DImage ref;

              if (high_index != subset.end()) {
                     if (i == *high_index) {
                            ++low_index;
                            ++high_index;
                            continue;
                     }

                     float w = float(*high_index - i) / (*high_index - *low_index);
                     FAddWeighted lerp(w);
                     ref = mia::filter(lerp, *images[*high_index], *images[*low_index]);

                     if (m_params.save_ref) {
                            stringstream refname;
                            refname << "ref" << setw(4) << setfill('0') << i << ".v";
                            save_image(refname.str(), ref);
                            cvmsg() << "Save reference to " << refname.str() << "\n";
                     }
              } else
                     ref = images[*low_index];

              cvmsg() << "Register image " << i << "\n";
              P3DTransformation transform = nr.run(images[i], ref);
              images[i] = (*transform)(*images[i]);
              transforms[i] = transform;
       }
}

C3DMyocardPeriodicRegistration::C3DMyocardPeriodicRegistration(const RegistrationParams& params):
       m_params(params),
       m_ref(0)
{
}

vector<P3DTransformation> C3DMyocardPeriodicRegistration::run(C3DImageSeries& images, size_t preskip, size_t postskip)
{
       vector<P3DTransformation> transforms(images.size());
       vector<size_t> subset = get_prealigned_subset(images, preskip, postskip);
       run_initial_pass(images, transforms, subset);
       run_final_pass(images, transforms, subset);
       return transforms;
}

size_t C3DMyocardPeriodicRegistration::get_ref_idx()const
{
       return m_ref;
}


C3DMyocardPeriodicRegistration::RegistrationParams::RegistrationParams():
       mg_levels(3),
       max_candidates(20),
       save_ref(false),
       global_reference(-1),
       max_delta(0)
{
}

int do_main( int argc, char *argv[] )
{
       // IO parameters
       string in_filename;
       string registered_filebase("reg%04d.v");
       string reference_index_file;
       size_t skip = 0;
       size_t preskip = 20;
       size_t postskip	= 2;
       C3DMyocardPeriodicRegistration::RegistrationParams params;
       CCmdOptionList options(g_description);
       options.set_group("\nFile-IO");
       options.add(make_opt( in_filename, "in-file", 'i',
                             "input images following the naming pattern nameXXXX.ext",
                             CCmdOptionFlags::required_input, &C3DImageIOPluginHandler::instance()));
       options.add(make_opt( registered_filebase, "out-file", 'o',
                             "file name base for registered files given as C-format string",
                             CCmdOptionFlags::required_output, &C3DImageIOPluginHandler::instance()));
       options.add(make_opt(params.save_ref, "save-references", 0,
                            "Save synthetic references to files refXXXX.v"));
       options.set_group("\nPreconditions & Preprocessing");
       options.add(make_opt(skip, "skip", 'k',
                            "Skip images at the begin of the series"));
       options.add(make_opt(preskip, "preskip", 0,
                            "Skip images at the beginning+skip of the series when searching for high contrats image"));
       options.add(make_opt(postskip, "postskip", 0,
                            "Skip images at the end of the series when searching for high contrats image"));
       options.add(make_opt(params.max_candidates, "max-candidates", 0,
                            "maximum number of candidates for global reference image"));
       options.add(make_opt(params.series_select_cost, "image:cost=[ngf:eval=ds]",
                            "cost-series", 'S',
                            "Const function to use for the analysis of the series"));
       options.add(make_opt(reference_index_file, "ref-idx", 0,
                            "save reference index number to this file"));
       options.add(make_opt(params.global_reference, "global-reference", 'R',
                            "save reference index number to this file"));
       options.add(make_opt(params.max_delta, "max-subset-delta", 'D',
                            "Maximum delta between two elements of the prealigned subset"));
       options.set_group("\nRegistration");
       options.add(make_opt( params.minimizer, "gsl:opt=gd,step=0.01", "optimizer", 'O',
                             "Optimizer used for minimization"));
       options.add(make_opt( params.mg_levels, "mr-levels", 'l', "multi-resolution levels"));
       options.add(make_opt( params.transform_creator, "spline", "transForm", 'f',
                             "transformation type"));
       options.add(make_opt(params.pass1_cost, "image:cost=[ngf:eval=ds]", "cost-subset", '1',
                            "Cost function for registration during the subset registration"));
       options.add(make_opt(params.pass2_cost, "image:cost=ssd", "cost-final", '2',
                            "Cost function for registration during the final registration"));

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

       size_t start_filenum = 0;
       size_t end_filenum  = 0;
       size_t format_width = 0;
       string src_basename =
              get_filename_pattern_and_range(in_filename, start_filenum,
                                             end_filenum, format_width);
       C3DImageSeries in_images;

       for (size_t i = start_filenum; i < end_filenum; ++i) {
              string src_name = create_filename(src_basename.c_str(), i);
              P3DImage image = load_image<P3DImage>(src_name);

              if (!image)
                     throw create_exception<runtime_error>( "image ", src_name, " not found");

              cvdebug() << "read '" << src_name << "\n";
              in_images.push_back(image);
       }

       if (skip >= in_images.size()) {
              throw create_exception<invalid_argument>( "Try to skip ",  skip,
                            " images, but input set has only ",
                            in_images.size(), " images.");
       }

       C3DImageSeries series(in_images.begin() + skip, in_images.end());

       if (preskip > series.size() - 2)
              throw create_exception<invalid_argument>( "Try to skip ", preskip, " images th the beginning of a series ",
                            series.size(), " of relevant images.");

       if (postskip  >= series.size() - preskip)
              throw create_exception<invalid_argument>( "Try to skip ", postskip, " images at the end of a series ",
                            series.size(), " of relevant images");

       C3DMyocardPeriodicRegistration mpr(params);
       vector<P3DTransformation> transforms = mpr.run(series, preskip, postskip);

       if (!reference_index_file.empty()) {
              ofstream refidxfile(reference_index_file.c_str(), ios_base::out );
              refidxfile << mpr.get_ref_idx() + skip;
       }

       copy(series.begin(), series.end(), in_images.begin() + skip);
       bool success = true;
       auto ii = in_images.begin();

       for (size_t i = start_filenum; i < end_filenum; ++i, ++ii) {
              string out_name = create_filename(registered_filebase.c_str(), i);
              cvmsg() << "Save image " << i << " to " << out_name << "\n";
              success &= save_image(out_name, *ii);
       }

       return success ? EXIT_SUCCESS : EXIT_FAILURE;
}


#include <mia/internal/main.hh>
MIA_MAIN(do_main)
