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

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include <mia/core/histogram.hh>
#include <mia/core/cmdlineparser.hh>
#include <mia/core/kmeans.hh>
#include <mia/2d.hh>

#include <memory>
#include <vector>

using namespace mia;
using namespace std;

typedef vector<C2DFImage> C2DFImageVec;

const SProgramDescription g_description = {
       {pdi_group, "Analysis, filtering, combining, and segmentation of 2D images"},
       {pdi_short, "Run a k-means segmentation of a 2D image."},
       {
              pdi_description, "This program runs the segmentation of a 2D image by applying "
              "a localized k-means approach that helps to overcome intensity inhomogeneities "
              "in the image. The approach evaluates a global k-means clustering, and then "
              "separates the image into overlapping regions where more k-means iterations "
              "are run only including the locally present classes, i.e. the classes that "
              "relatively contain more pixels than a given threshold."
       },
       {
              pdi_example_descr, "Run the segmentation on image test.png using three classes, "
              "local regions of 40 pixels (grid width 20 pixels), and a class ignore threshold of 0.01."
       },
       {pdi_example_code, "-i test.png -o label.png -n 3 -g 20 -t 0.01"}
};


class FKmeans: public TFilter<vector<double>>
{
public:
       FKmeans(const vector<double>& in_classes, vector<C2DSBImage>& class_output, double rel_cluster_threshold);
       void set_range(const C2DBounds& start, const C2DBounds& end, int xtarget, int ytarget);

       template <typename T>
       vector<double> operator () (const T2DImage<T>& image);

private:
       vector<double> m_in_classes;
       vector<C2DSBImage>& m_class_output;
       double m_rel_cluster_threshold;
       C2DBounds m_start;
       C2DBounds m_end;
       int m_xtarget;
       int m_ytarget;
};

class FKMeansLocal: public TFilter<P2DImage>
{
public:
       FKMeansLocal(const vector<C2DDImage>& class_centers);

       template <typename T>
       P2DImage operator () (const T2DImage<T>& image) const;
private:
       const vector<C2DDImage>& m_class_centers;
};


