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

#include <mia/core/spacial_kernel.hh>
#include <mia/3d/transform/vectorfield.hh>

NS_MIA_USE

using namespace std;
using namespace ::boost;
using namespace boost::unit_test;

struct GridTransformFixture {
       GridTransformFixture():
              size(64, 32, 48),
              r(size.x - 1, size.y - 1, size.z - 1),
              field(size, C3DInterpolatorFactory("bspline:d=1", "mirror")),
              m_scale(1.0 / r.x, 1.0 / r.y, 1.0 / r.z)
       {
              C3DGridTransformation::field_iterator i = field.field_begin();
              C3DFVector v;

              for (v.z = 0; v.z < size.z; ++v.z)
                     for (v.y = 0; v.y < size.y; ++v.y)
                            for (v.x = 0; v.x < size.x; ++v.x, ++i) {
                                   *i = C3DFVector( fx(v), fy(v), fz(v));
                            }
       }
       C3DBounds size;
       C3DBounds r;
       C3DGridTransformation field;
protected:
       float fx(const C3DFVector& v);
       float fy(const C3DFVector& v);
       float fz(const C3DFVector& v);

       float dfx_x(const C3DFVector& v);
       float dfy_x(const C3DFVector& v);
       float dfz_x(const C3DFVector& v);

       float dfx_y(const C3DFVector& v);
       float dfy_y(const C3DFVector& v);
       float dfz_y(const C3DFVector& v);

       float dfx_z(const C3DFVector& v);
       float dfy_z(const C3DFVector& v);
       float dfz_z(const C3DFVector& v);

       float dfx_xx(float x, float y, float z);
       float dfx_xy(float x, float y, float z);
       float dfx_yx(float x, float y, float z);
       float dfx_yy(float x, float y, float z);
       float dfy_xx(float x, float y, float z);
       float dfy_xy(float x, float y, float z);
       float dfy_yx(float x, float y, float z);
       float dfy_yy(float x, float y, float z);

       C3DFVector scaled(C3DFVector x) const;

       C3DFVector m_scale;
};

C3DFVector GridTransformFixture::scaled(C3DFVector x) const
{
       return m_scale * x;
}

BOOST_AUTO_TEST_CASE(test_vf_plugin)
{
       auto vf_creator = BOOST_TEST_create_from_plugin<C3DGridTransformCreatorPlugin>("vf");
       BOOST_CHECK(vf_creator);
       C3DBounds size(2, 3, 4);
       auto vf = vf_creator->create(size);
       BOOST_CHECK_EQUAL(vf->get_size(), size);
}



BOOST_FIXTURE_TEST_CASE(test_gridtransform_basic_props, GridTransformFixture)
{
       BOOST_CHECK_EQUAL(field.degrees_of_freedom(), size.product()  * 3);
       C3DFVector testx0(22, 22, 23);
       C3DFVector result0 = field.get_displacement_at(testx0);
       BOOST_CHECK_CLOSE(result0.x, fx(testx0), 1);
       BOOST_CHECK_CLOSE(result0.y, fy(testx0), 1);
       BOOST_CHECK_CLOSE(result0.z, fz(testx0), 1);
       C3DFVector testx1(32.5, 22, 23);
       C3DFVector result1 = field.get_displacement_at(testx1);
       BOOST_CHECK_CLOSE(result1.x, fx(testx1), 1);
       BOOST_CHECK_CLOSE(result1.y, fy(testx1), 1);
       BOOST_CHECK_CLOSE(result1.z, fz(testx1), 1);
       C3DFVector testx2(32, 22, 23);
       C3DFVector result2 = field.get_displacement_at(testx2);
       BOOST_CHECK_CLOSE(result2.x, fx(testx2), 1);
       BOOST_CHECK_CLOSE(result2.y, fy(testx2), 1);
       BOOST_CHECK_CLOSE(result2.z, fz(testx2), 1);
       // the non-linearity of the test function is quite high ...
       C3DFVector testx(32.2, 21.9, 22.9);
       C3DFVector result = field.get_displacement_at(testx);
       BOOST_CHECK_CLOSE(result.x, fx(testx), 4);
       BOOST_CHECK_CLOSE(result.y, fy(testx), 4);
       BOOST_CHECK_CLOSE(result.z, fz(testx), 4);
       result = field(testx);
       BOOST_CHECK_CLOSE(result.x, testx.x - fx(testx), 0.1);
       BOOST_CHECK_CLOSE(result.y, testx.y - fy(testx), 0.1);
       BOOST_CHECK_CLOSE(result.z, testx.z - fz(testx), 0.1);
       // 4.25 * 0.25 + 6.25 *0.75
       C3DFVector ip_loc( 2.25, 1.75, 3.2);
       auto ip = field.get_displacement_at(ip_loc);
       BOOST_CHECK_CLOSE(ip.x, fx(ip_loc), 0.02);
       BOOST_CHECK_CLOSE(ip.y, fy(ip_loc), 0.02);
       BOOST_CHECK_CLOSE(ip.z, fz(ip_loc), 0.02);
}

