/*******************************************************************************
 * ALBERTA:  an Adaptive multi Level finite element toolbox using
 *           Bisectioning refinement and Error control by Residual
 *           Techniques for scientific Applications
 *
 * file:     diag_precon.c
 *
 * description:  Diagonal precondioning for a DOF_MATRIX (of any type).
 *
 *******************************************************************************
 *
 *  authors:   Alfred Schmidt
 *             Zentrum fuer Technomathematik
 *             Fachbereich 3 Mathematik/Informatik
 *             Universitaet Bremen
 *             Bibliothekstr. 2
 *             D-28359 Bremen, Germany
 *
 *             Kunibert G. Siebert
 *             Institut fuer Mathematik
 *             Universitaet Augsburg
 *             Universitaetsstr. 14
 *             D-86159 Augsburg, Germany
 *
 *             Claus-Justus Heine
 *             Abteilung fuer Angewandte Mathematik
 *             Albert-Ludwigs-Universitaet Freiburg
 *             Hermann-Herder-Str. 10
 *             D-79104 Freiburg im Breisgau, Germany
 *
 *  http://www.alberta-fem.de
 *
 *  (c) by A. Schmidt and K.G. Siebert (1996-2003),
 *         C.-J. Heine (1998-2009)
 *
 ******************************************************************************/

#ifdef HAVE_CONFIG_H
# include "config.h"
#endif

#include "alberta_intern.h"
#include "alberta.h"

struct diag_precon_data
{
  PRECON              precon;

  const DOF_MATRIX    *matrix;
  const DOF_SCHAR_VEC *mask;
  size_t              dim;

  const REAL          *diag_1;

  struct diag_precon_data *next;
};

/* Perform the actual precondioning: scale by the inverse of the diagonal. */
static void diag_precon(void *ud, int dim, REAL *r)
{
  FUNCNAME("diag_precon_s");
  struct diag_precon_data *data = (struct diag_precon_data *)ud;
  const REAL *diag_1 = data->diag_1;
  int i;

  DEBUG_TEST_EXIT(dim == (int)data->dim,
		  "argument dim =% d != diag_precon_data->dim = %d\n",
		  dim, data->dim);

  for (i = 0; i < dim; i++) {
    r[i] *= diag_1[i];
  }
}