int do_main(int argc, char *argv[])
{
       string in_filename;
       string out_filename;
       string out_filename2;
       string debug_filename;
       int blocksize = 20;
       unsigned n_classes = 3;
       double rel_cluster_threshold = 0.0001;
       const auto& image2dio = C2DImageIOPluginHandler::instance();
       CCmdOptionList opts(g_description);
       opts.set_group("File-IO");
       opts.add(make_opt( in_filename, "in-file", 'i', "image to be segmented", CCmdOptionFlags::required_input, &image2dio ));
       opts.add(make_opt( out_filename, "out-file", 'o', "class label image based on merging local labels", CCmdOptionFlags::output, &image2dio ));
       opts.add(make_opt( out_filename2, "out-file2", 'O', "class label image based on classifying by using the interpolated class centers ",
                          CCmdOptionFlags::output, &image2dio ));
       opts.add(make_opt( debug_filename, "debug-centers", 'D', "Save the class center images to a debug file (type is double)",
                          CCmdOptionFlags::output, &image2dio ));
       opts.set_group("Parameters");
       opts.add(make_opt( blocksize, EParameterBounds::bf_min_closed, {3},
                          "grid-spacing", 'g', "Spacing of the grid used to modulate the intensity inhomogeneities"));
       opts.add(make_opt( n_classes, EParameterBounds::bf_closed_interval, {2u, 127u},
                          "nclasses", 'n', "Number of intensity classes to segment"));
       opts.add(make_opt( rel_cluster_threshold, EParameterBounds::bf_min_closed | EParameterBounds::bf_max_open,
       {0.0, 1.0}, "relative-cluster-threshold", 't', "Number of intensity classes to segment"));

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

       if (out_filename.empty() && out_filename2.empty()) {
              throw invalid_argument("You must specify at least one output file");
       }

       auto in_image = load_image2d(in_filename);
       stringstream kfilter_ss;
       kfilter_ss << "kmeans:c=" << n_classes;
       auto full_kmeans = run_filter(in_image, kfilter_ss.str().c_str());
       vector <C2DSBImage> class_output(4, C2DSBImage(in_image->get_size()));

       for (auto image : class_output) {
              fill(image.begin(), image.end(), -1);
       }

       auto full_classes_ptr = full_kmeans->get_attribute(ATTR_IMAGE_KMEANS_CLASSES);
       const CVDoubleAttribute& full_classes = dynamic_cast<const CVDoubleAttribute&>(*full_classes_ptr);
       FKmeans local_kmeans(full_classes, class_output, rel_cluster_threshold);
       // true block size is selected so that the start and end pixles are on the boundaries
       unsigned nx = (in_image->get_size().x + blocksize - 1) / blocksize;
       unsigned ny = (in_image->get_size().y  + blocksize - 1) / blocksize;
       double scalex = double(in_image->get_size().x - 1) / double (nx - 1);
       double scaley = double(in_image->get_size().y - 1) / double (ny - 1);
       cvmsg() << "using a scale factor of " << scalex << " x " << scaley << "\n";
       int xtarget = 0;
       int ytarget = 0;
       vector<C2DDImage> local_class_center_img(n_classes, C2DDImage(C2DBounds(nx, ny)));
       vector<C2DDImage::iterator> ilcc;

       for (unsigned i = 0; i < n_classes; ++i) {
              ilcc.push_back(local_class_center_img[i].begin());
       }

       for (unsigned iy = 0; iy < ny; ++iy, ytarget = (ytarget + 1) % 2) {
              for (unsigned ix = 0; ix < nx; ++ix, xtarget = (xtarget + 1) % 2) {
                     C2DBounds start(static_cast<unsigned>(ix * scalex), static_cast<unsigned>(iy * scaley));
                     C2DBounds end(static_cast<unsigned>(2 * (ix + 1) * scalex), static_cast<unsigned>( 2 * (iy + 1) * scaley));

                     if (start.x > in_image->get_size().x)
                            start.x = 0;

                     if (start.y > in_image->get_size().y)
                            start.y = 0;

                     if (end.x > in_image->get_size().x)
                            end.x = in_image->get_size().x;

                     if (end.y > in_image->get_size().y)
                            end.y = in_image->get_size().y;

                     local_kmeans.set_range(start, end, xtarget, ytarget);
                     auto local_class_centers = mia::accumulate(local_kmeans, *in_image);

                     for (unsigned i = 0; i < n_classes; ++i) {
                            *ilcc[i]++ = local_class_centers[i];
                     }
              }
       }

       if (!out_filename.empty())  {
              // merge the class labels;
              cvmsg() << "Create result image of size" << in_image->get_size() << "\n";
              C2DUBImage result_labels(in_image->get_size(), *in_image);
              auto iout = result_labels.begin();
              auto eout = result_labels.end();
              vector<C2DSBImage::const_iterator> ii;

              for (int i = 0; i < 4; ++i) {
                     ii.push_back(class_output[i].begin());
              }

              vector<unsigned> counts(n_classes);

              while (iout != eout)  {
                     fill(counts.begin(), counts.end(), 0);

                     for (int i = 0; i < 4; ++i) {
                            if (*ii[i] >= 0) {
                                   ++counts[*ii[i]];
                            }

                            ++ii[i];
                     }

                     unsigned max_count = 0;
                     int max_class = -1;
                     unsigned k = 0;

                     while ( (max_count << 1) < n_classes && k < n_classes ) {
                            if (max_count < counts[k]) {
                                   max_count = counts[k];
                                   max_class = k;
                            }

                            ++k;
                     }

                     *iout = max_class;
                     ++iout;
              }

              save_image(out_filename, result_labels);
       }

       if (!out_filename2.empty()) {
              cvmsg() << "Save interpolated class center based labeling to " << out_filename2 << "\n";
              stringstream scale_filter_ss;
              scale_filter_ss << "scale:sx=" << in_image->get_size().x
                              << ",sy=" << in_image->get_size().y << ",interp=[bspline:d=1]";
              auto scale_filter = produce_2dimage_filter(scale_filter_ss.str().c_str());
              C2DImageIOPPH::Data scaled_images;
              transform(local_class_center_img.begin(), local_class_center_img.end(),
                        back_inserter<C2DImageIOPPH::Data>(scaled_images),
              [ scale_filter ](const C2DDImage & img) {
                     return scale_filter->filter(img);
              });
              transform(scaled_images.begin(), scaled_images.end(), local_class_center_img.begin(),
              [](P2DImage img) {
                     return dynamic_cast<const C2DDImage&>(*img);
              });

              if (!debug_filename.empty()) {
                     image2dio.save(debug_filename, scaled_images);
              }

              FKMeansLocal means_local(local_class_center_img);
              auto result = mia::filter(means_local, *in_image);
              save_image(out_filename2, result);
       }

       return EXIT_SUCCESS;
}


