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

#define VSTREAM_DOMAIN "StackDT"

#include <limits>
#include <mia/3d/stackdisttrans.hh>

#include <mia/core/export_handler.hh>

#include <mia/core/errormacro.hh>
#include <mia/core/ioplugin.cxx>
#include <mia/core/iohandler.cxx>


NS_MIA_BEGIN

using namespace std;

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

const char *const C2DStackDistanceTransform::type_descr = "dt";
const char *const C2DStackDistanceTransform::data_descr = "stack2d";


C2DStackDistanceTransform *C2DStackDistanceTransform::clone() const
{
       return new C2DStackDistanceTransform(*this);
}

C2DStackDistanceTransform::C2DStackDistanceTransform(const C2DImage& slice, const C3DFVector& voxel_size):
       m_size(slice.get_size()),
       m_voxel_size(voxel_size),
       m_k(m_size.x * m_size.y, 0),
       m_zdt(m_size.x * m_size.y)
{
       TRACE_FUNCTION;
       const C2DBitImage *src = dynamic_cast<const C2DBitImage *>(&slice);

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

       auto izdt = m_zdt.begin();
       SParabola p = {0, 0, -g_far, g_far};
       auto pixel = src->begin();

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

/*
  Implements Algorithm DT(f) of the paper
*/
void C2DStackDistanceTransform::read( const C2DImage& slice, int q)
{
       TRACE_FUNCTION;
       const C2DBitImage *src = dynamic_cast<const C2DBitImage *>(&slice);

       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!");

       C2DFVector pixel_size = slice.get_pixel_size();
       C2DFVector predev_pixel(m_voxel_size.x,  m_voxel_size.y);

       if ( pixel_size != predev_pixel)
              cvwarn() << "Input image voxel size and predefined voxel size differ: "
                       << pixel_size  <<  " != " << predev_pixel
                       << " the latter takes precedence \n";

       auto si = src->begin();
       auto ei = src->end();
       auto k = m_k.begin();
       auto p = m_zdt.begin();

       // for each point of the input mask run the DT(f) algorithm
       // this takes care of the z-diatance parts
       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);
                     assert(*k >= 0);
                     parabola = (*p)[*k];
                     s  = d (f, q, parabola.fv, parabola.v);
              }

              ++(*k);
              const int psize = p->size();

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

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

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

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

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


vector<C2DStackDistanceTransform::DistanceFromPoint>
C2DStackDistanceTransform::get_slice(size_t s, const C2DImage& image) const
{
       TRACE_FUNCTION;
       vector<C2DStackDistanceTransform::DistanceFromPoint>  result;
       auto src = dynamic_cast<const C2DBitImage&>(image);

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

       C2DFVector pixel_size = image.get_pixel_size();
       C2DFVector predev_pixel(m_voxel_size.x,  m_voxel_size.y);

       if ( pixel_size != predev_pixel)
              cvwarn() << "Input image voxel size and predefined voxel size differ: "
                       << pixel_size  <<  " != " << predev_pixel
                       << " the latter takes precedence \n";

       // evaluate the distance image for this slice
       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) * m_voxel_size.z;
              *i = delta * delta + (*p)[k].fv;
              ++i;
              ++p;
       }

       // run the 2D distance transform folding operation
       dt2d(slice_tmp);
       // now for each point in the input that is not zero, evaluate the
       // distance to the original mask
       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(DistanceFromPoint(C3DBounds(x, y, s), sqrt(*iref)));

       return result;
}

bool C2DStackDistanceTransform::save(const string& /*filename*/)
{
       return false;
}


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

void C2DStackDistanceTransform::dt1d(vector<float>& r, float scale)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]) * scale;
              r[q] = delta * delta + f[v[k]];
       }
}

void C2DStackDistanceTransform::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, m_voxel_size.x);
              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, m_voxel_size.y);
              image.put_data_line_y(x, buffer);
       }
}

template<> const  char *const
TPluginHandler<C2DStackDistanceTransformIO>::m_help =
       "These plug-ins provide support for loading and storing the results of a distance transform of a stack of 2D images."
       ;

template class TIOPlugin<C2DStackDistanceTransform>;
template class THandlerSingleton<TIOPluginHandler<C2DStackDistanceTransformIO>>;
template class TIOPluginHandler<C2DStackDistanceTransformIO>;
template class TPluginHandler<C2DStackDistanceTransformIO>;

NS_MIA_END
