/*
 * ALBERTA:  an Adaptive multi Level finite element toolbox  using
 *           Bisectioning refinement and Error control by Residual
 *           Techniques
 *
 * File: SSOR_precon.c
 *
 * (SSOR) preconditioner for DOF[_DOWB]_MATRICES.
 *
 * NOTE: this duplicates the code found in ssor.c, however: we
 * pre-invert the diagonal for the sake of efficiency.
 *
 ******************************************************************************
 *
 *  author:     Claus-Justus Heine
 *              Abteilung fuer Angewandte Mathematik
 *              Albert-Ludwigs-Universitaet Freiburg
 *              Hermann-Herder-Str. 10
 *              79104 Freiburg
 *              Germany
 *              claus@mathematik.uni-freiburg.de
 *
 *  http://www.mathematik.uni-freiburg.de/IAM/ALBERTA
 *
 *  (c) by C.-J. Heine (2006-2009)
 *
 */

#if HAVE_CONFIG_H
# include <config.h>
#endif

#include "alberta.h"

struct SSOR_precon_s_data
{
  PRECON              precon;
  REAL                omega;
  int                 sym;    /* 0: sor, 1: ssor */
  int                 n_iter;

  const DOF_MATRIX    *matrix;
  const DOF_SCHAR_VEC *bound;
  bool                auto_bound;

  DOF                 dim;
  DOF                 max_dim;
  REAL                *inv_diag;
  REAL                *f;

  struct SSOR_precon_s_data  *next;
};

static bool init_SSOR_precon_s(void *precon_data)
{
  FUNCNAME("init_SSOR_precon_s");
  struct SSOR_precon_s_data *data  = (struct SSOR_precon_s_data *)precon_data;
  MATRIX_ROW_REAL         **row  = (MATRIX_ROW_REAL **)data->matrix->matrix_row;
  const S_CHAR            *b     = data->bound ? data->bound->vec : NULL;
  const DOF_ADMIN         *admin = data->matrix->row_fe_space->admin;
  int                     dim    = admin->size_used;
  REAL                    *inv_diag;

  if (data->max_dim < dim) {
    data->inv_diag =
      MEM_REALLOC(data->inv_diag, data->max_dim, dim, REAL);
    data->f       =
      MEM_REALLOC(data->f, data->max_dim, dim, REAL);
    data->max_dim = dim;
  }

  data->dim = dim;
  inv_diag = (REAL *)data->inv_diag;

  if (b) {
    FOR_ALL_DOFS(admin, {
	inv_diag[dof] =
	  (row[dof] && b[dof] <= 0 && ABS(row[dof]->entry[0]) > 1.0E-20)
	  ? 1.0/row[dof]->entry[0] : 1.0;
      });
  } else {
    FOR_ALL_DOFS(admin, {
	inv_diag[dof] = (row[dof] && ABS(row[dof]->entry[0]) > 1.0E-20)
	  ? 1.0/row[dof]->entry[0] : 1.0;
      });
  }

  FOR_ALL_FREE_DOFS(admin,
		    if (dof >= admin->size_used) {
		      break;
		    }
		    inv_diag[dof] = 1.0);

  return true;
}