static inline
void __init_diag_precon(DOF_REAL_VEC_D *dv_vec,
			const DOF_MATRIX *A,
			const DOF_SCHAR_VEC *mask)
{
  const FE_SPACE *fe_space = A->row_fe_space;
  DOF i, dof, size_used    = fe_space->admin->size_used;
  
  if (fe_space->rdim != 1 && fe_space->bas_fcts->rdim == 1) {
    /* Potential DOWxDOW matrix case */
    REAL_D *diag_1 = (REAL_D *)dv_vec->vec;
    switch (A->type) {
    case MATENT_REAL_DD: {
      if (A->is_diagonal) {
	DOF     *diag_cols = A->diag_cols->vec;
	REAL_DD *diag      = A->diagonal.real_dd->vec;

	for (dof = 0; dof < size_used; dof++) {
	  if (ENTRY_USED(diag_cols[dof]) &&
	      (mask == NULL || mask->vec[dof] < DIRICHLET)) {
	    for (i = 0; i < DIM_OF_WORLD; i++) {
	      if (ABS(diag[dof][i][i]) > 1.0E-20) {
		diag_1[dof][i] = fabs(1.0/diag[dof][i][i]);
	      } else {
		diag_1[dof][i] = 1.0;
	      }
	    }
	  } else {
	    SET_DOW(1.0, diag_1[dof]);
	  }
	}
      } else {
	MATRIX_ROW_REAL_DD **row = (MATRIX_ROW_REAL_DD **)A->matrix_row;

	for (dof = 0; dof < size_used; dof++) {
	  if (row[dof] && (mask == NULL || mask->vec[dof] < DIRICHLET)) {
	    for (i = 0; i < DIM_OF_WORLD; i++) {
	      if (ABS(row[dof]->entry[0][i][i]) > 1.0E-20) {
		diag_1[dof][i] = fabs(1.0/row[dof]->entry[0][i][i]);
	      } else {
		diag_1[dof][i] = 1.0;
	      }
	    }
	  } else {
	    SET_DOW(1.0, diag_1[dof]);
	  }
	}
      }
      break;
    }
    case MATENT_REAL_D: {
      if (A->is_diagonal) {
	DOF    *diag_cols = A->diag_cols->vec;
	REAL_D *diag      = A->diagonal.real_d->vec;

	for (dof = 0; dof < size_used; dof++) {
	  if (ENTRY_USED(diag_cols[dof]) &&
	      (mask == NULL || mask->vec[dof] < DIRICHLET)) {
	    for (i = 0; i < DIM_OF_WORLD; i++) {
	      if (ABS(diag[dof][i]) > 1.0E-20) {
		diag_1[dof][i] = fabs(1.0/diag[dof][i]);
	      } else {
		diag_1[dof][i] = 1.0;
	      }
	    }
	  } else {
	    SET_DOW(1.0, diag_1[dof]);
	  }
	}
      } else {
	MATRIX_ROW_REAL_D **row = (MATRIX_ROW_REAL_D **)A->matrix_row;

	for (dof = 0; dof < size_used; dof++) {
	  if (row[dof] && (mask == NULL|| mask->vec[dof] < DIRICHLET)) {
	    for (i = 0; i < DIM_OF_WORLD; i++) {
	      if (ABS(row[dof]->entry[0][i]) > 1.0E-20) {
		diag_1[dof][i] = fabs(1.0/row[dof]->entry[0][i]);
	      } else {
		diag_1[dof][i] = 1.0;
	      }
	    }
	  } else {
	    SET_DOW(1.0, diag_1[dof]);
	  }
	}
      }
      break;
    }
    case MATENT_REAL: {
      if (A->is_diagonal) {
	DOF  *diag_cols = A->diag_cols->vec;
	REAL *diag      = A->diagonal.real->vec;

	for (dof = 0; dof < size_used; dof++) {
	  if (ENTRY_USED(diag_cols[dof]) &&
	      (mask == NULL || mask->vec[dof] < DIRICHLET) && 
	      ABS(diag[dof]) > 1.0E-20) {
	    SET_DOW(fabs(1.0/diag[dof]), diag_1[dof]);
	  } else {
	    SET_DOW(1.0, diag_1[dof]);
	  }
	}
      } else {
	MATRIX_ROW_REAL **row = (MATRIX_ROW_REAL **)A->matrix_row;

	for (dof = 0; dof < size_used; dof++) {
	  if (row[dof] && (mask == NULL || mask->vec[dof] < DIRICHLET) &&
	      ABS(row[dof]->entry[0]) > 1.0E-20) {
	    SET_DOW(fabs(1.0/row[dof]->entry[0]), diag_1[dof]);
	  } else {
	    SET_DOW(1.0, diag_1[dof]);
	  }
	}
      }
      break;
    }
    default:
      ERROR_EXIT("Unknown or invalid MATENT_TYPE: %d\n", A->type);
    }
  } else {
    /* ordinary scalar case */
    if (A->is_diagonal) {
      DOF *diag_cols = A->diag_cols->vec;
      REAL *diag     = A->diagonal.real->vec;

      for (dof = 0; dof < size_used; dof++) {
	if (ENTRY_USED(diag_cols[dof])  &&
	    (mask == NULL || mask->vec[dof] < DIRICHLET) && 
	    ABS(diag[dof]) > 1.0E-20) {
	  dv_vec->vec[dof] = fabs(1.0/diag[dof]);
	} else {
	  dv_vec->vec[dof] = 1.0;
	}
      }
    } else {
      MATRIX_ROW_REAL **row = (MATRIX_ROW_REAL **)A->matrix_row;

      for (dof = 0; dof < size_used; dof++) {
	if (row[dof] && (mask == NULL || mask->vec[dof] < DIRICHLET) && 
	    ABS(row[dof]->entry[0]) > 1.0E-20) {
	  dv_vec->vec[dof] = fabs(1.0/A->matrix_row[dof]->entry.real[0]);
	} else {
	  dv_vec->vec[dof] = 1.0;
	}
      }
    }
  }
}

