/* -*- 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 <iostream>
#include <iomanip>
#include <fstream>

#include <string>
#include <stdexcept>
#include <dlfcn.h>

#include <mia/core.hh>
#include <mia/2d/imageio.hh>
#include <mia/3d/2dimagefifofilter.hh>
#include <mia/3d/vector.hh>


using namespace std;
NS_MIA_USE;

const SProgramDescription g_description = {
       {pdi_group, "Processing of series of 2D images in a 3D fashion (out-of-core)"},
       {pdi_short, "Evaluate the distance between two binary shapes."},

       {
              pdi_description, "This program takes two binary masks as input and evaluates the distance of one "
              "mask with respect to the other in voxel space. The output is given as text file "
              "with the coordinates of the source voxels and their distance to the reference "
              "mask. Correction for voxel size must be done after processing."
       },

       {
              pdi_example_descr, "Evaluate the distance of the mask given in the images srcXXXX.png to the mask "
              "given in refXXXX.png and save the result to distances.txt"
       },

       {pdi_example_code, "-i srcXXXX.png -r refXXXX.png -o distances.txt"}
};


const float g_far = numeric_limits<float>::max();

class C3DDT
{
public:
       C3DDT();
       void initialise(const C2DImage& wimage);

       void read( const C2DImage& wimage, short q);

       void get_slice(size_t s, const C2DImage& image, vector<pair<C3DBounds, float>>& result) const;

       bool test();

       bool test_initialise(const C2DBounds& size);
private:

       float d(float fq, float q, float fv, float v)const;

       void dt1d(vector<float>& f)const;
       void dt2d(C2DFImage& image)const;

       struct SParabola {
              short k;
              short v;
              float z;
              float fv;
       };

       C2DBounds _M_size;
       vector<short> _M_k;
       vector< vector<SParabola>> _M_zdt;
};

C3DDT::C3DDT()
{
}

float C3DDT::d(float fq, float q, float fv, float v) const
{
       const float d = ( fq  - fv + q * q - v * v) / (q - v) * 0.5;
       return d;
}


void C3DDT::dt1d(vector<float>& r)const
{
       TRACE_FUNCTION;
       vector<float> f(r);
       vector<int> v(f.size());
       vector<float> z(f.size() + 1);
       int k = 0;
       v[0] = 0;
       z[0] = -numeric_limits<float>::max();
       z[1] = +numeric_limits<float>::max();

       for (size_t q = 1; q < f.size(); q++) {
              float s  = d(f[q], q, f[v[k]], v[k]);

              while (s <= z[k]) {
                     --k;
                     s  = d(f[q], q, f[v[k]], v[k]);
              }

              ++k;
              v[k] = q;
              z[k] = s;
              z[k + 1] = numeric_limits<float>::max();
       }

       k = 0;

       for (size_t q = 0; q < f.size(); ++q) {
              while (z[k + 1] < q)
                     ++k;

              float delta = float(q) - v[k];
              r[q] = delta * delta + f[v[k]];
       }
}



void C3DDT::initialise(const C2DImage& wimage)
{
       TRACE_FUNCTION;
       _M_size = wimage.get_size();
       _M_k.resize(_M_size.x * _M_size.y);
       fill(_M_k.begin(), _M_k.end(), 0.0f);
       _M_zdt.resize(_M_size.x * _M_size.y);
       auto src = dynamic_cast<const C2DBitImage&>(wimage);
       auto pixel = src.begin();
       auto k = _M_zdt.begin();
       SParabola p = {0, 0, -numeric_limits<float>::max(), 0.0f};

       for (size_t i = 0; i < _M_zdt.size(); ++i, ++pixel, ++k) {
              p.fv = *pixel ? 0.0f : g_far;
              k->push_back(p);
       }
}


void C3DDT::read(const C2DImage& wimage, short q)
{
       TRACE_FUNCTION;
       const C2DBitImage *src = dynamic_cast<const C2DBitImage *>(&wimage);

       if (!src)
              throw runtime_error("input image is not a bit image!");

       if (src->size() != _M_zdt.size())
              throw runtime_error("input image has different size!");

       auto si = src->begin();
       auto ei = src->end();
       auto k = _M_k.begin();
       auto p = _M_zdt.begin();

       while (si != ei) {
              float f = *si ? 0.0f : g_far;
              SParabola& parabola = (*p)[*k];
              float s  = d (f, q, parabola.fv, parabola.v);

              while (s <= parabola.z) {
                     --(*k);
                     parabola = (*p)[*k];
                     s  = d (f, q, parabola.fv, parabola.v);
              }

              ++(*k);

              if (*k > (int)p->size()) {
                     cverr() << "k = " << *k << " but p->size() = " << p->size() << "\n";
                     assert(0 && "can't do");
              }

              SParabola new_p = {*k, q, s, f};

              if ( *k == (int)p->size() ) {
                     p->push_back(new_p);
              } else {
                     (*p)[*k] = new_p;

                     if (*k < (int)p->size() - 1)
                            p->resize(*k + 1);
              }

              ++si;
              ++k;
              ++p;
       }
}

void C3DDT::dt2d(C2DFImage& image)const
{
       TRACE_FUNCTION;
       vector<float> buffer(image.get_size().x);

       for (size_t y = 0; y < image.get_size().y; ++y) {
              image.get_data_line_x(y, buffer);
              dt1d(buffer);
              image.put_data_line_x(y, buffer);
       }

       buffer.resize(image.get_size().y);

       for (size_t x = 0; x < image.get_size().x; ++x) {
              image.get_data_line_y(x, buffer);
              dt1d(buffer);
              image.put_data_line_y(x, buffer);
       }
}

void C3DDT::get_slice(size_t s, const C2DImage& image, vector<pair<C3DBounds, float>>& result) const
{
       TRACE_FUNCTION;
       const C2DBitImage *src = dynamic_cast<const C2DBitImage *>(&image);

       if (!src) {
              throw create_exception<invalid_argument>( "input image ", s, "not of type bit");
       }

       if (src->get_size() != _M_size) {
              throw create_exception<invalid_argument>( "input image ", s, "has a dffernt size then reference");
       }

       C2DFImage slice_tmp(_M_size);
       auto i = slice_tmp.begin();
       auto e = slice_tmp.end();
       auto p = _M_zdt.begin();

       while (i != e) {
              size_t k = 0;

              while ( k < p->size() - 1 && (*p)[ k + 1 ].z  < s) {
                     ++k;
              }

              float delta = float(s) - (*p)[k].v;
              *i = delta * delta + (*p)[k].fv;
              ++i;
              ++p;
       }

       dt2d(slice_tmp);
       auto isrc = src->begin();
       auto iref = slice_tmp.begin();

       for (size_t y = 0; y < _M_size.y; ++y)
              for (size_t x = 0; x < _M_size.x; ++x, ++iref, ++isrc)
                     if (*isrc)
                            result.push_back(pair<C3DBounds, float>(C3DBounds(x, y, s), sqrt(*iref)));
}

bool C3DDT::test_initialise(const C2DBounds& size)
{
       if (_M_zdt.size() != size.x * size.y) {
              cverr() << " buffer size should be " << size.x *size.y
                      << " but is " << _M_zdt.size() << "\n";
              return false;
       }

       auto ki = _M_zdt.begin();
       auto ke = _M_zdt.end();
       bool result = true;

       while (ki != ke) {
              if (ki->size() != 1) {
                     cverr() << "not all slots are properly initialised in _M_zdt\n";
                     return false;
              }

              const SParabola& p = (*ki)[0];

              if (p.k != 0) {
                     cverr() << "got a wrong k-initializer\n";
                     result = false;
              }

              if (p.v != 0) {
                     cverr() << "got a wrong v-initializer\n";
                     result = false;
              }

              if (p.z != -numeric_limits<float>::max()) {
                     cverr() << "got a wrong z-initializer\n";
                     result = false;
              }

              if ( ki == _M_zdt.begin() ) {
                     if (p.fv != 0.0 ) {
                            cverr() << "got a wrong (0,0,0) fv-initializer\n";
                            result = false;
                     }
              } else {
                     if (p.fv != g_far) {
                            cverr() << "got a wrong general fv-initializer: " << p.fv << " != " << g_far << "\n";
                            result = false;
                     }
              }

              ++ki;
       }

       return result;
}

/*
static void print_diff(const C2DUSImage& image, unsigned short *ref)
{
	size_t k = 0;
	for (C2DUSImage::const_iterator i = image.begin();
	     i != image.end(); ++i, ++k)
		cverr() << *i << " vs. ref:" << ref[k] << "\n";
}
*/