BOOST_FIXTURE_TEST_CASE(test_gridtransform_derivative_1, GridTransformFixture)
{
       C3DFVector v(10, 11, 12);
       C3DFMatrix dv =  field.derivative_at(10, 11, 12);
       BOOST_CHECK_CLOSE(dv.x.x, 1.0f - dfx_x(v), 1);
       BOOST_CHECK_CLOSE(dv.x.y,      - dfy_x(v), 1);
       BOOST_CHECK_CLOSE(dv.x.z,      - dfz_x(v), 1);
       BOOST_CHECK_CLOSE(dv.y.x,      - dfx_y(v), 1);
       BOOST_CHECK_CLOSE(dv.y.y, 1.0f - dfy_y(v), 1);
       BOOST_CHECK_CLOSE(dv.y.z,      - dfz_y(v), 1);
       BOOST_CHECK_CLOSE(dv.z.x,      - dfx_z(v), 1);
       BOOST_CHECK_CLOSE(dv.z.y,      - dfy_z(v), 1);
       BOOST_CHECK_CLOSE(dv.z.z, 1.0f - dfz_z(v), 1);
}

BOOST_FIXTURE_TEST_CASE(test_gridtransform_derivative_at, GridTransformFixture)
{
       C3DFVector v(10.2, 11.4, 12.6);
       C3DFMatrix dv =  field.derivative_at(v);
       BOOST_CHECK_CLOSE(dv.x.x, 1.0f - dfx_x(v), 1);
       BOOST_CHECK_CLOSE(dv.x.y,      - dfy_x(v), 1);
       BOOST_CHECK_CLOSE(dv.x.z,      - dfz_x(v), 1);
       BOOST_CHECK_CLOSE(dv.y.x,      - dfx_y(v), 1);
       BOOST_CHECK_CLOSE(dv.y.y, 1.0f - dfy_y(v), 1);
       BOOST_CHECK_CLOSE(dv.y.z,      - dfz_y(v), 1);
       BOOST_CHECK_CLOSE(dv.z.x,      - dfx_z(v), 1);
       BOOST_CHECK_CLOSE(dv.z.y,      - dfy_z(v), 1);
       BOOST_CHECK_CLOSE(dv.z.z, 1.0f - dfz_z(v), 1);
}

BOOST_FIXTURE_TEST_CASE(test_gridtransform_derivative_at_centered, GridTransformFixture)
{
       C3DFVector v(10.5, 11.5, 12.5);
       C3DFMatrix dv =  field.derivative_at(v);
       BOOST_CHECK_CLOSE(dv.x.x, 1.0f - dfx_x(v), 1);
       BOOST_CHECK_CLOSE(dv.x.y,      - dfy_x(v), 1);
       BOOST_CHECK_CLOSE(dv.x.z,      - dfz_x(v), 1);
       BOOST_CHECK_CLOSE(dv.y.x,      - dfx_y(v), 1);
       BOOST_CHECK_CLOSE(dv.y.y, 1.0f - dfy_y(v), 1);
       BOOST_CHECK_CLOSE(dv.z.x,      - dfx_z(v), 1);
       BOOST_CHECK_CLOSE(dv.y.z,      - dfz_y(v), 1);
       BOOST_CHECK_CLOSE(dv.z.y,      - dfy_z(v), 1);
       BOOST_CHECK_CLOSE(dv.z.z, 1.0f - dfz_z(v), 1);
}