/* Just a copy of ssor_s() with precomputed inverted diagonal ... */
static void SSOR_precon_s(void *ud, int dim, REAL *r)
{
  struct SSOR_precon_s_data *data = (struct SSOR_precon_s_data *)ud;
  const S_CHAR    *b        = data->bound ? data->bound->vec : NULL;  
  REAL            *inv_diag = data->inv_diag;
  REAL            *f        = data->f;
  REAL            omega     = data->omega;
  REAL            omega1    = 1.0 - omega;
  REAL            accu;
  int             iter;
  DOF             dof;
  
  /* set initial guess to zero */
  for (dof = 0; dof < data->dim; dof ++) {
    if (data->matrix->matrix_row[dof] == NULL) {
      f[dof] = 0.0;
      continue;
    }
    if (!b || b[dof] < DIRICHLET) {
      f[dof] = r[dof];
      r[dof] = 0.0;
    }
  }

  for (iter = 0; iter < data->n_iter; iter++) {

    /* forward Gauss-Seidel iteration */
    for (dof = 0; dof < data->dim; dof ++) {
      if (data->matrix->matrix_row[dof] == NULL) {
	continue;
      }

      if (!b || b[dof] < DIRICHLET) { 
	accu = f[dof]; /* remember rhs */

	FOR_ALL_MAT_COLS(REAL, data->matrix->matrix_row[dof], {
	    if (col_dof != dof) {
	      accu -= row->entry[col_idx] * r[col_dof];
	    }
	  });
	r[dof] = omega*inv_diag[dof]*accu + omega1*r[dof];
      }
    }

    if (false && !data->sym) {
      continue;  /* sor preconditioning */
    }

    /* backwards Gauss-Seidel iteration */
    for (dof = data->dim-1; dof >= 0; dof--) {
      if (data->matrix->matrix_row[dof] == NULL) {
	continue;
      }

      if ((!b || b[dof] < DIRICHLET)) { 

	accu = f[dof];

	FOR_ALL_MAT_COLS(REAL, data->matrix->matrix_row[dof], {
	    if (col_dof != dof) {
	      accu -= row->entry[col_idx] * r[col_dof];
	    }
	  });
	r[dof] = omega*data->inv_diag[dof]*accu + omega1*r[dof];
      }
    }
  }
}

static void exit_SSOR_precon_s(void *precon_data)
{
  struct SSOR_precon_s_data *data = (struct SSOR_precon_s_data *)precon_data;

  if (data->inv_diag) {
    MEM_FREE(data->inv_diag, data->max_dim, REAL);
    data->inv_diag = NULL;
    MEM_FREE(data->f, data->max_dim, REAL);
    data->f = NULL;

    data->max_dim = 0;
  }
  data->dim = 0;
}

static const PRECON *get_SSOR_precon_s(const DOF_MATRIX *A,
				       const DOF_SCHAR_VEC *bound,
				       REAL omega,
				       /* bool sym, */
				       int n_iter)
{
  static struct SSOR_precon_s_data *first = NULL;
  struct SSOR_precon_s_data        *data;

  for (data = first; data; data = data->next) {
    if (data->inv_diag == NULL && data->max_dim == 0) {
      break;
    }
  }

  if (!data) {
    data = MEM_CALLOC(1, struct SSOR_precon_s_data);
    data->next     = first;
    first          = data;
  }

  data->precon.precon_data = data;
  data->precon.init_precon = init_SSOR_precon_s;
  data->precon.precon      = SSOR_precon_s;
  data->precon.exit_precon = exit_SSOR_precon_s;

  data->omega  = omega;
  data->sym    = true /* sym */;
  data->n_iter = n_iter;
  data->matrix = A;
  data->bound  = bound;

  return &data->precon;
}

struct SSOR_precon_d_data
{
  PRECON              precon;
  REAL                omega;
  int                 sym;    /* 0: sor, 1: ssor */
  int                 n_iter;

  const DOF_MATRIX    *matrix;
  const DOF_SCHAR_VEC *bound;

  DOF                 dim;
  DOF                 max_dim;
  REAL_D              *inv_diag;
  REAL_D              *f;

  struct SSOR_precon_d_data *next;
};