bool C3DDT::test()
{
       float in_1d[16]  = { 1, 1, 1, 0,
                            1, 1, 1, 1,
                            1, 1, 0, 1,
                            1, 1, 1, 1
                          };
       float out_1d[16] = { 9, 4, 1, 0,
                            1, 4, 9, 9,
                            4, 1, 0, 1,
                            4, 9, 16, 25
                          };
       vector<float> src(16);
       transform(&in_1d[0], &in_1d[16], src.begin(),
       [](float x) {
              return numeric_limits<float>::max() * x;
       });
       dt1d(src);

       if (!equal(src.begin(), src.end(), &out_1d[0])) {
              cverr() << "1D distance transform failed\n";

              for (size_t i = 0; i < 16; ++i) {
                     cverr() << out_1d[i] << " vs " << src[i] << "\n";
              }

              return false;
       }

       float in_2d[16] =  { 1, 1, 1, 0,
                            1, 1, 1, 1,
                            1, 1, 0, 1,
                            1, 1, 1, 1
                          };
       float out_2d[16] = { 8, 4, 1, 0,
                            5, 2, 1, 1,
                            4, 1, 0, 1,
                            5, 2, 1, 2
                          };
       C2DFImage src_img(C2DBounds(4, 4));
       transform(&in_2d[0], &in_2d[16], src_img.begin(), [](float x) {
              return numeric_limits<float>::max() * x;
       });
       dt2d(src_img);

       if (!equal(src_img.begin(), src_img.end(), &out_2d[0])) {
              cverr() << "2D distance transform failed\n";
              return false;
       }

       bool in_3d_1[9] = { 1, 0, 0, 0, 0, 0, 0, 0, 0};
       bool in_3d_2[9] = { 0, 0, 0, 0, 0, 0, 0, 0, 0};
       /*
       float  out_3d_1[9] = { 0,          1,           2,
       		       1, sqrt(2.0f),  sqrt(5.0f),
       		       2, sqrt(5.0f), sqrt(8.0f)};

       float  out_3d_2[9] = {          1, sqrt(2.0f), sqrt(5.0f),
       		       sqrt(2.0f), sqrt(3.0f), sqrt(6.0f),
       		             0.0f, sqrt(6.0f),         3};
       */
       vector<pair<C3DBounds, float>> ref_result;
       ref_result.push_back(pair<C3DBounds, float>(C3DBounds(0, 0, 0), 0));
       ref_result.push_back(pair<C3DBounds, float>(C3DBounds(1, 0, 0), 1));
       ref_result.push_back(pair<C3DBounds, float>(C3DBounds(2, 0, 0), 2));
       ref_result.push_back(pair<C3DBounds, float>(C3DBounds(0, 1, 0), 1));
       ref_result.push_back(pair<C3DBounds, float>(C3DBounds(1, 1, 0), sqrt(2.0f)));
       ref_result.push_back(pair<C3DBounds, float>(C3DBounds(2, 1, 0), sqrt(5.0f)));
       ref_result.push_back(pair<C3DBounds, float>(C3DBounds(0, 2, 0), 2));
       ref_result.push_back(pair<C3DBounds, float>(C3DBounds(1, 2, 0), sqrt(5.0f)));
       ref_result.push_back(pair<C3DBounds, float>(C3DBounds(2, 2, 0), sqrt(8.0f)));
       //ref_result.push_back(pair<C3DBounds, float>(C3DBounds(0,0,1), 0));
       ref_result.push_back(pair<C3DBounds, float>(C3DBounds(1, 0, 1), sqrt(2.0f)));
       ref_result.push_back(pair<C3DBounds, float>(C3DBounds(2, 0, 1), sqrt(5.0f)));
       ref_result.push_back(pair<C3DBounds, float>(C3DBounds(0, 1, 1), sqrt(2.0f)));
       ref_result.push_back(pair<C3DBounds, float>(C3DBounds(1, 1, 1), sqrt(3.0f)));
       ref_result.push_back(pair<C3DBounds, float>(C3DBounds(2, 1, 1), sqrt(6.0f)));
       //ref_result.push_back(pair<C3DBounds, float>(C3DBounds(0,2,1), 0));
       ref_result.push_back(pair<C3DBounds, float>(C3DBounds(1, 2, 1), sqrt(6.0f)));
       ref_result.push_back(pair<C3DBounds, float>(C3DBounds(2, 2, 1), 3));
       bool ref_3d_1[9] = { 1, 1, 1, 1, 1, 1, 1, 1, 1};
       bool ref_3d_2[9] = { 0, 1, 1, 1, 1, 1, 0, 1, 1};
       C2DBounds Size2D(3, 3);
       C2DBitImage *src1 = new C2DBitImage(Size2D);
       C2DBitImage *src2 = new C2DBitImage(Size2D);
       C2DBitImage *ref1 = new C2DBitImage(Size2D);
       C2DBitImage *ref2 = new C2DBitImage(Size2D);
       copy(&in_3d_1[0], &in_3d_1[9], src1->begin());
       copy(&in_3d_2[0], &in_3d_2[9], src2->begin());
       copy(&ref_3d_1[0], &ref_3d_1[9], ref1->begin());
       copy(&ref_3d_2[0], &ref_3d_2[9], ref2->begin());
       P2DImage s1(src1);
       P2DImage s2(src2);
       P2DImage wref1(ref1);
       P2DImage wref2(ref2);
       initialise(*s1);

       if (!test_initialise(Size2D)) {
              cverr() << "initialisation failed!\n";
              return false;
       }

       read(*s2, 1);
       read(*s1, 2);
       read(*s2, 3);
       read(*s1, 4);
       // should implement another test
       vector<pair<C3DBounds, float>> result;
       get_slice(0, *wref1, result);
       get_slice(1, *wref2, result);

       if (!equal(result.begin(), result.end(), ref_result.begin())) {
              cverr() << "part 3 failed\n";
              return false;
       }

       return true;
}


