/* -*- 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 <ostream>
#include <fstream>
#include <map>
#include <ctime>
#include <cstdlib>
#include <mia/core.hh>
#include <mia/core/fft1d_r2c.hh>
#include <queue>

#include <mia/2d/imageio.hh>
#include <mia/2d/filter.hh>
#include <mia/core/ica.hh>
#include <mia/2d/ica.hh>
#include <mia/core/slopeclassifier.hh>

NS_MIA_USE;
using namespace std;

const SProgramDescription g_description = {
       {pdi_group, "Tools for Myocardial Perfusion Analysis"},
       {pdi_short, "Run an ICA analysis on a series of 2D images"},

       {
              pdi_description, "This program is used to run a ICA on a series of myocardial "
              "perfusion images to create sythetic references that can be used for motion correction "
              "by image registration. If the aim is to run a full motion compensation then it is "
              "better to create a segmentation set and use mia-2dmyoica-nonrigid. If the input data "
              "is given by means of a segmentation set, then on can also use mia-2dmyocard-icaseries.\n"
              "This program is essentially used to test different options on how to run the ICA for"
              "reference image creation."
       },

       {
              pdi_example_descr, "Evaluate the synthetic references from images imageXXXX.exr and save them to "
              "refXXXX.exr by using five independend components, mean stripping, normalizing, "
              "and skipping 2 images."
       },


       {pdi_example_code, "-i imageXXXX.exr -o ref -k 2 -C 5 -m -n"}
};

unique_ptr<C2DImageSeriesICA> get_ica(vector<C2DFImage>& series, bool strip_mean,
                                      size_t& components, bool ica_normalize, int max_iterations,
                                      const CIndepCompAnalysisFactory& icatool)
{
       srand(time(NULL));
       unique_ptr<C2DImageSeriesICA> ica(new C2DImageSeriesICA(icatool, series, false));

       if (components > 0) {
              ica->set_max_iterations(max_iterations);
              ica->run(components, strip_mean, ica_normalize);
       } else {
              // maybe one can use the correlation and create an initial guess by combining
              // highly correlated curves.
              float min_cor = 0.0;

              for (int i = 7; i > 3; --i) {
                     unique_ptr<C2DImageSeriesICA> l_ica(new C2DImageSeriesICA(icatool, series, false));
                     l_ica->set_max_iterations(max_iterations);
                     l_ica->run(i, strip_mean, ica_normalize);
                     CSlopeClassifier cls(l_ica->get_mixing_curves(), strip_mean);
                     float max_slope = /*log2(i) * */ cls.get_max_slope_length_diff();
                     cvinfo() << "Components = " << i << " max_slope = " << max_slope << "\n";

                     if (min_cor < max_slope) {
                            min_cor = max_slope;
                            components = i;
                            ica.swap(l_ica);
                     }
              }
       }

       return ica;
}

void save_feature_image_unnamed(const string& base, int id, size_t max_comp,
                                const C2DImageSeriesICA& ica)
{
       stringstream fname;
       fname << base << "-" << "-" << max_comp <<  id << ".exr";

       if (!save_image(fname.str(), ica.get_feature_image(id)))
              throw create_exception<runtime_error>( "unable to save ", fname.str(), "\n");
}


void save_feature_image(const string& base, const string& descr, int id, size_t max_comp,
                        const C2DImageSeriesICA& ica)
{
       stringstream fname;

       if (id >= 0) {
              fname << base << "-" << descr << "-" << max_comp << ".exr";
              cvinfo() << "save id=" << id << " as feature '" << descr << "'\n";

              if (!save_image(fname.str(), ica.get_feature_image(id)))
                     throw create_exception<runtime_error>( "unable to save ", fname.str(), "\n");
       } else {
              fname << base << "-" << descr << "-" << max_comp << ".exr";
              P2DImage image(ica.get_mean_image().clone());

              if (!save_image(fname.str(), image))
                     throw create_exception<runtime_error>( "unable to save ", fname.str(), "\n");
       }
}

void save_feature_image(const string& base, const string& descr, size_t max_comp, P2DImage image)
{
       stringstream fname;
       fname << base << "-" << descr << "-" << max_comp << ".exr";

       if (!save_image(fname.str(), image))
              throw create_exception<runtime_error>( "unable to save '", fname.str(), "'");
}