static bool init_SSOR_precon_d(void *precon_data)
{
  FUNCNAME("init_SSOR_precon_d");
  struct SSOR_precon_d_data *data = (struct SSOR_precon_d_data *)precon_data;
  MATRIX_ROW              **row  = data->matrix->matrix_row;
  const S_CHAR            *b     = data->bound ? data->bound->vec : NULL;
  const DOF_ADMIN         *admin = data->matrix->row_fe_space->admin;
  int                     dim    = admin->size_used;
  REAL_D                  *inv_diag;
  int i;

  if (data->max_dim < dim) {
    data->inv_diag =
      MEM_REALLOC(data->inv_diag, data->max_dim, dim, REAL_D);
    data->f       =
      MEM_REALLOC(data->f, data->max_dim, dim, REAL_D);
    data->max_dim = dim;
  }

  data->dim = dim;
  inv_diag = (REAL_D *)data->inv_diag;

#undef MAT_BODY
#define MAT_BODY(F, CC, C, S, TYPE)					\
  if (b) {								\
    FOR_ALL_DOFS(admin, {						\
	SET_DOW(1.0, inv_diag[dof]);					\
	if (!row[dof] || b[dof] >= DIRICHLET) {				\
	  continue;							\
	}								\
	F##DIV_DOW(CC row[dof]->entry.S[0], inv_diag[dof], inv_diag[dof]); \
	for (i = 0; i < DIM_OF_WORLD; ++i) {				\
	  if (!isfinite(inv_diag[dof][i])) {				\
	    inv_diag[dof][i] = 1.0;					\
	  }								\
	}								\
      });								\
  } else {								\
    FOR_ALL_DOFS(admin, {						\
	SET_DOW(1.0, inv_diag[dof]);					\
	if (!row[dof]) {						\
	  continue;							\
	}								\
	F##DIV_DOW(CC row[dof]->entry.S[0], inv_diag[dof], inv_diag[dof]); \
	for (i = 0; i < DIM_OF_WORLD; ++i) {				\
	  if (!isfinite(inv_diag[dof][i])) {				\
	    inv_diag[dof][i] = 1.0;					\
	  }								\
	}								\
      });								\
  }
  MAT_EMIT_BODY_SWITCH(data->matrix->type);

  FOR_ALL_FREE_DOFS(admin, {
      if (dof >= admin->size_used) {
	break;
      }
      SET_DOW(1.0, inv_diag[dof]);
    });
  
  return true;
}

/* Just a copy of ssor_s() with precomputed inverted diagonal ... */
static void SSOR_precon_d(void *ud, int dim, REAL *_r)
{
  struct SSOR_precon_d_data *data = (struct SSOR_precon_d_data *)ud;
  const S_CHAR    *b        = data->bound ? data->bound->vec : NULL;  
  REAL_D          *inv_diag = data->inv_diag;
  REAL_D          *f        = data->f;
  REAL_D          *r        = (REAL_D *)_r;
  REAL            omega     = data->omega;
  REAL            omega1    = 1.0 - omega;
  REAL_D          accu;
  int             iter, d;
  DOF             dof;
  
  /* set initial guess to zero */
  for (dof = 0; dof < data->dim; dof ++) {
    if (data->matrix->matrix_row[dof] == NULL) {
      SET_DOW(0.0, f[dof]);
      continue;
    }
    if (!b || b[dof] < DIRICHLET) {
      COPY_DOW(r[dof], f[dof]);
      SET_DOW(0.0, r[dof]);
    }
  }

#undef MAT_BODY
#define MAT_BODY(F, CC, C, S, TYPE)					\
  for (iter = 0; iter < data->n_iter; iter++) {				\
									\
    /* forward Gauss-Seidel iteration */				\
    for (dof = 0; dof < data->dim; dof ++) {				\
      if (data->matrix->matrix_row[dof] == NULL) {			\
	continue;							\
      }									\
									\
      if (!b || b[dof] < DIRICHLET) {					\
	COPY_DOW(f[dof], accu);						\
									\
	FOR_ALL_MAT_COLS(TYPE, data->matrix->matrix_row[dof], {		\
	    if (col_dof != dof) {					\
	      F##GEMV_DOW(-1.0, CC row->entry[col_idx], r[col_dof],	\
			  1.0, accu);					\
	    } else {                                                    \
	      F##GEMV_ND_DOW(-1.0, CC row->entry[col_idx], r[col_dof],	\
                             1.0, accu);                                \
            }                                                           \
	  });                                                           \
	for (d = 0; d < DIM_OF_WORLD; d++) {				\
	  accu[d] *= inv_diag[dof][d];					\
	}								\
	AXPBY_DOW(omega, accu, omega1, r[dof], r[dof]);			\
      }									\
    }									\
									\
    if (false && !data->sym) {						\
      continue;  /* sor preconditioning */				\
    }									\
									\
    /* backwards Gauss-Seidel iteration */				\
    for (dof = data->dim-1; dof >= 0; dof--) {				\
      if (data->matrix->matrix_row[dof] == NULL) {			\
	continue;							\
      }									\
									\
      if ((!b || b[dof] < DIRICHLET)) {					\
									\
	COPY_DOW(f[dof], accu);						\
									\
	FOR_ALL_MAT_COLS(TYPE, data->matrix->matrix_row[dof], {		\
	    if (col_dof != dof) {					\
	      F##GEMV_DOW(-1.0, CC row->entry[col_idx], r[col_dof],	\
			  1.0, accu);					\
	    } else {                                                    \
              F##GEMV_ND_DOW(-1.0, CC row->entry[col_idx], r[col_dof],  \
                             1.0, accu);                               \
            }                                                           \
	  });								\
	for (d = 0; d < DIM_OF_WORLD; d++) {				\
	  accu[d] *= inv_diag[dof][d];					\
	}								\
	AXPBY_DOW(omega, accu, omega1, r[dof], r[dof]);			\
      }									\
    }									\
  }
  MAT_EMIT_BODY_SWITCH(data->matrix->type);
}

