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

#include <mia/core/msgstream.hh>
#include <mia/2d/fuzzyclustersolver_cg.hh>


NS_MIA_BEGIN
// this needs to be tested
bool fborder (long index, long nx, long ny)
{
       const long x = index % nx;

       if ( x < 2 || x > nx - 3 )
              return true;

       index /= nx;
       const long y = index % ny;

       if ( y < 2 || y > ny - 3 )
              return true;

       return false;
}



C2DSolveCG::C2DSolveCG (C2DFImage& w1, C2DFImage& f1, C2DFImage& gain_image, double l1, double l2, double r_res, double m_res):
       m_lambda1(l1),
       m_lambda2(l2),
       m_iter(0),
       m_nx(w1.get_size().x),
       m_ny(w1.get_size().y),
       m_count(m_nx * m_ny),
       m_weight_imagePtr(&w1(0, 0)),
       m_fptr(&f1(0, 0)),
       m_gain_image_ptr(&gain_image(0, 0)),

       m_b(m_count),
       m_v(m_count),
       m_r(m_count),

       m_rho(m_count),   // p^(k)
       m_g(m_count),
       m_Ag(m_count),   // speichert A * p

       m_scale(m_count),
       m_scale2(m_count),
       m_border(m_count),

       m_r1rho1(0.0),
       m_r2rho2(0.0),
       m_normr0(0.0),
       m_q(0.0),
       m_e(0.0),
       m_sprod(0.0),

       m_min_res(m_res),
       m_relres(r_res)

{
       TRACE_FUNCTION;
       init ();
}


C2DSolveCG::~C2DSolveCG()
{
}



// Umweg, da keine Methoden parallel aufgerufen werden koennen
void solver(C2DSolveCG *s, long *maxit, double *normr, double *firstnormr0)
{
       s->solvepar(maxit, normr, firstnormr0);
}


int C2DSolveCG::solve(long max_iterations, double *firstnormr0)
{
       double normr;
#ifdef PARALLEL
       m_set_procs(NPROCS);
       m_fork((void (*)())solver, this, &max_iterations, &normr, firstnormr0);
       m_kill_procs();
#else
       solver(this, &max_iterations, &normr, firstnormr0);
#endif
       return normr > 1;
}


// loest das Gleichungssystem mittels CG-Algorithmus
void C2DSolveCG::solvepar(long *maxit, double *normr, double *firstnormr0)
{
       TRACE_FUNCTION;
       // Prozessornummer
       int me;
       // Start und Ende fuer Vektoroperationen
       long start;
       long ende;
       // Setze me, start und ende
       me = 0;
       start = 0;
       ende = m_count;
       // Initialisierung
       cvmsg() << "mult = " << start << ", " << ende << "\n";
       multA(m_v, m_r, start, ende);
       // Berechnung normr0;
       double tmp_normr0 = 0;

       for (long i = start; i < ende; i++) {
              m_r[i] += m_b[i];
              double square = m_r[i] / m_scale[i];
              tmp_normr0 += square * square;
       }

       m_normr0 = tmp_normr0;

       if (me == 0)
              *normr = m_normr0 = sqrt(m_normr0);

       if (*firstnormr0 == 1.0 && m_normr0 > 1.0)
              *firstnormr0 = m_normr0;

       // Iterationen
       while (*normr >= m_min_res && m_iter < *maxit && *normr / *firstnormr0 > m_relres) {
              if (me == 0) {
                     m_iter++;
                     copy(m_r.begin(), m_r.end(), m_rho.begin());
                     m_r2rho2 = m_r1rho1;
              }

              // r1rho1 = r * rho
              if (me == 0) m_r1rho1 = 0;

              double tmp_r1rho1 = 0;

              for (long i = start; i < ende; i++) {
                     tmp_r1rho1 += m_r[i] * m_rho[i];
              }

              m_r1rho1 = tmp_r1rho1;

              if (m_iter >= 2 && me == 0) m_e = m_r1rho1 / m_r2rho2;
              else m_e = 0;

              for (long i = start; i < ende; i++) m_g[i] = -m_rho[i] + m_e * m_g[i];

              multA(m_g, m_Ag, start, ende);

              // q = r1rho1 / (g * Ag)
              if (me == 0) m_sprod = 0;

              double tmp_sprod = 0;

              for (long i = start; i < ende; i++) tmp_sprod += m_g[i] * m_Ag[i];

              m_sprod = tmp_sprod;

              if (me == 0) m_q = m_r1rho1 / m_sprod;

              // v = v1 + q * g, r = r1 + q * Ag
              if (me == 0) *normr = 0;

              double tmp_normr = 0;

              for (long i = start; i < ende; i++) {
                     m_v[i] += m_q * m_g[i];
                     m_r[i] += m_q * m_Ag[i];
                     tmp_normr += pow(m_r[i] / m_scale[i], 2);
              }

              *normr = tmp_normr;

              if (me == 0)
                     *normr = sqrt(*normr);

              //*\todo implement some progression alert here */
       }

       return;
}


