/* -*- 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/3d/fuzzyclustersolver_cg.hh>

NS_MIA_BEGIN
// this needs to be tested
bool fborder (long index, long nx, long ny, long nz)
{
       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;

       index /= nx;
       const long z = index / ny;

       if ( z < 2 || z > nz - 3 )
              return true;

       return false;
}



solve_sCG::solve_sCG (C3DFImage& w1, C3DFImage& f1, C3DFImage& gain_image, double l1, double l2, double r_res, double m_res):
       __lambda1(l1),
       __lambda2(l2),
       __iter(0),
       __nx(w1.get_size().x),
       __ny(w1.get_size().y),
       __nz(w1.get_size().z),
       __count(__nz * __ny * __nx),
       __weight_imagePtr(&w1(0, 0, 0)),
       __fptr(&f1(0, 0, 0)),
       __gain_image_ptr(&gain_image(0, 0, 0)),

       __b(new double[__count]),
       __v(new double[__count]),
       __r(new double[__count]),

       __rho(new double[__count]),   // p^(k)
       __g(new double[__count]),
       __Ag(new double[__count]),   // speichert A * p

       __scale(new double[__count]),
       __scale2(new double[__count]),
       __border(new bool[__count]),

       __r1rho1(0.0),
       __r2rho2(0.0),
       __normr0(0.0),
       __q(0.0),
       __e(0.0),
       __sprod(0.0),

       __min_res(m_res),
       __relres(r_res)

{
       init ();
}


solve_sCG::~solve_sCG()
{
       delete[] __b;
       delete[] __v;
       delete[] __scale;
       delete[] __scale2;
       delete[] __border;
       // free pointers
       delete[] __r;
       delete[] __rho;
       delete[] __g;
       delete[] __Ag;
}



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


int solve_sCG::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 solve_sCG::solvepar(long *maxit, double *normr, double *firstnormr0)
{
       // Prozessornummer
       int me;
       // Start und Ende fuer Vektoroperationen
       long start;
       long ende;
       // Setze me, start und ende
       me = 0;
       start = 0;
       ende = __count;
       // Initialisierung
       multA(__v, __r, start, ende);
       // Berechnung normr0;
       double tmp_normr0 = 0;

       for (long i = start; i < ende; i++) {
              __r[i] += __b[i];
              double square = __r[i] / __scale[i];
              tmp_normr0 += square * square;
       }

       __normr0 = tmp_normr0;

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

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

       // Iterationen
       while (*normr >= __min_res && __iter < *maxit && *normr / *firstnormr0 > __relres) {
              if (me == 0) {
                     __iter++;
                     memcpy(__rho, __r, __count * sizeof(double));
                     __r2rho2 = __r1rho1;
              }

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

              double tmp_r1rho1 = 0;

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

              __r1rho1 = tmp_r1rho1;

              if (__iter >= 2 && me == 0) __e = __r1rho1 / __r2rho2;
              else __e = 0;

              for (long i = start; i < ende; i++) __g[i] = -__rho[i] + __e * __g[i];

              multA(__g, __Ag, start, ende);

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

              double tmp_sprod = 0;

              for (long i = start; i < ende; i++) tmp_sprod += __g[i] * __Ag[i];

              __sprod = tmp_sprod;

              if (me == 0) __q = __r1rho1 / __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++) {
                     __v[i] += __q * __g[i];
                     __r[i] += __q * __Ag[i];
                     tmp_normr += pow(__r[i] / __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 solve_sCG::multA(double *x, double *result, long start, long ende)
{
       long i;
       bool *bord2;
       double *x2, *sc2, *res2;

       for (i = start; i < ende; i++) {
              double s1, s2, s3;
              bord2 = &__border[i];
              res2  = &result[i];
              x2    = &x[i];

              if (*bord2) {
                     *res2 = *x2;
                     continue;
              }

              sc2 = &__scale2[i];
              s1 = VALUE(-__nx)    + VALUE(-1) + VALUE(+1) + VALUE(+__nx) +
                   VALUE(-__nx * __ny) + VALUE(+__nx * __ny);
              s2 = VALUE(-__nx - 1)     + VALUE(-__nx + 1)	+ VALUE(+__nx - 1)     + VALUE(+__nx + 1)    +
                   VALUE(-__nx * __ny - 1)  + VALUE(-__nx * __ny + 1)  + VALUE(-__nx * __ny - __nx) + VALUE(-__nx * __ny + __nx) +
                   VALUE(+__nx * __ny - 1)  + VALUE(+__nx * __ny + 1)  + VALUE(+__nx * __ny - __nx) + VALUE(+__nx * __ny + __nx);
              s3 = VALUE( - 2 * __nx) + VALUE( - 2)       + VALUE( + 2) +
                   VALUE( + 2 * __nx) + VALUE( - 2 * __nx * __ny) + VALUE( + 2 * __nx * __ny);
              *res2 = (*x2 + ((s3 + 2 * s2 - 12 * s1) * __lambda2 - __lambda1 * s1) * *sc2);
       }
}


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

void solve_sCG::multA_float(float *x, float *result)
{
       bool *bord2;
       float *x2, *res2;

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

              if (*bord2) {
                     *res2 = (__weight_imagePtr[i] + 6 * __lambda1 + 42 * __lambda2) * *x2;
                     continue;
              }

              s1 = VALUE2(-__nx)    + VALUE2(-1) + VALUE2(+1) + VALUE2(+__nx) +
                   VALUE2(-__nx * __ny) + VALUE2(+__nx * __ny);
              s2 = VALUE2(-__nx - 1)     + VALUE2(-__nx + 1)	 + VALUE2(+__nx - 1)     + VALUE2(+__nx + 1)	+
                   VALUE2(-__nx * __ny - 1)  + VALUE2(-__nx * __ny + 1)  + VALUE2(-__nx * __ny - __nx) + VALUE2(-__nx * __ny + __nx) +
                   VALUE2(+__nx * __ny - 1)  + VALUE2(+__nx * __ny + 1)  + VALUE2(+__nx * __ny - __nx) + VALUE2(+__nx * __ny + __nx);
              s3 = VALUE2( - 2 * __nx) + VALUE2( - 2)	+ VALUE2( + 2) +
                   VALUE2( + 2 * __nx) + VALUE2( - 2 * __nx * __ny) + VALUE2( + 2 * __nx * __ny);
              *res2 = (__weight_imagePtr[i] + 6 * __lambda1 + 42 * __lambda2) * *x2
                      + ((s3 + 2 * s2 - 12 * s1) * __lambda2 - __lambda1 * s1);
       }
}


// loest die PDE  (w + __lambda1 * H1 + __lambda2 * H2)*m = f
void solve_sCG::init()
{
       if (!__b || !__v || !__scale || !__scale2 || !__border)
              throw runtime_error("not enough memory");

       // set b and scaling
       for (unsigned long i = 0; i < __count; i++) {
              __border[i] = (fborder(i, __nx, __ny, __nz) ? 1 : 0);
       }

       for (unsigned long i = 0; i < __count; i++) {
              // Wenn Element Randelement ist, setze b entsprechend
              if (__border[i]) {
                     __b[i] = -__gain_image_ptr[i];
                     __scale[i] = 1.0;
                     __scale2[i] = 0;
                     __v[i] = __gain_image_ptr[i] / __scale[i];
                     continue;
              }

              // set b[i]
              __b[i] = -__fptr[i];
              // ziehe Anteile aus Randelementen ab
              // scaling
              __scale[i] = __scale2[i] = 1.0 / sqrt(__weight_imagePtr[i] + 6 * __lambda1 + 42 * __lambda2);
              __b[i] *= __scale[i];
              __v[i] = __gain_image_ptr[i] / __scale[i];
       }

       return;
}


void solve_sCG::get_solution(C3DFImage& m)
{
       C3DFImage::iterator __gain_image_ptr = m.begin();

       for (unsigned long i = 0; i < __count; i++, ++__gain_image_ptr)
              *__gain_image_ptr = (float)__v[i] * __scale[i];

       return;
}


void solve_sCG::add_to_solution(C3DFImage *e)
{
       C3DFImage::iterator eptr = e->begin();

       for (unsigned long i = 0; i < __count; i++)
              __v[i] += eptr[i] / __scale[i];

       return;
}

NS_MIA_END