static void exit_SSOR_precon_d(void *precon_data)
{
  struct SSOR_precon_d_data *data = (struct SSOR_precon_d_data *)precon_data;

  if (data->inv_diag) {
    MEM_FREE(data->inv_diag, data->max_dim, REAL_D);
    data->inv_diag = NULL;
    MEM_FREE(data->f, data->max_dim, REAL_D);
    data->f = NULL;

    data->max_dim = 0;
  }
  data->dim = 0;
}

static const PRECON *get_SSOR_precon_d(const DOF_MATRIX *A,
				       const DOF_SCHAR_VEC *bound,
				       REAL omega,
				       /* bool sym, */
				       int n_iter)
{
  FUNCNAME("get_SSOR_precon_d");
  static struct SSOR_precon_d_data *first = NULL;
  struct SSOR_precon_d_data        *data;

  for (data = first; data; data = data->next) {
    if (data->inv_diag == NULL && data->max_dim == 0) {
      break;
    }
  }

  if (!data) {
    data = MEM_CALLOC(1, struct SSOR_precon_d_data);
    data->next     = first;
    first          = data;
  }

  data->precon.precon_data = data;
  data->precon.init_precon = init_SSOR_precon_d;
  data->precon.precon      = SSOR_precon_d;
  data->precon.exit_precon = exit_SSOR_precon_d;

  data->omega  = omega;
  data->sym    = true /* sym */;
  data->n_iter = n_iter;
  data->matrix = A;
  data->bound  = bound;

  return &data->precon;
}

/****************************************************************************/

const PRECON *get_SSOR_precon(const DOF_MATRIX *A,
			      const DOF_SCHAR_VEC *mask,
			      REAL omega, /* bool sym, */ int n_iter)
{
  if (A->is_diagonal) {
    /* makes no sense with diagonal matrices */
    return get_diag_precon(A, mask);
  }

  if (A->row_fe_space->rdim != 1 &&
      A->row_fe_space->bas_fcts->rdim == 1) {
    return get_SSOR_precon_d(A, mask, omega, /* sym, */ n_iter);
  } else {    
    return get_SSOR_precon_s(A, mask, omega, /* sym, */ n_iter);
  }
}