/* Initialize diagonal preconditioning for all kind of matrices. Now
 * that the fe-spaces carry a dimension we do no longer need separate
 * versions for scalar/vector valued problems.
 */
static bool init_diag_precon(void *diag_precon_data)
{
  FUNCNAME("init_diag_precon_s");
  struct diag_precon_data *data = (struct diag_precon_data *)diag_precon_data;
  const DOF_SCHAR_VEC *mask     = data->mask;
  const DOF_MATRIX    *A        = data->matrix;
  const FE_SPACE      *fe_space = A->row_fe_space;
  int                 dim       = dof_real_vec_d_length(fe_space);
  DOF_REAL_VEC_D      vecs[CHAIN_LENGTH(fe_space)], *dv_vec;

  if (data->dim < (size_t)dim) {
    data->diag_1 = MEM_REALLOC(data->diag_1, data->dim, dim, REAL);
    data->dim = dim;
  }

  dv_vec = init_dof_real_vec_d_skel(vecs, "diag skel", fe_space);
  distribute_to_dof_real_vec_d_skel(dv_vec, data->diag_1);

  CHAIN_DO(fe_space, const FE_SPACE) {
    if (mask) {
      __init_diag_precon(dv_vec, A, mask);
    } else {
      __init_diag_precon(dv_vec, A, NULL);
    }
    mask = mask ? CHAIN_NEXT(mask, DOF_SCHAR_VEC) : NULL;
    dv_vec = CHAIN_NEXT(dv_vec, DOF_REAL_VEC_D);
    A = COL_CHAIN_NEXT(A, const DOF_MATRIX);
    A = ROW_CHAIN_NEXT(A, const DOF_MATRIX);
  } CHAIN_WHILE(fe_space, const FE_SPACE);

  FOREACH_FREE_DOF_DOW(fe_space,
		       if (dof >= fe_space->admin->size_used) {
			 continue;
		       }
		       dv_vec->vec[dof] = 0.0,
		       if (dof >= fe_space->admin->size_used){
			 continue;
		       }
		       SET_DOW(0.0, ((REAL_D *)dv_vec->vec)[dof]),
		       dv_vec = CHAIN_NEXT(dv_vec, DOF_REAL_VEC_D));

  return true;
}

static void exit_diag_precon(void *diag_precon_data)
{
  struct diag_precon_data *data = (struct diag_precon_data *)diag_precon_data;

  if (data->diag_1) {
    MEM_FREE(data->diag_1, data->dim, REAL);
    data->diag_1 = NULL;
  }
  MEM_FREE(data, 1, diag_precon_data);
}

const PRECON *get_diag_precon(const DOF_MATRIX *A, const DOF_SCHAR_VEC *mask)
{
  struct diag_precon_data *data;
  
  TEST_EXIT(FE_SPACE_EQ_P(A->row_fe_space, A->col_fe_space),
	    "Row and column FE_SPACEs don't match!\n");
  
  data = MEM_CALLOC(1, struct diag_precon_data);

  data->precon.precon_data = data;
  data->precon.init_precon = init_diag_precon;
  data->precon.precon      = diag_precon;
  data->precon.exit_precon = exit_diag_precon;

  data->matrix = A;
  data->mask   = mask;
  data->dim    = 0; /* is already 0 */

  return &data->precon;
}