void save_feature_image_png(const string& base, const string& descr, P2DImage image)
{
       stringstream fname;
       fname << base << "-" << descr << ".png";

       if (!save_image(fname.str(), image))
              throw create_exception<runtime_error>( "unable to save ", fname.str(), "\n");
}


void save_coefs(const string&  coefs_name, const C2DImageSeriesICA& ica)
{
       CSlopeClassifier::Columns mix = ica.get_mixing_curves();
       ofstream coef_file(coefs_name.c_str());

       for (size_t r = 0; r < mix[0].size(); ++r) {
              for (size_t c = 0; c < mix.size(); ++c) {
                     coef_file   << setw(10) << mix[c][r] << " ";
              }

              coef_file << "\n";
       }

       if (!coef_file.good())
              throw create_exception<runtime_error>( "unable to save coefficients to '", coefs_name, "'");
}

CIndepCompAnalysis::IndexSet get_all_without_periodic(const CSlopeClassifier::Columns& curves, bool strip_mean)
{
       CSlopeClassifier cls(curves, strip_mean);
       int periodic_index = cls.get_periodic_idx();
       CIndepCompAnalysis::IndexSet result;

       for (int i = 0; i < (int)curves.size(); ++i) {
              if (i != periodic_index)
                     result.insert(i);
       }

       return result;
}

class GetRegionCenter: public TFilter<C2DFVector>
{
public:
       template <typename T>
       C2DFVector operator() (const T2DImage<T>& image) const
       {
              C2DFVector result;
              size_t n = 0;
              typename T2DImage<T>::const_iterator i = image.begin();

              for (size_t y = 0; y < image.get_size().y; ++y)
                     for (size_t x = 0; x < image.get_size().x; ++x, ++i) {
                            if (*i) {
                                   result.x += x;
                                   result.y += y;
                                   ++n;
                            }
                     }

              if (!n)
                     throw create_exception<invalid_argument>( "GetRegionCenter: provided an empty region");

              return result / float(n);
       };
};


class GetRegionSize: public TFilter<size_t>
{
public:
       template <typename T>
       size_t operator() (const T2DImage<T>& image) const
       {
              size_t n = 0;
              typename T2DImage<T>::const_iterator i = image.begin();

              for (size_t y = 0; y < image.get_size().y; ++y)
                     for (size_t x = 0; x < image.get_size().x; ++x, ++i) {
                            if (*i)
                                   ++n;
                     }

              return n;
       };
};

class GetClosestRegionLabel: public TFilter<int>
{
public:
       GetClosestRegionLabel(const C2DFVector& point): m_point(point) {}
       template <typename T>
       int operator() (const T2DImage<T>& image) const;
private:
       C2DFVector m_point;
       struct Collector {
              size_t size;
              C2DFVector center;
       };

};

typedef pair<float, size_t> element;