BOOST_FIXTURE_TEST_CASE(test_gridtransform_derivative, GridTransformFixture)
{
       for (size_t z = 0; z < size.z; ++z)
              for (size_t y = 0; y < size.y; ++y)
                     for (size_t x = 0; x < size.x; ++x) {
                            C3DFVector v(x, y, z);
                            C3DFMatrix dv =  field.derivative_at(x, y, z);

                            if ( x > 0 && x < size.x - 1) {
                                   BOOST_CHECK_CLOSE(dv.x.x, 1.0f - dfx_x(v), 1);
                                   BOOST_CHECK_CLOSE(dv.x.y,      - dfy_x(v), 1);
                                   BOOST_CHECK_CLOSE(dv.x.z,      - dfz_x(v), 1);
                            } else {
                                   BOOST_CHECK_EQUAL(dv.x.x, 1.0f);
                                   BOOST_CHECK_EQUAL(dv.x.y, 0);
                                   BOOST_CHECK_EQUAL(dv.x.z, 0);
                            }

                            if ( y > 0 && y < size.y - 1) {
                                   BOOST_CHECK_CLOSE(dv.y.x,      - dfx_y(v), 1);
                                   BOOST_CHECK_CLOSE(dv.y.y, 1.0f - dfy_y(v), 1);
                                   BOOST_CHECK_CLOSE(dv.y.z,      - dfz_y(v), 1);
                            } else {
                                   BOOST_CHECK_EQUAL(dv.y.x, 0);
                                   BOOST_CHECK_EQUAL(dv.y.y, 1.0f);
                                   BOOST_CHECK_EQUAL(dv.y.z, 0);
                            }

                            if ( z > 0 && z < size.z - 1) {
                                   BOOST_CHECK_CLOSE(dv.z.x,      - dfx_z(v), 1);
                                   BOOST_CHECK_CLOSE(dv.z.y,      - dfy_z(v), 1);
                                   BOOST_CHECK_CLOSE(dv.z.z, 1.0f - dfz_z(v), 1);
                            } else {
                                   BOOST_CHECK_EQUAL(dv.z.x, 0);
                                   BOOST_CHECK_EQUAL(dv.z.y, 0);
                                   BOOST_CHECK_EQUAL(dv.z.z, 1.0);
                            }
                     }
}


BOOST_FIXTURE_TEST_CASE(test_gridtransform_set_identity, GridTransformFixture)
{
       field.set_identity();
       C3DGridTransformation::const_iterator i = field.begin();

       for (size_t z = 0; z < size.z; ++z)
              for (size_t y = 0; y < size.y; ++y)
                     for (size_t x = 0; x < size.x; ++x, ++i) {
                            BOOST_CHECK_EQUAL( *i, C3DFVector(x, y, z));
                     }
}

BOOST_AUTO_TEST_CASE( test_grid3d_iterator )
{
       C3DBounds size(5, 10, 7);
       C3DInterpolatorFactory ipf("bspline:d=1", "mirror");
       C3DGridTransformation gt(size, ipf);
       C3DGridTransformation::field_iterator i = gt.field_begin();
       C3DFVector value(0.12, 0.32, 0.36);

       while (i != gt.field_end()) {
              *i = value;
              value.x += 0.12f;
              value.y += 0.32f;
              value.z += 0.36f;
              ++i;
       }

       C3DTransformation::const_iterator k = gt.begin();
       float n = 1.0f;

       for (size_t z = 0; z < size.z; ++z)
              for (size_t y = 0; y < size.y; ++y)
                     for (size_t x = 0; x < size.x; ++x, ++k, n += 1.0f) {
                            BOOST_CHECK_CLOSE( k->x, (float)x - n * 0.12f, 0.1f );
                            BOOST_CHECK_CLOSE( k->y, (float)y - n * 0.32f, 0.1f );
                            BOOST_CHECK_CLOSE( k->z, (float)z - n * 0.36f, 0.1f );
                     }
}

