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

NS_MIA_BEGIN
using std::vector;
using std::copy;
using std::numeric_limits;

struct F3DDistanceTransform : public TFilter <C3DFImage> {
       template <typename T>
       C3DFImage operator () ( const T3DImage<T>& f) const;
};

template <typename T>
C3DFImage F3DDistanceTransform::operator () ( const T3DImage<T>& image) const
{
       C3DFImage result(image.get_size(), image);
       vector<float> buffer(image.get_size().x);
       vector<T> in_buffer(image.get_size().x);

       for (size_t z = 0; z < image.get_size().z; ++z) {
              for (size_t y = 0; y < image.get_size().y; ++y) {
                     image.get_data_line_x(y, z, in_buffer);
                     copy(in_buffer.begin(), in_buffer.end(), buffer.begin());
                     distance_transform_inplace(buffer);
                     result.put_data_line_x(y, z, buffer);
              }
       }

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

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

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

       for (size_t y = 0; y < image.get_size().y; ++y) {
              for (size_t x = 0; x < image.get_size().x; ++x) {
                     result.get_data_line_z(x, y, buffer);
                     distance_transform_inplace(buffer);
                     result.put_data_line_z(x, y, buffer);
              }
       }

       return result;
}

C3DFImage EXPORT_3D distance_transform(const C3DImage& f)
{
       F3DDistanceTransform dtf;
       return mia::filter(dtf, f);
}



struct C3DDistanceImpl {
       C3DDistanceImpl(const C2DImage& slice);

       void push_slice(int z, const C2DImage& slice);

       float get_distance_at(const C3DFVector& p) const;

       C2DFImage get_distance_slice(int z) const;

       struct SParabola {
              int v;
              float z;
              float fv;
       };

       C2DBounds m_size;
       vector<int> m_k;
       vector< vector<SParabola>> m_zdt;

       struct FSlicePusher: public TFilter<void> {
              FSlicePusher(vector<int>& k, vector< vector<SParabola>>& zdt):
                     m_k(k), m_zdt(zdt), m_q(-1)
              {
              }

              template <typename T>
              void operator ()(const T2DImage<T>& f);

              void set_slice(int z);

       private:
              vector<int>& m_k;
              vector< vector<SParabola>>& m_zdt;
              int m_q;
       };

       struct FPushFirstSlice: public TFilter<void> {
              FPushFirstSlice(vector< vector<SParabola>>& zdt):
                     m_zdt(zdt)
              {
              }

              template <typename T>
              void operator ()(const T2DImage<T>& f);
              vector< vector<SParabola>>& m_zdt;
       };
private:
       float get_local_distance(unsigned int x, unsigned  int y, const C3DFVector& p)const;

       FSlicePusher m_pusher;
};

inline float d(float fp, float p, float fq, float q)
{
       return  ( fp  - fq + p * p - q * q) / (p - q) * 0.5;
}

void C3DDistanceImpl::FSlicePusher::set_slice(int z)
{
       m_q = z;
}