// Vektor-Matrix-Multiplikation
#define VALUE(o) (sc2[o] * x2[o])

void C2DSolveCG::multA(vector<double>& x, vector<double>& result, long start, long ende)
{
       assert(x.size() == result.size());
       assert(x.size() == m_border.size());
       assert(x.size() == m_scale2.size());
       assert(x.size() == m_count);
       auto bord2 = m_border.begin();
       auto res2 = result.begin();
       auto x2 = x.begin();
       auto sc2 = m_scale2.begin();
       vector<double> help(x.size());
       transform(x.begin(), x.end(), m_scale2.begin(), help.begin(),
       [](double x, double s) {
              return x * s;
       });
       auto value = help.begin();

       for (int i = start; i < ende; i++, ++bord2, ++res2, ++x2, ++value, ++sc2) {
              if (*bord2) {
                     *res2 = *x2;
              } else {
                     double s1 = value[-m_nx      ] + value[-1] + value[+1] + value[+m_nx];
                     double s2 = value[-m_nx   - 1] + value[-m_nx + 1]	+ value[+ m_nx - 1]  + value[+m_nx + 1];
                     double s3 = value[ - 2 * m_nx] + value[ - 2     ]  + value[+ 2]        + value[ + 2 * m_nx];
                     *res2 += ((s3 + 2 * s2 - 8 * s1) * m_lambda2 - m_lambda1 * s1) * *sc2;
              }
       }
}


#define VALUE2(o) ((bord2[o])?0:x2[o])

void C2DSolveCG::multA_float(float *x, float *result)
{
       auto bord2 = m_border.begin();
       float *x2, *res2;

       for (unsigned long i = 0; i < m_count; i++, ++bord2) {
              double s1, s2, s3;
              res2  = result + i;
              x2    = x + i;

              if (*bord2) {
                     *res2 = (m_weight_imagePtr[i] + 4 * m_lambda1 + 20 * m_lambda2) * *x2;
                     continue;
              }

              s1 = VALUE2(-m_nx)    + VALUE2(-1) + VALUE2(+1) + VALUE2(+m_nx);
              s2 = VALUE2(-m_nx - 1)     + VALUE2(-m_nx + 1)	 + VALUE2(+m_nx - 1)     + VALUE2(+m_nx + 1);
              s3 = VALUE2( - 2 * m_nx) + VALUE2( - 2)	+ VALUE2( + 2) + VALUE2( + 2 * m_nx);
              *res2 = (m_weight_imagePtr[i] + 4 * m_lambda1 + 8 * m_lambda2) * *x2  + ((s3 + 2 * s2 - 12 * s1) * m_lambda2 - m_lambda1 * s1);
       }
}

// loest die PDE  (w + m_lambda1 * H1 + m_lambda2 * H2)*m = f
void C2DSolveCG::init()
{
       TRACE_FUNCTION;

       // set b and scaling
       for (unsigned long i = 0; i < m_count; i++) {
              m_border[i] = (fborder(i, m_nx, m_ny) ? 1 : 0);
       }

       for (unsigned long i = 0; i < m_count; i++) {
              // Wenn Element Randelement ist, setze b entsprechend
              if (m_border[i]) {
                     m_b[i] = -m_gain_image_ptr[i];
                     m_scale[i] = 1.0;
                     m_scale2[i] = 0;
                     m_v[i] = m_gain_image_ptr[i] / m_scale[i];
                     continue;
              }

              // set b[i]
              m_b[i] = -m_fptr[i];
              // ziehe Anteile aus Randelementen ab
              // scaling
              m_scale[i] = m_scale2[i] = 1.0 / sqrt(m_weight_imagePtr[i] + 4 * m_lambda1 + 20 * m_lambda2);
              m_b[i] *= m_scale[i];
              m_v[i] = m_gain_image_ptr[i] / m_scale[i];
       }
}


void C2DSolveCG::get_solution(C2DFImage& m)
{
       C2DFImage::iterator m_gain_image_ptr = m.begin();

       for (unsigned long i = 0; i < m_count; i++, ++m_gain_image_ptr)
              *m_gain_image_ptr = (float)m_v[i] * m_scale[i];
}


void C2DSolveCG::add_to_solution(C2DFImage *e)
{
       C2DFImage::iterator eptr = e->begin();

       for (unsigned long i = 0; i < m_count; i++)
              m_v[i] += eptr[i] / m_scale[i];
}

NS_MIA_END