BOOST_AUTO_TEST_CASE( test_get_parameters )
{
       C3DFVector init_a[8] = {
              C3DFVector(1, 1, 1), C3DFVector(2, 1, 1),
              C3DFVector(1, 2, 1), C3DFVector(2, 2, 1),

              C3DFVector(1, 1, 2), C3DFVector(2, 1, 2),
              C3DFVector(1, 2, 2), C3DFVector(2, 2, 2)
       };
       C3DBounds size(2, 2, 2);
       C3DInterpolatorFactory ipf("bspline:d=1", "mirror");
       C3DGridTransformation a(size, ipf);
       std::copy(init_a, init_a + 8, a.field_begin());
       CDoubleVector params = a.get_parameters();
       BOOST_CHECK_EQUAL(params.size(), 24u);
       BOOST_REQUIRE(params.size() == 24u);
       double test_params[24] =  {1, 1, 1, 2, 1, 1, 1, 2, 1, 2, 2, 1,
                                  1, 1, 2, 2, 1, 2, 1, 2, 2, 2, 2, 2
                                 };

       for (int i = 0; i < 24; ++i)
              BOOST_CHECK_EQUAL(params[i], test_params[i]);
}

BOOST_AUTO_TEST_CASE( test_set_parameters )
{
       double init_params[24] =  {1, 1, 1, 2, 1, 1, 1, 2, 1, 2, 2, 1,
                                  1, 1, 2, 2, 1, 2, 1, 2, 2, 2, 2, 2
                                 };
       C3DFVector test_vectors[8] = {
              C3DFVector(1, 1, 1), C3DFVector(2, 1, 1),
              C3DFVector(1, 2, 1), C3DFVector(2, 2, 1),

              C3DFVector(1, 1, 2), C3DFVector(2, 1, 2),
              C3DFVector(1, 2, 2), C3DFVector(2, 2, 2)
       };
       C3DBounds size(2, 2, 2);
       C3DInterpolatorFactory ipf("bspline:d=1", "mirror");
       C3DGridTransformation a(size, ipf);
       CDoubleVector params(24);
       std::copy(init_params, init_params + 24, params.begin());
       a.set_parameters(params);
       auto ia = a.field_begin();

       for (int i = 0; i < 8; ++i, ++ia)
              BOOST_CHECK_EQUAL(*ia, test_vectors[i]);
}

BOOST_AUTO_TEST_CASE( test_gridtransform_add )
{
       C3DFVector init_a[18] = {
              C3DFVector(2, 3, 2), C3DFVector(0, 0, 0), C3DFVector(0, 0, 0),
              C3DFVector(0, 0, 0), C3DFVector(-1, -1, 0), C3DFVector(0, 0, 0),
              C3DFVector(0, 0, 0), C3DFVector(0, 0, 0), C3DFVector(0, 0, 0),

              C3DFVector(2, 3, 2), C3DFVector(0, 0, 0), C3DFVector(0, 0, 0),
              C3DFVector(0, 0, 0), C3DFVector(-1, -1, 0), C3DFVector(0, 0, 0),
              C3DFVector(0, 0, 0), C3DFVector(0, 0, 0), C3DFVector(0, 0, 0)
       };
       C3DFVector init_b[18] = {
              C3DFVector(0, 0, 0), C3DFVector(0, 0, 0), C3DFVector(0, 0, 0),
              C3DFVector(0, 0, 0), C3DFVector(1, 1, 0), C3DFVector(0.5, 0, 0),
              C3DFVector(0, 0, 0), C3DFVector(0, 0.5, 0), C3DFVector(4, 5, 0),

              C3DFVector(0, 0, 0), C3DFVector(0, 0, 0), C3DFVector(0, 0, 0),
              C3DFVector(0, 0, 0), C3DFVector(1, 1, 1), C3DFVector(0.5, 0, 0),
              C3DFVector(0, 0, 0), C3DFVector(0, 0.5, 0), C3DFVector(4, 5, 0)
       };
       C3DBounds size(3, 3, 2);
       C3DInterpolatorFactory ipf("bspline:d=1", "mirror");
       C3DGridTransformation a(size, ipf);
       std::copy(init_a, init_a + 18, a.field_begin());
       C3DGridTransformation b(size, ipf);
       std::copy(init_b, init_b + 18, b.field_begin());
       C3DGridTransformation c = a + b;
       BOOST_CHECK_EQUAL( c.get_displacement_at(C3DFVector(1, 1, 0)), C3DFVector(3, 4, 2));
       BOOST_CHECK_EQUAL( c.get_displacement_at(C3DFVector(1, 1, 1)), C3DFVector(3, 4, 3));
       BOOST_CHECK_EQUAL( c.get_displacement_at(C3DFVector(2, 1, 0)), C3DFVector(0.0, -0.5, 0));
       BOOST_CHECK_EQUAL( c.get_displacement_at(C3DFVector(1, 2, 0)), C3DFVector(-0.5, 0.0, 0));
}