P2DFilter create_LV_cropper(P2DImage rvlv_feature,
                            float LV_mask_amplify,
                            const string& feature_image_base,
                            C2DBounds& crop_start
                           )
{
       C2DImageFilterChain kmeans_filter_chain({"close:shape=[sphere:r=2]",
                                               "open:shape=[sphere:r=2]"
                                               });
       C2DImageFilterChain RV_filter_chain({"label", "selectbig"});
       C2DImageFilterChain LVcandidate_filter_chain({"binarize:min=0,max=0", "label"});
       P2DImage pre_kmeans = kmeans_filter_chain.run(rvlv_feature);
       P2DImage RV;
       P2DImage kmeans;
       size_t nc = 1;
       size_t npixels;

       do {
              ++nc;
              stringstream kfd;
              kfd << "kmeans:c=" << 2 * nc + 1;
              kmeans = run_filter(*pre_kmeans, kfd.str().c_str());
              stringstream kfb;
              kfb << "binarize:min=" << 2 * nc << ",max=" << 2 * nc;
              P2DImage kmeans_binarized = run_filter(*kmeans, kfb.str().c_str());
              RV = RV_filter_chain.run(kmeans_binarized);
              npixels = ::mia::filter(GetRegionSize(), *RV);
       } while (10 * npixels > rvlv_feature->get_size().x * rvlv_feature->get_size().y && nc < 5);

       if (nc == 5)
              return P2DFilter();

       P2DImage LV_candidates = LVcandidate_filter_chain.run(kmeans);
       C2DFVector RV_center = ::mia::filter(GetRegionCenter(), *RV);
       int label = ::mia::filter(GetClosestRegionLabel(RV_center), *LV_candidates);
       stringstream binarize_lv;
       binarize_lv << "binarize:min=" << label << ",max=" << label;
       cvinfo() << "LV label=  " << label << "\n";
       P2DImage LV = run_filter(*LV_candidates, binarize_lv.str().c_str());
       C2DFVector LV_center = ::mia::filter(GetRegionCenter(), *LV);
       cvinfo() << "RV center = " << RV_center << "\n";
       cvinfo() << "LV center = " << LV_center << "\n";
       float r = LV_mask_amplify * (LV_center - RV_center).norm();
       cvinfo() << "r = " << r << "\n";
       stringstream mask_lv;
       crop_start = C2DBounds(int(LV_center.x - r), int(LV_center.y - r));

       // this is ugly and should be replaced
       if (crop_start.x > LV_center.x || crop_start.y > LV_center.y)
              return P2DFilter();

       mask_lv << "crop:start=[" << crop_start
               << "],end=[" << C2DBounds(int(LV_center.x + r), int(LV_center.y + r)) << "]";
       cvinfo() << "crop region = '" << mask_lv.str() << "'\n";

       if (!feature_image_base.empty()) {
              save_feature_image_png("feature", "kmeans", kmeans);
              save_feature_image_png(feature_image_base, "RV", RV);
              save_feature_image_png(feature_image_base, "LV", LV);
       }

       return C2DFilterPluginHandler::instance().produce(mask_lv.str().c_str());
}


CIndepCompAnalysis::IndexSet get_LV_RV_Perfusion(const CSlopeClassifier::Columns& curves)
{
       priority_queue<element> sorted;
       element e;

       for (CSlopeClassifier::Columns::const_iterator i = curves.begin();
            i != curves.end(); ++i, ++e.second) {
              CSlopeStatistics stat(*i, 0);
              e.first = - stat.get_curve_length() / stat.get_range();
              sorted.push(e);
       }

       CIndepCompAnalysis::IndexSet result;
       result.insert(sorted.top().second);
       cvinfo() << "add " << sorted.top().first << " as " << sorted.top().second << "\n";
       sorted.pop();
       result.insert(sorted.top().second);
       cvinfo() << "add " << sorted.top().first << " as " << sorted.top().second << "\n";
       sorted.pop();
       cvinfo() << "add " << sorted.top().first << " as " << sorted.top().second << "\n";
       result.insert(sorted.top().second);
       return result;
}

struct convert2byte {
       unsigned char operator() (float x) const
       {
              if ( x < 0)
                     return 0;

              if ( x > 255)
                     return 255;

              return (unsigned char)x;
       }
};

P2DImage convert_to_ubyte(const C2DFImage&  pref)
{
       C2DUBImage *result = new C2DUBImage(pref.get_size());
       transform(pref.begin(), pref.end(), result->begin(), convert2byte());
       return P2DImage(result);
}

