/* -*- 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 <numeric>
#include <mia/2d/groundtruthproblem.hh>

NS_MIA_BEGIN

using namespace std;

GroundTruthProblem::GroundTruthProblem(double a, double b,
                                       const C2DBounds& slice_size,
                                       size_t nframes,
                                       const CDoubleVector& left_side,
                                       const  CCorrelationEvaluator::result_type& corr):
       m_a(a),
       m_b(b),
       m_slice_size(slice_size),
       m_nframes(nframes),
       m_frame_size(slice_size.x * slice_size.y),
       m_left_side(left_side),
       m_corr(corr)
{
       add(property_gradient);
}

size_t GroundTruthProblem::do_size() const
{
       return m_left_side.size();
}

// this function can certainly be optimized
double GroundTruthProblem::evaluate_slice_gradient(CDoubleVector::const_iterator ii,
              vector<double>::iterator iout)
{
       double value = 0.0;
       {
              auto icorH = m_corr.horizontal.begin();
              auto icorV = m_corr.vertical.begin();

              // evaluate first row
              for (size_t x = 0; x < m_slice_size.x; ++x) {
                     const double delta = ii[x] - ii[x + m_slice_size.x] ;
                     iout[x] = 2 * delta * icorV[x];
                     value += delta * delta * icorV[x];
              }

              for (size_t x = 0; x < m_slice_size.x - 1; ++x) {
                     const double delta =  ii[x] - ii[x + 1];
                     iout[x] += 2 * delta * icorH[x];
                     value += delta * delta * icorH[x];
              }

              for (size_t x = 1; x < m_slice_size.x; ++x) {
                     const double delta = ii[x] - ii[x - 1];
                     iout[x] += 2 * delta * icorH[x - 1];
                     value += delta * delta * icorH[x - 1];
              }
       }

       for (size_t y = 1;  y < m_slice_size.y - 1; ++y)  {
              auto iii = ii + y * m_slice_size.x;
              auto iiout = iout + y * m_slice_size.x;
              auto icorH = m_corr.horizontal.begin_at(0, y);
              auto icorVm = m_corr.vertical.begin_at(0, y - 1);
              auto icorVp = m_corr.vertical.begin_at(0, y);

              for (size_t x = 0; x < m_slice_size.x; ++x) {
                     double delta = iii[x] - iii[x + m_slice_size.x];
                     iiout[x] = 2 * delta * icorVp[x];
                     value += delta * delta * icorVp[x];
              }

              for (size_t x = 0; x < m_slice_size.x; ++x) {
                     double delta = iii[x] - iii[x - m_slice_size.x];
                     iiout[x] += 2 * delta * icorVm[x];
                     value += delta * delta * icorVm[x];
              }

              for (size_t x = 0; x < m_slice_size.x - 1; ++x) {
                     double delta = iii[x] - iii[x + 1];
                     iiout[x] += 2 * delta * icorH[x];
                     value += delta * delta * icorH[x];
              }

              for (size_t x = 1; x < m_slice_size.x; ++x) {
                     double delta = iii[x] - iii[x - 1];
                     iiout[x] += 2 * delta * icorH[x - 1];
                     value +=  delta * delta * icorH[x - 1];
              }
       }

       {
              auto iii = ii + (m_slice_size.y - 1) * m_slice_size.x;
              auto iiout = iout + (m_slice_size.y - 1) * m_slice_size.x;
              auto icorH = m_corr.horizontal.begin_at(0, m_slice_size.y - 1);
              auto icorVm = m_corr.vertical.begin_at(0, m_slice_size.y - 2);
              // evaluate last row

              for (size_t x = 0; x < m_slice_size.x; ++x) {
                     double delta = iii[x] - iii[x - m_slice_size.x];
                     iiout[x] = 2 * delta * icorVm[x];
                     value +=  delta * delta * icorVm[x];
              }

              for (size_t x = 0; x < m_slice_size.x - 1; ++x) {
                     double delta = iii[x] - iii[x + 1];
                     iiout[x] += 2 * delta * icorH[x];
                     value +=  delta * delta * icorH[x];
              }

              for (size_t x = 1; x < m_slice_size.x; ++x) {
                     double delta = iii[x] - iii[x - 1];
                     iiout[x] += 2 * delta * icorH[x - 1];
                     value +=  delta * delta * icorH[x - 1];
              }
       }

       return value;
}

double GroundTruthProblem::evaluate_spacial_gradients(const CDoubleVector& x)
{
       double result = 0.0;
       m_spacial_gradient.resize(x.size());

       for (size_t f = 0; f < m_nframes; ++f) {
              result += evaluate_slice_gradient(x.begin() + f * m_frame_size,
                                                m_spacial_gradient.begin() + f * m_frame_size);
       }

       return result;
}

double  GroundTruthProblem::evaluate_time_gradients(const CDoubleVector& x)
{
       /*
         The finite difference test shows that in the first and in the last
         slice the gradient needs the factor 5, while in the other slices a
         factor 6 is used (the latter follows from a proper derivation)
        */
       double value = 0.0;
       m_time_derivative.resize(x.size());
       auto im2 = x.begin();
       auto im = x.begin();
       auto ii = x.begin();
       auto ip = x.begin() + m_frame_size;
       auto ip2 = x.begin() + 2 * m_frame_size;
       auto ih = m_time_derivative.begin();

       for (size_t k = 0; k < m_frame_size; ++k, ++ih, ++ip, ++ii, ++ip2) {
              const double v =  2 * *ii - 2 * *ip ;
              value += v * v;
              *ih = 5 * *ii +  *ip2 - 6 * *ip;
       }

       for (size_t k = 0; k < m_frame_size; ++k, ++ih, ++ip, ++ii, ++ip2, ++im) {
              const double imp = *ip + *im;
              const double v =  2 * *ii - imp;
              value += v * v;
              *ih = 9 *  *ii +  *ip2 - 6 * *im - 4 * *ip;
       }

       for (size_t f = 2; f < m_nframes - 2; ++f) {
              for (size_t k = 0; k < m_frame_size; ++k, ++ih, ++ip, ++im, ++ii, ++ip2, ++im2) {
                     const double imp = *ip + *im;
                     const double v =  2 * *ii - imp;
                     value += v * v;
                     *ih = 6 * *ii +  (*ip2 + *im2)  - 4 * imp;
              }
       }

       for (size_t k = 0; k < m_frame_size; ++k, ++ih, ++im, ++ii, ++ip, ++im2) {
              const double imp = *ip + *im;
              double v =  2 * *ii - imp;
              value += v * v;
              *ih = 9 *  *ii +  *im2 - 6 * *ip - 4 * *im;
       }

       for (size_t k = 0; k < m_frame_size; ++k, ++ih, ++im, ++ii, ++im2) {
              double v =  2 * *ii -  2 * *im;
              value += v * v;
              *ih = 5 * *ii +  *im2 - 6 * *im;
       }

       return value;
}


