/* -*- 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 <limits>
#include <queue>
#include <ostream>
#include <fstream>
#include <set>
#include <map>
#include <ctime>
#include <mia/3d/fifof/label.hh>


NS_BEGIN(label_2dstack_filter)

NS_MIA_USE;
using namespace std;

C2DLabelStackFilter::C2DLabelStackFilter(const string& mapfile, P2DShape n):
       C2DImageFifoFilter(1, 1, 0),
       m_neigbourhood(n),
       m_start_label(1),
       m_last_label(1),
       m_map_file(mapfile),
       m_first_pass(true)
{
}

C2DLabelStackFilter::~C2DLabelStackFilter()
{
}

static int slice = 0;
void C2DLabelStackFilter::do_initialize(::boost::call_traits<mia::P2DImage>::param_type x)
{
       m_first_pass = true;
       m_joints.clear();
       m_out_buffer = C2DUIImage(x->get_size());
       m_last_label = m_start_label;
}

void C2DLabelStackFilter::set_start_label(unsigned int start_label)
{
       m_start_label = start_label;
}


void C2DLabelStackFilter::grow( int x, int y, C2DBitImage& input, unsigned int l)
{
       vector<C2DBounds> seed;
       seed.push_back(C2DBounds(x, y));

       while (!seed.empty()) {
              C2DBounds loc = seed.back();
              seed.pop_back();

              for (auto si = m_neigbourhood->begin();
                   si != m_neigbourhood->end(); ++si) {
                     size_t px = si->x + loc.x;
                     size_t py = si->y + loc.y;

                     if (px >= input.get_size().x ||
                         py >= input.get_size().y)
                            continue;

                     unsigned int lold = m_out_buffer(px, py);

                     if (lold) {
                            if (l != lold) {
                                   m_joints.add_pair(l, lold);
                                   m_out_buffer(px, py) = l;
                            }
                     } else if (input(px, py)) {
                            m_out_buffer(px, py) = l;
                            input(px, py) = false;
                            seed.push_back(C2DBounds(px, py));
                     }
              }
       }
}


void C2DLabelStackFilter::label_new_regions(C2DBitImage& input)
{
       auto ii = input.begin();
       auto usi = m_out_buffer.begin();

       for (size_t y = 0; y < input.get_size().y; ++y)
              for (size_t x = 0; x < input.get_size().x; ++x, ++usi, ++ii) {
                     if (*ii) {
                            cvdebug() << "(" << x << ", " << y << "," << slice <<  "):"
                                      << m_last_label << "\n";

                            if (m_last_label < numeric_limits<unsigned int>::max())
                                   *usi = m_last_label++;
                            else
                                   throw create_exception<invalid_argument>("C2DLabelStackFilter: number of connected components is about to "
                                                 "exeed the  supported limit of ",
                                                 numeric_limits<unsigned int>::max(),
                                                 ", sorry can't continue\n");

                            *ii = false;
                            grow(x, y, input, *usi);
                     }
              }
}

void C2DLabelStackFilter::label(C2DBitImage& input)
{
       // first grow all regions that are already labeled from the last slice
       auto usi = m_out_buffer.begin();

       for (size_t y = 0; y < input.get_size().y; ++y)
              for (size_t x = 0; x < input.get_size().x; ++x, ++usi) {
                     if ( *usi )
                            grow(x, y, input, *usi);
              }

       // then label all not yet connected regions
       label_new_regions(input);
}

void C2DLabelStackFilter::new_label(C2DBitImage& input)
{
       m_out_buffer = C2DUIImage(input.get_size(), input);
       label_new_regions(input);
}

void  C2DLabelStackFilter::re_label(C2DBitImage& input)
{
       auto usi = m_out_buffer.begin();
       auto use = m_out_buffer.end();
       auto ii = input.begin();

       // maintain old labeling for new slice, and clean the input at
       // labeled positions
       while (usi != use)  {
              if (!*ii)
                     *usi = 0;
              else if (*usi)
                     *ii = 0;

              ++ii;
              ++usi;
       }

       // now continue normal labeling
       label(input);
}

void C2DLabelStackFilter::do_push(::boost::call_traits<mia::P2DImage>::param_type x)
{
       const C2DBitImage *inp = dynamic_cast<const C2DBitImage *>(x.get());

       if (!inp)
              throw invalid_argument("C2DLabelStackFilter: input images must be binary");

       C2DBitImage input(*inp);

       if (m_first_pass) {
              new_label(input);
              m_first_pass = false;
       } else {
              re_label(input);
       }

       ++slice;
}

void CLabelRemapper::clear()
{
       m_raw_map.clear();
}

P2DImage C2DLabelStackFilter::do_filter()
{
       return P2DImage(new C2DUIImage(m_out_buffer));
}

void CLabelRemapper::add_pair(unsigned int a, unsigned int b)
{
       if (a > b)
              m_raw_map.insert(T2DVector<unsigned int>(a, b));
       else
              m_raw_map.insert(T2DVector<unsigned int>(b, a));
}

struct greater_than {
       typedef T2DVector<unsigned int> value_type;
       bool operator() (const T2DVector<unsigned int>& lhs,
                        const T2DVector<unsigned int>& rhs)
       {
              return lhs.x > rhs.x || ((lhs.x  == rhs.x)  && lhs.y > rhs.y);
       }
};

CLabelMap CLabelRemapper::get_map() const
{
       CLabelMap result;
       priority_queue<T2DVector<unsigned int>, vector<T2DVector<unsigned int>>, greater_than> sorted;
       cvdebug() << "got " << m_raw_map.size() << " joints\n";

       for (auto i = m_raw_map.begin();  i != m_raw_map.end();  ++i)
              sorted.push(*i);

       while (!sorted.empty()) {
              auto v = sorted.top();
              cvdebug() << "Top = " << v << "\n";
              sorted.pop();
              // first check, if the value we map to is already mapped to a lower number
              // if yes, take this mapping for the new value
              auto m = result.find(v.y);

              if (m != result.end()) {
                     result[v.x] = m->second;
              } else {
                     // now test of the value to be mapped is already available
                     // and if not add the new mapping, if yes, add a mapping for the target
                     m = result.find(v.x);

                     if (m == result.end()) {
                            result[v.x] = v.y;
                     } else {
                            result[v.y] = m->second;
                     }
              }
       }

       return result;
}


void C2DLabelStackFilter::post_finalize()
{
       m_target = m_joints.get_map();

       if (!m_map_file.empty()) {
              ofstream outfile(m_map_file.c_str(), ios_base::out );
              m_target.save(outfile);

              if (!outfile.good()) {
                     throw create_exception<runtime_error>("C2DLabelStackFilter: failed to save labale join map to '",
                                                           m_map_file, "'");
              }
       }
}

const CLabelMap& C2DLabelStackFilter::get_joints() const
{
       return m_target;
}


class C2DLabelFifoFilterPlugin : public C2DFifoFilterPlugin
{
public:
       C2DLabelFifoFilterPlugin();
private:

       virtual const string do_get_descr() const;
       virtual C2DImageFifoFilter *do_create()const;

       mia::P2DShape m_neighbourhood;
       string m_mapfile;
};

C2DLabelFifoFilterPlugin::C2DLabelFifoFilterPlugin():
       C2DFifoFilterPlugin("label")
{
       add_parameter("n", make_param(m_neighbourhood, "4n", false,
                                     "2D neighbourhood shape to define connectedness"));
       add_parameter("map", new CStringParameter(m_mapfile, CCmdOptionFlags::required_input,
                     "Mapfile to save label numbers that are joined"));
}

const string C2DLabelFifoFilterPlugin::do_get_descr() const
{
       return "Stack Label filter";
}

C2DImageFifoFilter *C2DLabelFifoFilterPlugin::do_create()const
{
       return new C2DLabelStackFilter(m_mapfile, m_neighbourhood);
}

extern "C" EXPORT CPluginBase *get_plugin_interface()
{
       return new C2DLabelFifoFilterPlugin();
}


NS_END