FKmeans::FKmeans(const vector<double>& in_classes, vector<C2DSBImage>& class_output,
                 double rel_cluster_threshold):
       m_in_classes(in_classes),
       m_class_output(class_output),
       m_rel_cluster_threshold(rel_cluster_threshold)
{
}

void FKmeans::set_range(const C2DBounds& start, const C2DBounds& end, int xtarget, int ytarget)
{
       m_start = start;
       m_end = end;
       m_xtarget = xtarget;
       m_ytarget = ytarget;
}


template <typename T>
vector<double> FKmeans::operator () (const T2DImage<T>& image)
{
       auto ibegin = image.begin_range(m_start, m_end);
       auto iend = image.end_range(m_start, m_end);
       auto i = ibegin;
       size_t n = 0;
       vector<size_t> cluster_sizes(m_in_classes.size());
       size_t l = m_in_classes.size() - 1;

       while ( i != iend ) {
              const unsigned c = kmeans_get_closest_clustercenter(m_in_classes, l, *i);
              ++cluster_sizes[c];
              ++i;
              ++n;
       }

       vector<double> rel_cluster_sizes(m_in_classes.size());
       transform(cluster_sizes.begin(), cluster_sizes.end(), rel_cluster_sizes.begin(),
       [n](double x) {
              return x / n;
       });
       cvinfo() << "Block [" << m_start << "],[" <<  m_end << "] Relative cluster sizes: " << rel_cluster_sizes << "\n";
       vector<short> map_idx;
       vector<double> result;

       for (unsigned i = 0; i < m_in_classes.size(); ++i) {
              if (rel_cluster_sizes[i] > m_rel_cluster_threshold) {
                     result.push_back(m_in_classes[i]);
                     map_idx.push_back(i);
              }
       }

       vector<int> class_relation(n);
       int biggest_class = -1;

       // iterate to update the class centers
       for (size_t  l = 1; l < 4; l++) {
              if (kmeans_step(ibegin, iend, class_relation.begin(), result,
                              result.size() - 1, biggest_class))
                     break;
       }

       transform(class_relation.begin(), class_relation.end(),
                 m_class_output[m_ytarget * 2 + m_xtarget].begin_range(m_start, m_end),
       [map_idx](int idx) {
              return map_idx[idx];
       });
       // put the classes back into the right size
       vector<double> remapped_result(m_in_classes);

       for (unsigned i = 0; i < result.size(); ++i) {
              remapped_result[map_idx[i]] = result[i];
       }

       cvinfo() << "remapped cc=" << remapped_result << "\n";
       return remapped_result;
}

FKMeansLocal::FKMeansLocal(const vector<C2DDImage>& class_centers):
       m_class_centers(class_centers)
{
}

template <typename T>
P2DImage FKMeansLocal::operator () (const T2DImage<T>& image) const
{
       C2DUBImage *result_labels = new C2DUBImage(image.get_size(), image);
       auto i = image.begin();
       auto iout = result_labels->begin();
       auto eout = result_labels->end();
       vector<C2DDImage::const_iterator> ii;
       unsigned n_classes = m_class_centers.size();

       for (unsigned i = 0; i < n_classes; ++i) {
              ii.push_back(m_class_centers[i].begin());
       }

       vector<double> centers( n_classes );

       while (iout != eout)  {
              for (unsigned i = 0; i < n_classes; ++i) {
                     centers[i] = *ii[i];
                     ++ii[i];
              }

              *iout = kmeans_get_closest_clustercenter(centers, n_classes - 1, *i);
              ++iout;
              ++i;
       }

       return P2DImage(result_labels);
}


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