BOOST_FIXTURE_TEST_CASE( test_gridtransform_max, GridTransformFixture )
{
       const float max = sqrt(49 + 49 + 49 * 2.25 * 2.25);
       BOOST_CHECK_CLOSE(field.get_max_transform(), max, 1);
}

#if 0
BOOST_FIXTURE_TEST_CASE( test_gridtransform_pertuberate, GridTransformFixture )
{
       C3DFVectorfield v(size);
       C3DFVector vv(10.0, 20.0);
       fill(v.begin(), v.end(), vv);
       float gamma = field.pertuberate(v);
       C3DFVector lmg(184, 107);
       C3DFVector mg(vv.x - vv.x * dfx_x(lmg.x, lmg.y) - vv.y * dfx_y(lmg.x, lmg.y),
                     vv.y - vv.x * dfy_x(lmg.x, lmg.y) - vv.y * dfy_y(lmg.x, lmg.y));
       BOOST_CHECK_CLOSE(gamma, mg.norm(), 0.1);

       for (size_t y = 1; y < size.y - 1; ++y)
              for (size_t x = 1; x < size.x - 1; ++x) {
                     const C3DFVector& iv = v(x, y);
                     BOOST_CHECK_CLOSE(iv.x, vv.x - vv.x * dfx_x(x, y) - vv.y * dfx_y(x, y), 1);
                     BOOST_CHECK_CLOSE(iv.y, vv.y - vv.x * dfy_x(x, y) - vv.y * dfy_y(x, y), 1);
              }
}

BOOST_FIXTURE_TEST_CASE( test_gridtransform_get_jacobian, GridTransformFixture )
{
       C3DFVectorfield v(size);
       fill(v.begin(), v.end(), C3DFVector(10.0, 20.0));
       C3DFVector lmg(64, 32);
       float j = field.get_jacobian(v, 1.0);
       C3DFMatrix J = field.derivative_at(lmg.x, lmg.y);
       BOOST_CHECK_CLOSE(j, J.x.x * J.y.y - J.x.y * J.y.x, 0.1);
}
#endif

BOOST_FIXTURE_TEST_CASE( test_grid_clone, GridTransformFixture )
{
       P3DTransformation clone(field.clone());
       C3DGridTransformation& spclone = dynamic_cast<C3DGridTransformation&>(*clone);
       BOOST_CHECK_EQUAL(spclone.get_size(), field.get_size());
       C3DGridTransformation::const_iterator ic = spclone.begin();
       C3DGridTransformation::const_iterator ec = spclone.end();
       C3DGridTransformation::const_iterator io = field.begin();

       while (ic != ec ) {
              BOOST_CHECK_EQUAL(*ic, *io);
              ++ic;
              ++io;
       }
}



float GridTransformFixture::fx(const C3DFVector& v)
{
       const C3DFVector s = scaled(v);
       const C3DFVector h( 1.0 - s.x, 2.0 - 2 * s.y, 3.0 - 3 * s.z);
       return 	0.5 * h.norm2();
}

float GridTransformFixture::fy(const C3DFVector& v)
{
       const C3DFVector s = scaled(v);
       const C3DFVector h( 1.0 - s.x, 2.0 - 2 * s.y, 3.0 - 3 * s.z);
       return 0.5 * 2.25 * h.norm2();
}

