/* -*- 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/3d/imageio.hh>
#include <mia/3d/filter.hh>
#include <mia/core.hh>
#include <mia/core/labelmap.hh>

#include <map>
#include <set>

using namespace mia;
using namespace std;

const SProgramDescription g_description = {
       {pdi_group, "Analysis, filtering, combining, and segmentation of 3D images"},
       {pdi_short, "Evaluate distance statistics for two labeled images."},
       {
              pdi_description, "This program is used to evaluate the distance between equally labelled "
              "areas in two images. The output file is a csv file containing the distances for "
              "each labeled coordinate in the test image in the following form:\n"
              "\n"
              "  label,n-samples,distance,distance,...\n"
              "\nA simple R program to get some per label statistics from this data would look like\n\n"
              "    args <- commandArgs(TRUE)\n"
              "    data <- read.csv(args[1], header=FALSE)\n"
              "    for ( i in 1:dim(data)[1] )  {\n"
              "        line <- data[i,]\n"
              "        label <- as.integer(line[1])\n"
              "        end_range <- 2 + as.numeric(line[2])\n"
              "        s <- as.numeric(line[3:end_range])\n"
              "        result <- sprintf(\"%d %8.4f %8.4f %8.4f %8.4f\\n\"\n"
              "            label, mean(s), sqrt(var(s)),  median(s), max(s))\n"
              "        cat(result)\n"
              "    }\n\n"
       },
       {
              pdi_example_descr, "Evaluate the distances for each label available in image.v to the "
              "corresponding labels in the image reference.v ans store the result "
              "a coma separated list of values, i.e. distances.csv."
       },
       {pdi_example_code, "-i image.v -r reference.v -o distances.csv" }
};


unsigned short translate_label( unsigned short l, const CLabelMap& map)
{
       if (map.empty())
              return l;

       auto idx = map.find(l);

       if (idx != map.end())
              return idx->second;

       throw create_exception<runtime_error>("unmapped label '", l, "' encountered");
}



typedef vector<float> CDistanceResult;


typedef map<unsigned short, CDistanceResult> CDistanceResultMap;

class  FEvaluateDistances: public TFilter<CDistanceResultMap>
{
public:
       FEvaluateDistances(const CLabelMap& lm);

       template <typename T, typename S>
       CDistanceResultMap operator ()(const T3DImage<T>& test, const T3DImage<S>& ref) const;
private:
       const CLabelMap m_lm;
};

template <typename T, typename S, bool, bool>
struct __dispatch_filter {
       static CDistanceResultMap apply(const T3DImage<T>& MIA_PARAM_UNUSED(test),
                                       const T3DImage<S>& MIA_PARAM_UNUSED(ref),
                                       const CLabelMap& MIA_PARAM_UNUSED(lm))
       {
              throw create_exception<invalid_argument>("CDistanceResultMap: input data not supported,"
                            "got ", __type_descr<T>::value, " and ",
                            __type_descr<S>::value, " but require integral scalar types");
       }
};

template <typename T>
P3DImage get_mask( const T3DImage<T>& image, unsigned short label)
{
       stringstream descr;
       descr << "binarize:min=" << label << ",max=" << label;
       return  run_filter(image, descr.str().c_str());
}

template <typename T>
C3DFImage get_ref_distance(const T3DImage<T>& image, unsigned short label, const C3DFilter& distance)
{
       auto mask = get_mask( image, label);
       auto d = distance.filter(*mask);
       return dynamic_cast<C3DFImage&>(*d);
}


template <typename T, typename S>
struct __dispatch_filter<T, S, true, true> {
       static CDistanceResultMap apply(const T3DImage<T>& test, const T3DImage<S>& ref, const CLabelMap& lm)
       {
              set<T> test_labels;
              CDistanceResultMap result;
              P3DFilter distance = produce_3dimage_filter("distance");

              // add all the test labels
              for ( auto t : test)
                     if (t != 0)
                            test_labels.insert(t);

              for (auto tl : test_labels) {
                     auto rl = translate_label(tl, lm);
                     auto test_maskp = get_mask(test, tl);
                     const C3DBitImage& test_mask = dynamic_cast<const C3DBitImage&>(*test_maskp);
                     C3DFImage ref_mask = get_ref_distance(ref, rl, *distance);
                     CDistanceResult distances;
                     auto it = test_mask.begin();
                     auto et = test_mask.end();
                     auto ir = ref_mask.begin();

                     while (it != et) {
                            if (*it)
                                   distances.push_back(*ir);

                            ++it;
                            ++ir;
                     }

                     result[tl] = distances;
              }

              return result;
       }
};

FEvaluateDistances::FEvaluateDistances(const CLabelMap& lm):
       m_lm(lm)
{
}

template <typename T, typename S>
CDistanceResultMap FEvaluateDistances::operator ()(const T3DImage<T>& test, const T3DImage<S>& ref) const
{
       const auto T_integral = is_integral<T>::value;
       const auto S_integral = is_integral<S>::value;
       return __dispatch_filter<T, S, T_integral, S_integral>::apply(test, ref, m_lm);
}


int do_main( int argc, char *argv[] )
{
       string in_filename;
       string ref_filename;
       string label_translate_filename;
       string out_filename;
       const auto& imageio = C3DImageIOPluginHandler::instance();
       stringstream filter_names;
       CCmdOptionList options(g_description);
       options.set_group("File-IO");
       options.add(make_opt( in_filename, "in-labels", 'i', "input label image",
                             CCmdOptionFlags::required_input, &imageio));
       options.add(make_opt( ref_filename, "ref-labels", 'r', "reference label image",
                             CCmdOptionFlags::required_input, &imageio));
       options.add(make_opt( label_translate_filename, "label-map", 'l', "optional mapping of label numbers",
                             CCmdOptionFlags::input));
       options.add(make_opt( out_filename, "out-file", 'o', "output file name to write the distances to. "
                             "The output file is a csv file, containing distances listed for each label.",
                             CCmdOptionFlags::required_output));

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

       auto test_labels = load_image3d(in_filename);
       auto ref_labels = load_image3d(ref_filename);
       CLabelMap lmap;

       if (!label_translate_filename.empty()) {
              ifstream ifs(label_translate_filename);
              lmap = CLabelMap(ifs);
       }

       FEvaluateDistances evaluater(lmap);
       auto result = mia::filter(evaluater, *test_labels, *ref_labels);
       // save result
       ofstream outf(out_filename);

       if (outf.bad())
              throw create_exception<runtime_error>("Error opening file '", out_filename, "' for writing");

       for (auto i : result)
              outf << i.first << "," << i.second.size() << "," << i.second << "\n";

       if (outf.bad())
              throw create_exception<runtime_error>("Error opening file '", out_filename, "' for writing");

       outf.close();
       return EXIT_SUCCESS;
}

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