int do_main( int argc, char *argv[] )
{
       string src_name("data0000.exr");
       string out_name("ref");
       string out_type("exr");
       string crop_name("crop");
       string coefs_name;
       size_t first =  2;
       size_t last  = 60;
       size_t components = 0;
       bool strip_mean = false;
       bool ica_normalize = false;
       bool skip_only_periodic = false;
       string feature_image_base;
       string numbered_feature_image;
       float LV_mask = 0.0; // no mask
       int max_iterations = 0;
       PIndepCompAnalysisFactory icatool;
       const auto& imageio = C2DImageIOPluginHandler::instance();
       CCmdOptionList options(g_description);
       options.add(make_opt( src_name, "in-base", 'i', "input file name ofolloing pattern nameXXXX.ext X=numbers",
                             CCmdOptionFlags::required_input, &imageio));
       options.add(make_opt( coefs_name, "coefs", 0, "output mixing coefficients to this file"));
       options.add(make_opt( out_name, "out-base", 'o', "output file name base",
                             CCmdOptionFlags::output));
       options.add(make_opt( out_type, imageio.get_supported_suffix_set(), "type", 't',
                             "output file type"));
       options.add(make_opt( first, "skip", 'k', "skip images at beginning of series"));
       options.add(make_opt( last, "end", 'e', "last image in series"));
       options.add(make_opt( components, "components", 'C', "nr. of components, 0=estimate automatically"));
       options.add(make_opt( strip_mean, "strip-mean", 'm', "strip mean image from series"));
       options.add(make_opt( feature_image_base, "save-features", 'f', "save feature image"));
       options.add(make_opt(skip_only_periodic, "strip-periodic", 'p', "strip only periodic component"));
       options.add(make_opt( icatool, "internal", "fastica", 0, "FastICA implementationto be used"));
       options.add(make_opt(max_iterations, "max-ica-iterations", 'x', "max ICA solver iterations"));
       options.add(make_opt( ica_normalize, "ica-normalize", 'n', "ica_normalize feature images"));
       options.add(make_opt( numbered_feature_image, "all-features", 0, "save all feature images to"));
       options.add(make_opt( LV_mask, "LV-crop-amp", 'L', "LV crop mask amplification, 0.0 = don't crop"));

       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(src_name, start_filenum, end_filenum, format_width);

       if (start_filenum < first)
              start_filenum = first;

       if (end_filenum > last)
              end_filenum = last;

       if (start_filenum + components >= end_filenum) {
              throw create_exception<invalid_argument>( "require at least ", components, " images");
       }

       // load images
       vector<C2DFImage> series;
       FCopy2DImageToFloatRepn converter;

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

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

              cvdebug() << "read '" << src_name << "\n";
              series.push_back(converter(*image));
       }

       cvmsg() << "Got series of " << series.size() << " images\n";
       // always strip mean
       unique_ptr<C2DImageSeriesICA> ica =
              get_ica(series, strip_mean, components, ica_normalize, max_iterations, *icatool);
       CSlopeClassifier::Columns curves = ica->get_mixing_curves();
       CIndepCompAnalysis::IndexSet component_set = skip_only_periodic ?
                     get_all_without_periodic(curves, strip_mean) :
                     get_LV_RV_Perfusion(curves);
       // important: mean_stripped applies only if it was done before ICA
       CSlopeClassifier cls(curves, false);
       cvinfo() << "LV series       = " << cls.get_LV_idx() << "\n";
       cvinfo() << "RV series       = " << cls.get_RV_idx() << "\n";
       cvinfo() << "Periodic series = " << cls.get_periodic_idx() << "\n";
       cvinfo() << "Perfusion series = " << cls.get_perfusion_idx() << "\n";
       cvinfo() << "Baseline series = " << cls.get_baseline_idx() << "\n";
       P2DFilter cropper;
       C2DBounds crop_start;

       try {
              // create crop filter and store source files too.
              if (LV_mask > 0.0) {
                     C2DImageSeriesICA::IndexSet plus;
                     C2DImageSeriesICA::IndexSet minus;
                     plus.insert(cls.get_RV_idx());
                     minus.insert(cls.get_LV_idx());
                     P2DImage rvlv_feature = ica->get_delta_feature(plus, minus);
                     cropper = create_LV_cropper(rvlv_feature, LV_mask, feature_image_base, crop_start);

                     // feature images do not work, try peak images
                     if (!cropper) {
                            int RV_peak = cls.get_RV_peak();
                            int LV_peak = cls.get_LV_peak();

                            if (RV_peak < 0 || LV_peak < 0)
                                   throw std::runtime_error("Feature images doen't work, and peaks could not be identified");

                            cvinfo() << "Using RV peaks =" << RV_peak << " and " << LV_peak << "\n";
                            C2DFImage *prvlv_diff = new C2DFImage(series[0].get_size());
                            P2DImage rvlv_feature(prvlv_diff);
                            transform(series[RV_peak].begin(), series[RV_peak].end(),
                                      series[LV_peak].begin(), prvlv_diff->begin(),
                            [](double x, double y) {
                                   return y - x;
                            });
                            cropper = create_LV_cropper(rvlv_feature, LV_mask, feature_image_base, crop_start);
                     }

                     if (!cropper)
                            throw std::runtime_error("Neither feature images nor peak images work");

                     ofstream crop_idx((crop_name + ".txt").c_str());
                     crop_idx << crop_start;
              }
       } catch (std::exception& x) {
              cvwarn() << "unable to create cropping filter:'" << x.what() << "'\n";
              cvwarn() << "use original images instead\n";
       }

       for (size_t i = start_filenum; i < end_filenum; ++i) {
              C2DFImage *pref = new C2DFImage(ica->get_partial_mix(i - start_filenum, component_set));
              P2DImage reference(pref);

              // this is a quick hack that should be improved
              if (out_type == "png")
                     reference = convert_to_ubyte(*pref);

              if (LV_mask > 0.0) {
                     stringstream scrop_name;
                     scrop_name << crop_name << setw(format_width) << setfill('0') << i << "." << out_type;
                     string src_name = create_filename(src_basename.c_str(), i);

                     if (cropper) {
                            reference = cropper->filter(*reference);
                            P2DImage crop_source = cropper->filter(*load_image2d(src_name));

                            if (!save_image(scrop_name.str(), crop_source))
                                   throw create_exception<runtime_error>( "unable to save ", scrop_name.str(), "\n");
                     } else  {
                            if (!save_image(scrop_name.str(), load_image2d(src_name)))
                                   throw create_exception<runtime_error>( "unable to save ", scrop_name.str(), "\n");
                     }
              }

              stringstream fname;
              fname << out_name << setw(format_width) << setfill('0') << i << "." << out_type;

              if (!save_image(fname.str(), reference))
                     throw create_exception<runtime_error>( "unable to save ", fname.str(), "\n");

              cvdebug() << "wrote '" << fname.str() << "\n";
       }

       if (!feature_image_base.empty()) {
              if (cls.get_periodic_idx() >= 0)
                     save_feature_image(feature_image_base, "periodic", cls.get_periodic_idx(), components, *ica);

              save_feature_image(feature_image_base, "LV", cls.get_LV_idx(), components, *ica);
              save_feature_image(feature_image_base, "RV", cls.get_RV_idx(), components, *ica);
              {
                     C2DImageSeriesICA::IndexSet plus;
                     plus.insert(cls.get_RV_idx());
                     C2DImageSeriesICA::IndexSet minus;
                     minus.insert(cls.get_LV_idx());
                     save_feature_image(feature_image_base, "RV-LV", components,
                                        ica->get_delta_feature(plus, minus) );
              }

              if (cls.get_perfusion_idx() >= 0) {
                     save_feature_image(feature_image_base, "perfusion", cls.get_perfusion_idx(),
                                        components, *ica);
                     C2DImageSeriesICA::IndexSet plus;
                     plus.insert(cls.get_RV_idx());
                     C2DImageSeriesICA::IndexSet minus;
                     minus.insert(cls.get_LV_idx());
                     minus.insert(cls.get_perfusion_idx());
                     save_feature_image(feature_image_base, "RV-perf-RVLV", components,
                                        ica->get_delta_feature(plus, minus));
              }

              if ( cls.get_baseline_idx() >= 0) {
                     save_feature_image(feature_image_base, "baseline", cls.get_baseline_idx(), components, *ica);
              }

              if (strip_mean)
                     save_feature_image(feature_image_base, "mean", -1, components, *ica);
       }

       if (!coefs_name.empty())
              save_coefs(coefs_name, *ica);

       if (!numbered_feature_image.empty())
              for (size_t i = 0; i < components; ++i)
                     save_feature_image_unnamed(numbered_feature_image, i, components, *ica);

       return EXIT_SUCCESS;
};

template <typename T>
int GetClosestRegionLabel::operator() (const T2DImage<T>& image) const
{
       map<int, Collector> collector_map;
       auto i = image.begin();

       for (size_t y = 0; y < image.get_size().y; ++y)
              for (size_t x = 0; x < image.get_size().x; ++x, ++i) {
                     if (!*i)
                            continue;

                     auto ic = collector_map.find(*i);

                     if (ic ==  collector_map.end()) {
                            Collector col;
                            col.size = 1;
                            col.center = C2DFVector(x, y);
                            collector_map[*i] = col;
                     } else {
                            ++ic->second.size;
                            ic->second.center.x += x;
                            ic->second.center.y += y;
                     }
              }

       float min_weighted_distance = numeric_limits<float>::max();
       int label = 0;

       for (auto i = collector_map.begin(); i != collector_map.end(); ++i) {
              i->second.center /= i->second.size;
              float wdist = (i->second.center - m_point).norm() / i->second.size;

              if ( min_weighted_distance > wdist ) {
                     min_weighted_distance = wdist;
                     label = i->first;
              }
       }

       return label;
};


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