double  GroundTruthProblem::do_f(const CDoubleVector& x)
{
       double spacial_cost = evaluate_spacial_gradients(x);
       double temporal_cost = evaluate_time_gradients(x);
       double result = 0.0;

       for (auto ix = x.begin(), il = m_left_side.begin();
            ix != x.end(); ++ix, ++il) {
              double v = *ix - *il;
              result += v * v;
       }

       // inner product them with alpha
       result += m_a * spacial_cost;
       result += m_b * temporal_cost;
       result *=  0.5;
       cvmsg() << "GroundTruthProblem::f = " << result << "\n";
       return result;
}

void GroundTruthProblem::do_df(const CDoubleVector&  x, CDoubleVector&  g)
{
       evaluate_spacial_gradients(x);
       evaluate_time_gradients(x);
       auto ix = x.begin();
       auto il = m_left_side.begin();
       auto ddt = m_time_derivative.begin();
       auto dds = m_spacial_gradient.begin();

       for (auto ig = g.begin(); ig != g.end(); ++ig, ++ix, ++ddt, ++dds, ++il) {
              *ig =  (*ix - *il) + m_a * *dds + m_b * *ddt;
       }
}

double  GroundTruthProblem::do_fdf(const CDoubleVector&  x, CDoubleVector&  g)
{
       const double spacial_cost = evaluate_spacial_gradients(x);
       const double temporal_cost = evaluate_time_gradients(x);
       double result = 0.0;

       for (auto ix = x.begin(), il = m_left_side.begin();
            ix != x.end(); ++ix, ++il) {
              double v = *ix - *il;
              result += v * v;
       }

       result += m_a * spacial_cost;
       result += m_b * temporal_cost;
       auto ix = x.begin();
       auto il = m_left_side.begin();
       auto ddt = m_time_derivative.begin();
       auto dds = m_spacial_gradient.begin();

       for (auto ig = g.begin(); ig != g.end(); ++ig, ++ix, ++ddt, ++dds, ++il) {
              *ig =  *ix - *il + m_a * *dds + m_b * *ddt;
       }

       result *=  0.5;
       cvmsg() << "GroundTruthProblem::f = " << result << "\n";
       return result;
}

void GroundTruthProblem::set_alpha_beta(double a, double b)
{
       m_a = a;
       m_b = b;
}

const std::vector<double>& GroundTruthProblem::get_spacial_gradient() const
{
       return m_spacial_gradient;
}

const std::vector<double>& GroundTruthProblem::get_time_derivative() const
{
       return m_time_derivative;
}

NS_MIA_END
