/* -*- 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 "TRANSFORMS"

#include <ostream>
#include <fstream>
#include <mia/3d/transform.hh>
#include <mia/internal/autotest.hh>
#include <mia/3d/transformfactory.hh>

NS_MIA_USE;
namespace bfs = boost::filesystem;

class Cost3DMock
{
public:
       Cost3DMock(const C3DBounds& size);
       double value(const C3DTransformation& t) const;
       double value_and_gradient(C3DFVectorfield& gradient) const;

       double src_value(const C3DFVector& x)const;
       double ref_value(const C3DFVector& x)const;
       C3DFVector src_grad(const C3DFVector& x)const;
       C3DBounds m_size;
       C3DFVector m_center;
       float m_r;
};

class TransformGradientFixture
{
public:
       TransformGradientFixture();

       void run_test(C3DTransformation& t, double tol = 0.1)const;

       C3DBounds size;
       Cost3DMock cost;

       C3DFVector m_x;
       C3DFVectorfield gradient;
};



BOOST_AUTO_TEST_CASE (selftest_Cost3DMock)
{
       C3DBounds size(20, 30, 15);
       Cost3DMock cm(size);
       C3DFVector x(11.0, 16.0, 7.0);
       C3DFVector dx(.001, 0.0, 0);
       C3DFVector dy(.0, 0.001, 0);
       C3DFVector dz(.0, .0, 0.001);
       C3DFVector grad = cm.src_grad(x);
       C3DFVector fdgrad( (cm.src_value(x + dx) - cm.src_value(x - dx)) / 0.002,
                          (cm.src_value(x + dy) - cm.src_value(x - dy)) / 0.002,
                          (cm.src_value(x + dz) - cm.src_value(x - dz)) / 0.002
                        );
       BOOST_CHECK_CLOSE(grad.x, fdgrad.x, 0.1);
       BOOST_CHECK_CLOSE(grad.y, fdgrad.y, 0.1);
       BOOST_CHECK_CLOSE(grad.z, fdgrad.z, 0.1);
}

BOOST_FIXTURE_TEST_CASE (test_translate_Gradient, TransformGradientFixture)
{
       const C3DTransformCreatorHandler::Instance& handler =
              C3DTransformCreatorHandler::instance();
       P3DTransformationFactory creater = handler.produce("translate");
       P3DTransformation transform = creater->create(size);
       run_test(*transform);
}

BOOST_FIXTURE_TEST_CASE (test_affine_Gradient, TransformGradientFixture)
{
       const C3DTransformCreatorHandler::Instance& handler =
              C3DTransformCreatorHandler::instance();
       P3DTransformationFactory creater = handler.produce("affine");
       P3DTransformation transform = creater->create(size);
       run_test(*transform, 0.2);
}

BOOST_FIXTURE_TEST_CASE (test_rigid_Gradient, TransformGradientFixture)
{
       const C3DTransformCreatorHandler::Instance& handler =
              C3DTransformCreatorHandler::instance();
       P3DTransformationFactory creater = handler.produce("rigid");
       P3DTransformation transform = creater->create(size);
       run_test(*transform);
}

BOOST_FIXTURE_TEST_CASE (test_rigid_Gradient_center_shifted, TransformGradientFixture)
{
       const C3DTransformCreatorHandler::Instance& handler =
              C3DTransformCreatorHandler::instance();
       P3DTransformationFactory creater = handler.produce("rigid:origin=[<0.5,0.4,0.6>]");
       P3DTransformation transform = creater->create(size);
       run_test(*transform);
}


BOOST_FIXTURE_TEST_CASE (test_rotation_Gradient, TransformGradientFixture)
{
       const C3DTransformCreatorHandler::Instance& handler =
              C3DTransformCreatorHandler::instance();
       P3DTransformationFactory creater = handler.produce("rotation");
       P3DTransformation transform = creater->create(size);
       run_test(*transform);
}

BOOST_FIXTURE_TEST_CASE (test_rotation_Gradient_center_shifted, TransformGradientFixture)
{
       const C3DTransformCreatorHandler::Instance& handler =
              C3DTransformCreatorHandler::instance();
       P3DTransformationFactory creater = handler.produce("rotation:origin=[<0.5,0.4,0.6>]");
       P3DTransformation transform = creater->create(size);
       run_test(*transform);
}

TransformGradientFixture::TransformGradientFixture():
       size(20, 28, 16),
       cost(size),
       m_x(11, 16, 7),
       gradient(size)

{
       cost.value_and_gradient(gradient);
}

void TransformGradientFixture::run_test(C3DTransformation& t, double tol)const
{
       auto params = t.get_parameters();
       CDoubleVector trgrad(params.size());
       t.translate(gradient,  trgrad);
       double delta = 0.01;
       int n_close_zero = 0;
       int n_zero = 0;

       for (auto itrg =  trgrad.begin(),
            iparam = params.begin(); itrg != trgrad.end(); ++itrg, ++iparam) {
              *iparam += delta;
              t.set_parameters(params);
              double cost_plus = cost.value(t);
              *iparam -= 2 * delta;
              t.set_parameters(params);
              double cost_minus = cost.value(t);
              *iparam += delta;
              cvdebug() << cost_plus << ", " << cost_minus << "\n";
              double test_val = (cost_plus - cost_minus) / (2 * delta);
              cvdebug() << *itrg << " vs " << test_val << "\n";

              if (fabs(*itrg) < 1e-6 && fabs(test_val) < 1e-6) {
                     n_close_zero++;
                     continue;
              }

              if (*itrg == 0.0 && fabs(test_val) < 1e-5) {
                     n_zero++;
                     continue;
              }

              BOOST_CHECK_CLOSE(*itrg, test_val, tol);
       }

       cvmsg() << "N=" << trgrad.size() << "\n";
       cvmsg() << "value pairs < 1e-6 = " << n_close_zero << "\n";
       cvmsg() << "grad value zero, but finite difference below 1e-5 = " << n_zero << "\n";
}

Cost3DMock::Cost3DMock(const C3DBounds& size):
       m_size(size),
       m_center(0.5 * size.x, 0.5 * size.y, 0.5 * size.z ),
       m_r(m_center.norm())
{
}

double Cost3DMock::value(const C3DTransformation& t) const
{
       assert(m_size == t.get_size());
       double result = 0.0;
       auto it = t.begin();

       for (size_t z = 0; z < m_size.z; ++z)
              for (size_t y = 0; y < m_size.y; ++y)
                     for (size_t x = 0; x < m_size.x; ++x, ++it) {
                            double v = src_value(*it) - ref_value(C3DFVector(x, y, z));
                            result += v * v;
                     }

       return result;
}

double Cost3DMock::value_and_gradient(C3DFVectorfield& gradient) const
{
       assert(gradient.get_size() == m_size);
       double result = 0.0;
       auto ig = gradient.begin();

       for (size_t z = 0; z < m_size.z; ++z)
              for (size_t y = 0; y < m_size.y; ++y)
                     for (size_t x = 0; x < m_size.x; ++x, ++ig) {
                            C3DFVector pos(x, y, z);
                            double v = src_value(pos) - ref_value(pos);
                            result += v * v;
                            *ig = 2.0 * v * src_grad(pos);
                     }

       return result;
}

double Cost3DMock::src_value(const C3DFVector& x)const
{
       const C3DFVector p = x - m_center;
       return exp( - p.norm2() / m_r  );
}

C3DFVector Cost3DMock::src_grad(const C3DFVector& x)const
{
       return - 2.0f  * (x - m_center) / m_r  * src_value(x);
}

double Cost3DMock::ref_value(const C3DFVector& x)const
{
       const C3DFVector p = x - m_center - C3DFVector(1.1, 1.2, 1.1);
       return exp( - p.norm2() / m_r );
}