template <typename T>
void C3DDistanceImpl::FSlicePusher::operator ()(const T2DImage<T>& func)
{
       assert(m_q > 0);
       auto si = func.begin();
       auto ei = func.end();
       auto k = m_k.begin();
       auto p = m_zdt.begin();
       cvdebug() << "Start slice " << m_q << "\n";

       while (si != ei) {
              const float fq = *si;

              // if the function is at inf, there is no contribution
              if (fq < numeric_limits<float>::max()) {
                     SParabola& parabola = (*p)[*k];
                     cvdebug() << "f= " << fq
                               << ", k=" << *k
                               << ", mq=" << m_q
                               << ", fv=" << parabola.fv
                               << ", pv=" << parabola.v
                               << ", z=" << parabola.z
                               << ", p.size()=" << p->size()
                               << "\n";
                     float s  = d (fq, m_q, parabola.fv, parabola.v);

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

                     ++(*k);

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

                     SParabola new_p = {m_q, s, fq};

                     if ( *k == (int)p->size() ) {
                            cvdebug() << "add parbola at " << *k << " {q=" << m_q << ", z=" << s << ", fq=" << fq << "}\n";
                            p->push_back(new_p);
                     } else {
                            cvdebug() << "set parbola at " << *k << " {q=" << m_q << ", z=" << s << ", fq=" << fq << "}\n";
                            (*p)[*k] = new_p;

                            if (*k < (int)p->size() - 1) {
                                   cvinfo() << "C3DDistance::FSlicePusher: reducing from " << p->size() << " to " << *k << "\n";
                                   cvinfo() << "C3DDistance::FSlicePusher::operator: reducing column size should not happen\n";
                                   p->resize(*k + 1);
                            }
                     }
              }

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

       m_q = -1;
}

template <typename T>
void C3DDistanceImpl::FPushFirstSlice::operator ()(const T2DImage<T>& func)
{
       auto si = func.begin();
       auto ei = func.end();
       auto p = m_zdt.begin();

       while (si != ei) {
              SParabola parabola{0, -numeric_limits<float>::max(), static_cast<float>(*si)};
              p->push_back(parabola);
              ++si;
              ++p;
       }
}



C3DDistance::C3DDistance():
       impl(NULL)
{
}

C3DDistance::~C3DDistance()
{
       delete impl;
}


void C3DDistance::push_slice(int z, const C2DImage& slice)
{
       if (z == 0)
              impl = new C3DDistanceImpl(slice);
       else {
              assert(impl);
              impl->push_slice(z, slice);
       }
}

float C3DDistance::get_distance_at(const C3DFVector& p) const
{
       assert(impl);
       return impl->get_distance_at(p);
}


C3DDistanceImpl::C3DDistanceImpl(const C2DImage& slice):
       m_size(slice.get_size()),
       m_k(m_size.product(), 0),
       m_zdt(m_size.product()),
       m_pusher(m_k, m_zdt)
{
       FPushFirstSlice pfs(m_zdt);
       mia::accumulate(pfs, slice);
}


void C3DDistanceImpl::push_slice(int z, const C2DImage& slice)
{
       m_pusher.set_slice(z);
       mia::accumulate(m_pusher, slice);
}

float C3DDistanceImpl::get_local_distance(unsigned int x, unsigned  int y, const C3DFVector& p)const
{
       if (x >= m_size.x || y >= m_size.y)
              return numeric_limits<float>::max();

       auto&  zdt = m_zdt[y * m_size.x + x];

       if (zdt.size() == 1 && zdt[0].fv == numeric_limits<float>::max())
              return numeric_limits<float>::max();

       unsigned int k = 0;

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

       const float delta = p.z - zdt[k].v;
       const float dx = p.x - x;
       const float dy = p.y - y;
       return dx * dx + dy * dy + delta * delta + zdt[k].fv;
}

float C3DDistanceImpl::get_distance_at(const C3DFVector& p) const
{
       float distance = numeric_limits<float>::max();
       int center_x = (int)(p.x + 0.5);
       int center_y = (int)(p.y + 0.5);

       // limit search to support range
       if (center_x >= (int)m_size.x - 1)
              center_x = m_size.x - 1;

       if (center_x < 0)
              center_x = 0;

       if (center_y >= (int)m_size.y - 1)
              center_y = m_size.y - 1;

       if (center_y < 0)
              center_y = 0;

       int max_delta_x = m_size.x - center_x;

       if (max_delta_x < center_x)
              max_delta_x = center_x;

       int max_delta_y = m_size.y - center_y;

       if (max_delta_y < center_y)
              max_delta_y = center_y;

       const int max_line = sqrt(max_delta_y * max_delta_y +  max_delta_x * max_delta_x) + 1;
       int x = center_x;
       int y = center_y;
       float d = get_local_distance(x, y, p);

       if (d < distance)
              distance = d;

       // This breaking condition is sub-optimal, it should actually break at distance > r^2, since distance is squared,
       // but because we move on a grid, this doesn't work properly.
       // In other words: the search could be optimized by refining this breaking condition.
       for ( int r = 1; r < max_line && distance > r; ++r) {
              for (int d_outer = 0; (d_outer < r); ++d_outer) {
                     const int d_inner = sqrt(r * r - d_outer * d_outer);
                     x = center_x + d_outer;
                     y = center_y + d_inner;
                     d = get_local_distance(x, y, p);

                     if (d < distance)
                            distance = d;

                     x = center_x + d_outer;
                     y = center_y - d_inner;
                     d = get_local_distance(x, y, p);

                     if (d < distance)
                            distance = d;

                     if (d_outer != 0) {
                            x = center_x - d_outer;
                            y = center_y + d_inner;
                            d = get_local_distance(x, y, p);

                            if (d < distance)
                                   distance = d;

                            x = center_x - d_outer;
                            y = center_y - d_inner;
                            d = get_local_distance(x, y, p);

                            if (d < distance)
                                   distance = d;
                     }

                     if (d_inner != d_outer) {
                            x = center_x + d_inner;
                            y = center_y + d_outer;
                            d = get_local_distance(x, y, p);

                            if (d < distance)
                                   distance = d;

                            x = center_x - d_inner;
                            y = center_y + d_outer;
                            d = get_local_distance(x, y, p);

                            if (d < distance)
                                   distance = d;

                            if (d_outer != 0) {
                                   x = center_x + d_inner;
                                   y = center_y - d_outer;
                                   d = get_local_distance(x, y, p);

                                   if (d < distance)
                                          distance = d;

                                   x = center_x - d_inner;
                                   y = center_y - d_outer;
                                   d = get_local_distance(x, y, p);

                                   if (d < distance)
                                          distance = d;
                            }
                     }
              }
       }

       return sqrt(distance);
}

C2DFImage C3DDistanceImpl::get_distance_slice(int z) const
{
       C2DFImage result(m_size);
       auto i = result.begin();
       auto e = result.end();
       auto p = m_zdt.begin();

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

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

              float delta = float(z) - (*p)[k].v;
              *i = delta * delta + (*p)[k].fv;
              cvdebug() << "get parbola at " << k << " {z=" << (*p)[k].z << ", fq=" << (*p)[k].fv << "}\n";
              cvdebug() << "z=" << z << ", k= " << k << ", dist=" << *i << "\n";
              ++i;
              ++p;
       }

       return distance_transform(result);
}

C2DFImage C3DDistance::get_distance_slice(int z) const
{
       assert(impl);
       return impl->get_distance_slice(z);
}
NS_MIA_END
