/* -*- 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 <mia/core/msgstream.hh>
#include <mia/2d/transform/vectorfield.hh>
#include <mia/2d/vfio.hh>

NS_MIA_BEGIN
using namespace std;

C2DGridTransformation::C2DGridTransformation(const C2DBounds& size, const C2DInterpolatorFactory& ipf):
       C2DTransformation(ipf),
       m_field(size),
       m_upscale_interpolator_factory("bspline:d=1", "zero")
{
}

P2DTransformation C2DGridTransformation::do_upscale(const C2DBounds& size) const
{
       /* This implementation could be improved by using something like the spline interpolator
          that applies the scaling in a separable way. */
       TRACE("C2DGridTransformation::upscale");
       DEBUG_ASSERT_RELEASE_THROW(m_field.get_size().x != 0 && m_field.get_size().y != 0,
                                  "C2DGridTransformation::do_upscale: input field has a zero dimension");
       C2DGridTransformation *result = new C2DGridTransformation(size, get_interpolator_factory());
       unique_ptr<T2DInterpolator<C2DFVector>>
                                            interp(m_upscale_interpolator_factory.create(m_field));
       float x_mult = float(size.x) / (float)m_field.get_size().x;
       float y_mult = float(size.y) / (float)m_field.get_size().y;
       float ix_mult = 1.0f / x_mult;
       float iy_mult = 1.0f / y_mult;
       auto i = result->m_field.begin();

       for (unsigned int y = 0; y < size.y; y++) {
              for (unsigned int x = 0; x < size.x; x++, ++i) {
                     C2DFVector help(ix_mult * x, iy_mult * y);
                     C2DFVector val = (*interp)(help);
                     *i = C2DFVector(val.x * x_mult, val.y * y_mult);
              }
       }

       return P2DTransformation(result);
}

const C2DBounds& C2DGridTransformation::get_size() const
{
       return m_field.get_size();
}

void C2DGridTransformation::update(float step, const C2DFVectorfield& a)
{
       C2DFVectorfield::const_iterator inf = a.begin();
       C2DFVectorfield::const_iterator enf = a.end();
       C2DFVectorfield::iterator onf = m_field.begin();

       while (inf != enf)
              *onf++ += step * *inf++;
}

C2DTransformation *C2DGridTransformation::do_clone() const
{
       return new C2DGridTransformation(*this);
}

C2DTransformation *C2DGridTransformation::invert() const
{
       assert(0 && "not implemented");
       return new C2DGridTransformation(*this);
}


size_t C2DGridTransformation::degrees_of_freedom() const
{
       return 2 * m_field.size();
}


C2DFMatrix C2DGridTransformation::derivative_at(int x, int y) const
{
       C2DFMatrix result(C2DFVector(1.0, 0), C2DFVector(0, 1.0));
       result -= field_derivative_at(x, y);
       return result;
}

C2DFMatrix C2DGridTransformation::derivative_at(const C2DFVector& x) const
{
       C2DFMatrix result = C2DFMatrix::_1;
       const int ix = static_cast<int>(floor(x.x));
       const int iy = static_cast<int>(floor(x.y));
       const int sx = m_field.get_size().x;
       const int sy = m_field.get_size().y;

       if (ix > 0 && ix < sx - 2 &&
           iy > 0 && iy < sy - 2) {
              const float dx = x.x - ix;
              const float dy = x.y - iy;
              const float fx = 1.0 - dx;
              const float fy = 1.0 - dy;
              C2DFMatrix m[4];
              m[0] = field_derivative_at(ix, iy    );
              m[1] = field_derivative_at(ix + 1, iy    );
              m[2] = field_derivative_at(ix, iy + 1);
              m[3] = field_derivative_at(ix + 1, iy + 1);
              result -= fy * (fx * m[0] + dx * m[1]) +
                        dy * (fx * m[2] + dx * m[3]);
       }

       return result;
}

C2DFMatrix C2DGridTransformation::field_derivative_at(int x, int y) const
{
       C2DFMatrix result;
       const int sx = m_field.get_size().x;
       const int sy = m_field.get_size().y;

       if (y >= 0 && y < sy ) {
              if (x >= 0 && x < sx ) {
                     C2DFVectorfield::const_pointer center = &m_field[sx * y + x];

                     if (x > 0 && x < sx - 1) {
                            result.x = (center[ 1] - center[-1]) * 0.5f;
                     }

                     if (y > 0 && y < sy - 1 ) {
                            result.y = (center[ sx] - center[-sx]) * 0.5f;
                     }
              }
       }

       return result;
}

void C2DGridTransformation::set_identity()
{
       fill(m_field.begin(), m_field.end(), C2DFVector(0, 0));
}

float C2DGridTransformation::get_max_transform() const
{
       C2DFVectorfield::const_iterator i = m_field.begin();
       C2DFVectorfield::const_iterator e = m_field.end();
       assert(i != e);
       float value = i->norm2();
       ++i;

       while (i != e) {
              float v = i->norm2();

              if (value < v)
                     value = v;

              ++i;
       }

       return sqrt(value);
}


CDoubleVector C2DGridTransformation::get_parameters() const
{
       CDoubleVector result(m_field.size() * 2);
       auto r = result.begin();

       for (auto f = m_field.begin(); f != m_field.end(); ++f) {
              *r++ = f->x;
              *r++ = f->y;
       }

       return result;
}

void C2DGridTransformation::set_parameters(const CDoubleVector& params)
{
       assert(2 * m_field.size() == params.size());
       auto r = params.begin();

       for (auto f = m_field.begin(); f != m_field.end(); ++f) {
              f->x = *r++;
              f->y = *r++;
       }
}