float GridTransformFixture::fz(const C3DFVector& v)
{
       const C3DFVector s = scaled(v);
       const C3DFVector h( 1.0 - s.x, 2.0 - 2 * s.y, 3.0 - 3 * s.z);
       return 	-0.5 * h.norm2();
}

float GridTransformFixture::dfx_x(const C3DFVector& v)
{
       return m_scale.x * (m_scale.x * v.x - 1);
}

float GridTransformFixture::dfx_y(const C3DFVector& v)
{
       return m_scale.y * (2.0 * m_scale.y * v.y - 2.0) * 2.0;
}

float GridTransformFixture::dfx_z(const C3DFVector& v)
{
       return m_scale.z *  (3.0 * m_scale.z * v.z - 3.0) * 3.0;
}


float GridTransformFixture::dfy_x(const C3DFVector& v)
{
       return 2.25 * m_scale.x *  (m_scale.x * v.x - 1);
}

float GridTransformFixture::dfy_y(const C3DFVector& v)
{
       return 2.25  * m_scale.y *  (2.0 * m_scale.y * v.y - 2.0) * 2.0;
}

float GridTransformFixture::dfy_z(const C3DFVector& v)
{
       return 2.25 * m_scale.z *  (3.0 * m_scale.z * v.z - 3.0) * 3.0;
}

float GridTransformFixture::dfz_x(const C3DFVector& v)
{
       return - m_scale.x *  (m_scale.x * v.x - 1);
}

float GridTransformFixture::dfz_y(const C3DFVector& v)
{
       return - m_scale.y *  (2.0 * m_scale.y * v.y - 2.0) * 2.0;
}

float GridTransformFixture::dfz_z(const C3DFVector& v)
{
       return - m_scale.z *  (3.0 * m_scale.z * v.z - 3.0) * 3.0;
}

float GridTransformFixture::dfx_xx(float x, float y, float z)
{
       x *= m_scale.x;
       y *= m_scale.y;
       z *= m_scale.z;
       return - m_scale.x * m_scale.x * sinf(x - M_PI / 2.0) * ( 1.0 + sinf(2 * y  - M_PI / 2.0));
}

float GridTransformFixture::dfx_xy(float x, float y, float z)

{
       x *= m_scale.x;
       y *= m_scale.y;
       z *= m_scale.z;
       return 2.0 * m_scale.x * m_scale.y * cosf(x - M_PI / 2.0) * cosf(2 * y  - M_PI / 2.0);
}

float GridTransformFixture::dfx_yx(float x, float y, float z)
{
       x *= m_scale.x;
       y *= m_scale.y;
       z *= m_scale.z;
       return 2.0 * m_scale.x * m_scale.y * cosf(x - M_PI / 2.0) * cosf(2 * y  - M_PI / 2.0);
}

float GridTransformFixture::dfx_yy(float x, float y, float z)
{
       x *= m_scale.x;
       y *= m_scale.y;
       z *= m_scale.z;
       return -4.0 * m_scale.y * m_scale.y * (1.0 + sinf(x - M_PI / 2.0)) * sinf(2 * y  - M_PI / 2.0);
}

float GridTransformFixture::dfy_xx(float x, float y, float z)
{
       x *= m_scale.x;
       y *= m_scale.y;
       z *= m_scale.z;
       return -4.0 * m_scale.x * m_scale.x * cosf(2 * x) * (1.0 - cosf(y));
}

float GridTransformFixture::dfy_xy(float x, float y, float z)
{
       x *= m_scale.x;
       y *= m_scale.y;
       z *= m_scale.z;
       return -2.0 * m_scale.x * m_scale.y * sinf(2 * x) * sinf(y);
}

float GridTransformFixture::dfy_yx(float x, float y, float z)
{
       x *= m_scale.x;
       y *= m_scale.y;
       z *= m_scale.z;
       return -2.0 * m_scale.y * m_scale.x * sinf(2 * x) * sinf(y);
}

float GridTransformFixture::dfy_yy(float x, float y, float z)
{
       x *= m_scale.x;
       y *= m_scale.y;
       z *= m_scale.z;
       return - m_scale.y * m_scale.y * (1.0 - cosf(2 * x)) * cosf(y);
}