int do_main( int argc, char *argv[] )
{
       string src_filename;
       string ref_filename;
       string out_filename;
       string out_type;
       bool self_test = false;
       CCmdOptionList options(g_description);
       options.set_group("File IO");
       options.add(make_opt( src_filename, "in-file", 'i',
                             "input image(s) that contain the source pixel mask",
                             CCmdOptionFlags::required_input));
       options.add(make_opt( ref_filename, "ref-file", 'r',
                             "reference mask to evaluate the distance from",
                             CCmdOptionFlags::required_input));
       options.add(make_opt( out_filename, "out-file", 'o', "output file name",
                             CCmdOptionFlags::required_output));
       options.set_group("Test");
       options.add(make_opt( self_test, "self-test", 0, "run a self test" ));

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

       if (self_test) {
              C3DDT dt;
              return  dt.test() ? EXIT_SUCCESS : EXIT_FAILURE;
       }

       size_t start_filenum = 0;
       size_t end_filenum  = 0;
       size_t format_width = 0;
       string ref_basename = get_filename_pattern_and_range(ref_filename, start_filenum,
                             end_filenum, format_width);

       if (start_filenum >= end_filenum)
              throw invalid_argument(string("no files match pattern ") + ref_basename);

       //char new_line = cverb.show_debug() ? '\n' : '\r';
       bool first = true;
       C3DDT dt;
       cvdebug() << "Read reference images\n";

       for (size_t i = start_filenum, k = 0; i < end_filenum; ++i, ++k) {
              string src_name = create_filename(ref_basename.c_str(), i);
              cvmsg() << "\rRead " << src_name;
              auto in_image = load_image2d(src_name);

              if (!in_image) {
                     cverr() << "expected " << end_filenum - start_filenum <<
                             " images, got only" << i - start_filenum << "\n";
                     break;
              }

              if (first) {
                     dt.initialise(*in_image);
                     first = false;
              } else
                     dt.read(*in_image, k);
       }

       cvmsg() << "\n";
       size_t src_start_filenum = 0;
       size_t src_end_filenum  = 0;
       string src_basename = get_filename_pattern_and_range(src_filename, src_start_filenum, src_end_filenum, format_width);

       if (src_start_filenum >= src_end_filenum)
              throw invalid_argument(string("no files match pattern ") + src_basename);

       if (src_start_filenum != start_filenum || src_end_filenum != end_filenum)
              throw invalid_argument(string("reference range and in-range are not equal"));

       vector<pair<C3DBounds, float>> result;

       for (size_t i = start_filenum, k = 0; i < end_filenum; ++i, ++k) {
              string src_name = create_filename(src_basename.c_str(), i);
              cvmsg() << "\rRead " << src_name;
              auto in_image = load_image2d(src_name);

              if (in_image) {
                     dt.get_slice(k, *in_image, result);
              }
       }

       cvmsg() << "\n";
       ofstream os( out_filename.c_str(), ios::out );

       for (vector<pair<C3DBounds, float>>::const_iterator i = result.begin();
            i != result.end(); ++i)
              os << i->first.x << " " << i->first.y << " " << i->first.z << " " << i->second << '\n';

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