C2DGridTransformation::iterator_impl::iterator_impl(const C2DBounds& pos, const C2DBounds& size,
              C2DFVectorfield::const_iterator start):
       C2DTransformation::iterator_impl(pos, size),
       m_current(start)
{
       if (pos.y < size.y && pos.x < size.x )
              m_value = C2DFVector(get_pos()) - *m_current;
}

C2DTransformation::iterator_impl *C2DGridTransformation::iterator_impl::clone() const
{
       return new C2DGridTransformation::iterator_impl(get_pos(), get_size(), m_current);
}

const C2DFVector&  C2DGridTransformation::iterator_impl::do_get_value()const
{
       return m_value;
}

void C2DGridTransformation::iterator_impl::do_x_increment()
{
       ++m_current;
       m_value = C2DFVector(get_pos()) - *m_current;
}

void C2DGridTransformation::iterator_impl::do_y_increment()
{
       ++m_current;
       m_value = C2DFVector(get_pos()) - *m_current;
}


C2DFVector C2DGridTransformation::operator ()(const  C2DFVector& x) const
{
       return x - get_displacement_at(x);
}

C2DGridTransformation::const_iterator C2DGridTransformation::begin() const
{
       return const_iterator(new iterator_impl(C2DBounds(0, 0),
                                               m_field.get_size(), m_field.begin()));
}

C2DGridTransformation::const_iterator C2DGridTransformation::end() const
{
       return const_iterator(new iterator_impl( m_field.get_size(), m_field.get_size(),
                             m_field.end()));
}


C2DGridTransformation::field_iterator C2DGridTransformation::field_begin()
{
       return m_field.begin();
}

C2DGridTransformation::field_iterator C2DGridTransformation::field_end()
{
       return m_field.end();
}

C2DGridTransformation::const_field_iterator C2DGridTransformation::field_begin() const
{
       return m_field.begin();
}

C2DGridTransformation::const_field_iterator C2DGridTransformation::field_end()const
{
       return m_field.end();
}

void C2DGridTransformation::translate(const C2DFVectorfield& gradient, CDoubleVector& params) const
{
       assert(2 * params.size() != gradient.size());
       // translating the gradient also means multiplication with -1
       // because the transformation is I-u(x)
       auto r = params.begin();

       for (auto f = gradient.begin(); f != gradient.end(); ++f, r += 2) {
              r[0] = -f->x;
              r[1] = -f->y;
       }
}

float C2DGridTransformation::pertuberate(C2DFVectorfield& v) const
{
       C2DFVectorfield::iterator iv = v.begin();
       float max_gamma = 0.0f;

       for (size_t y = 0; y < v.get_size().y; ++y)
              for (size_t x = 0; x < v.get_size().x; ++x, ++iv) {
                     const C2DFMatrix j = field_derivative_at(x, y);
                     const C2DFVector ue = j * *iv;
                     *iv -= ue;
                     float gamma = iv->norm2();

                     if (gamma > max_gamma)
                            max_gamma = gamma;
              }

       return sqrt(max_gamma);
}

float C2DGridTransformation::get_jacobian(const C2DFVectorfield& v, float delta) const
{
       assert(v.get_size() == get_size());
       const int dx = v.get_size().x;
       float j_min = numeric_limits<float>::max();

       for (size_t y = 1; y < get_size().y - 1; ++y) {
              C2DFVectorfield::const_iterator iu = m_field.begin_at(1, y);
              C2DFVectorfield::const_iterator iv = v.begin_at(1, y);

              for (size_t x = 1; x < get_size().x - 1; ++x, ++iu, ++iv) {
                     C2DFVector jx((iu[1] - iu[-1]) + (delta * (iv[1] - iv[-1])));
                     C2DFVector jy((iu[dx] - iu[-dx]) + (delta * (iv[dx] - iv[-dx])));
                     const float j = (2.0 - jx.x) * (2.0 - jy.y) - jx.y * jy.x;

                     if ( j_min > j) {
                            j_min = j;
                     }
              }
       }

       return 0.25 * j_min;
}

EXPORT_2D C2DGridTransformation operator + (const C2DGridTransformation& a, const C2DGridTransformation& b)
{
       assert( a.get_size() == b.get_size());
       C2DGridTransformation result(a.get_size(), a.get_interpolator_factory());
       C2DFVectorfield::iterator ri = result.field_begin();
       C2DFVectorfield::const_iterator bi = b.field_begin();

       for (size_t y = 0; y < a.get_size().y; ++y)  {
              for (size_t x = 0; x < a.get_size().x; ++x, ++ri, ++bi)  {
                     const C2DFVector xi = C2DFVector(x, y) - *bi;
                     *ri = a.get_displacement_at(xi) +  *bi;
              }
       }

       return result;
}

C2DGridTransformCreator::C2DGridTransformCreator(const C2DInterpolatorFactory& ipf):
       C2DTransformCreator(ipf)
{
}

P2DTransformation C2DGridTransformCreator::do_create(const C2DBounds& size, const C2DInterpolatorFactory& ipf) const
{
       return P2DTransformation(new C2DGridTransformation(size, ipf));
}



/**
   Plugin class to create the creater.
 */
C2DGridTransformCreatorPlugin::C2DGridTransformCreatorPlugin():
       C2DTransformCreatorPlugin("vf")
{
}

C2DTransformCreator *C2DGridTransformCreatorPlugin::do_create(const C2DInterpolatorFactory& ipf) const
{
       return new C2DGridTransformCreator(ipf);
}

const std::string C2DGridTransformCreatorPlugin::do_get_descr() const
{
       return "This plug-in implements a transformation that defines a translation for "
              "each point of the grid defining the domain of the transformation.";
}


NS_MIA_END
