/*******************************************************************************
 * ALBERTA:  an Adaptive multi Level finite element toolbox using
 *           Bisectioning refinement and Error control by Residual
 *           Techniques for scientific Applications
 *
 * file:     block_precon.c
 *
 * description: Composition of a block-preconditioner from
 *              preconditioner for the single blocks.
 *
 *******************************************************************************
 *
 *  authors:   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 C.-J. Heine (2009)
 *
 ******************************************************************************/

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

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

typedef struct precon_node PRECON_NODE;
struct precon_node
{
  const PRECON   *precon;
  OEM_PRECON     type;
  size_t         dim;
  DOF_MATRIX     *A;      /* unchained flat copy of the original matrix block */
  DOF_SCHAR_VEC  *mask;
  DOF_REAL_VEC_D *accu;  /* for BlkSSORPrecon */
  DOF_MATRIX     *A_row; /* Row without diagonal block for SSORPrecon */
  DBL_LIST_NODE  chain;
};

typedef struct block_precon_data
{
  PRECON              precon;

  const DOF_MATRIX    *matrix;
  const DOF_SCHAR_VEC *mask;
  size_t              dim;
  
  OEM_PRECON          block_precon; /* how to combine the individual precons */
  DBL_LIST_NODE       chain;        /* the list of the individual precons */

  REAL                omega;  /* for BlkSSORPrecon */
  int                 n_iter; /* for BlkSSORPrecon */

  DOF_REAL_VEC_D      *rhs;    /* for BlkSSORPrecon */
  DOF_REAL_VEC_D      *r_skel; /* for BLKSSORPrecon */

  SCRATCH_MEM         scratch;
} BLOCK_PRECON_DATA;

/* Simplest method: just do block-diagonal precon.
 *
 *
 * P = diag(P_1, P_2, ...) where P_1, ... are listed in precon_chain.
 * 
 * Maybe we could give block-SSOR a chance:
 *
 * A = (a_{ij}) where a_{ij} are matrices, the precons listed in
 * precon_chain serve as approximate inverse to a_{ii}
 *
 */
static void block_diag_precon(void *pd, int dim, REAL *r)
{
  BLOCK_PRECON_DATA *data = (struct block_precon_data *)pd;
  struct precon_node *prec;
  
  CHAIN_FOREACH(prec, data, struct precon_node) {
    if (prec->type != NoPrecon) {
      prec->precon->precon(prec->precon->precon_data, prec->dim, r);
    }
    r += prec->dim;
  }
}

static void block_diag_exit_precon(void *pd)
{
  BLOCK_PRECON_DATA *data = (struct block_precon_data *)pd;
  struct precon_node *prec;
  
  CHAIN_FOREACH(prec, data, struct precon_node) {
    if (prec->type != NoPrecon) {
      prec->precon->exit_precon(prec->precon->precon_data);
    }
  }

  /* Release all our resources. */
  SCRATCH_MEM_ZAP(data->scratch);
}

static bool block_diag_init_precon(void *pd)
{
  BLOCK_PRECON_DATA *data = (struct block_precon_data *)pd;
  struct precon_node *prec;
  
  CHAIN_FOREACH(prec, data, struct precon_node) {
    if (prec->type != NoPrecon) {
      update_dof_matrix_sub_chain(prec->A);
      if (prec->mask) {
	update_dof_schar_vec_sub_chain(prec->mask);
      }
      if (!prec->precon->init_precon(prec->precon->precon_data)) {
	return false;
      }
      prec->dim = dof_real_vec_d_length(prec->A->row_fe_space);
    }
  }
  return true;
}

/* Use a block SSOR method, using preconditioners for the diagonal
 * elements to approximate the inverse of the diagonal blocks.
 */
static void block_SSOR_precon(void *pd, int dim, REAL *r)
{
  BLOCK_PRECON_DATA *data = (struct block_precon_data *)pd;
  struct precon_node *prec;
  int iter;
  REAL *rbase = r;

  /* FIXME: do we need to mask out boundary DOFs? */
  distribute_to_dof_real_vec_d_skel(data->r_skel, r);
  dof_copy_dow(data->r_skel, data->rhs);
  dset(dim, 0.0, r, 1);

  for (iter = 0; iter < data->n_iter; iter++) {
    /* Forward Gauss-Seidel iteration */
    r = rbase;
    CHAIN_FOREACH(prec, data, struct precon_node) {
      data->r_skel = CHAIN_NEXT(data->r_skel, DOF_REAL_VEC_D);
      dcopy(prec->dim, data->rhs->vec, 1, prec->accu->vec, 1);
      dof_gemv_dow(NoTranspose,
		   -1.0, prec->A_row, data->mask, data->r_skel,
		   1.0, prec->accu);

      /* Pseudo-invert the diagonal block */
      if (prec->type != NoPrecon) {
	prec->precon->precon(
	  prec->precon->precon_data, prec->dim, prec->accu->vec);
      }

      /* r[idx] = omega A_diag^{-1} accu + (1-omega) r */
      dscal(prec->dim, data->omega, prec->accu->vec, 1);
      dxpay(prec->dim, prec->accu->vec, 1, 1.0 - data->omega, r, 1);

      data->rhs  = CHAIN_NEXT(data->rhs, DOF_REAL_VEC_D);
      data->mask = data->mask ? CHAIN_NEXT(data->mask, DOF_SCHAR_VEC) : NULL;
      r += prec->dim;
    }
    /* Backward Gauss-Seidel iteration */
    r = rbase + dim;
    CHAIN_FOREACH_REV(prec, data, struct precon_node) {
      r -= prec->dim;
      data->rhs  = CHAIN_PREV(data->rhs, DOF_REAL_VEC_D);
      data->mask = data->mask ? CHAIN_PREV(data->mask, DOF_SCHAR_VEC) : NULL;
      dcopy(prec->dim, data->rhs->vec, 1, prec->accu->vec, 1);
      dof_gemv_dow(NoTranspose,
		   -1.0, prec->A_row, data->mask, data->r_skel,
		   1.0, prec->accu);

      /* r[idx] = omega A_diag^{-1} accu + (1-omega) r */
      if (prec->type != NoPrecon) {
	prec->precon->precon(
	  prec->precon->precon_data, prec->dim, prec->accu->vec);
      }
      dscal(prec->dim, data->omega, prec->accu->vec, 1);
      dxpay(prec->dim, prec->accu->vec, 1, 1.0 - data->omega, r, 1);

      data->r_skel = CHAIN_PREV(data->r_skel, DOF_REAL_VEC_D);
    }
  }
}

static void block_SSOR_exit_precon(void *pd)
{
  BLOCK_PRECON_DATA *data = (struct block_precon_data *)pd;
  struct precon_node *prec;
  
  CHAIN_FOREACH(prec, data, struct precon_node) {
    if (prec->type != NoPrecon) {
      prec->precon->exit_precon(prec->precon->precon_data);
    }
    free_dof_real_vec_d(prec->accu);
  }

  free_dof_real_vec_d(data->rhs);

  /* Release all our resources. */
  SCRATCH_MEM_ZAP(data->scratch);
}

static bool block_SSOR_init_precon(void *pd)
{
  BLOCK_PRECON_DATA *data = (struct block_precon_data *)pd;
  struct precon_node *prec;
  
  CHAIN_FOREACH(prec, data, struct precon_node) {
    update_dof_matrix_sub_chain(prec->A);
    update_dof_matrix_sub_chain(prec->A_row);
    if (prec->mask) {
      update_dof_schar_vec_sub_chain(prec->mask);
    }
    if (!prec->precon->init_precon(prec->precon->precon_data)) {
      return false;
    }
    prec->dim = dof_real_vec_d_length(prec->A->row_fe_space);
  }
  return true;
}

const PRECON *_AI_vget_block_diag_precon(const DOF_MATRIX *A,
					 const DOF_SCHAR_VEC *mask,
					 int info,
					 va_list ap)
{
  FUNCNAME("_AI_vget_block_diag_precon");
  PRECON_TYPE precon;
  OEM_PRECON sub_type = NoPrecon;
  int i, nblocks = COL_CHAIN_LENGTH(A);

  precon.type = BlkDiagPrecon;
  
  for (i = 0; i < nblocks && sub_type != PreconRepeat; i++) {
    TEST_EXIT(i < N_BLOCK_PRECON_MAX,
	      "Sorry, only up to %d x %d blocks are supported.\n",
	      N_BLOCK_PRECON_MAX, N_BLOCK_PRECON_MAX);
    precon.param.BlkDiag.precon[i].type =
      sub_type = (OEM_PRECON)va_arg(ap, int);
    if (sub_type == __SSORPrecon) {
      precon.param.BlkDiag.precon[i].param.__SSOR.omega  = va_arg(ap, REAL);
      precon.param.BlkDiag.precon[i].param.__SSOR.n_iter = va_arg(ap, int);
    }
  }

  return _AI_get_block_precon(A, mask, info, &precon);
}

const PRECON *_AI_vget_block_SSOR_precon(const DOF_MATRIX *A,
					 const DOF_SCHAR_VEC *mask,
					 int info,
					 va_list ap)
{
  FUNCNAME("_AI_vget_block_diag_precon");
  PRECON_TYPE precon;
  OEM_PRECON sub_type = NoPrecon;
  int i, nblocks = COL_CHAIN_LENGTH(A);
  
  precon.type = BlkSSORPrecon;
  precon.param.BlkSSOR.omega  = va_arg(ap, REAL);
  precon.param.BlkSSOR.n_iter = va_arg(ap, int);

  for (i = 0; i < nblocks && sub_type != PreconRepeat; i++) {
    TEST_EXIT(i < N_BLOCK_PRECON_MAX,
	      "Sorry, only up to %d x %d blocks are supported.\n",
	      N_BLOCK_PRECON_MAX, N_BLOCK_PRECON_MAX);
    precon.param.BlkDiag.precon[i].type =
      sub_type = (OEM_PRECON)va_arg(ap, int);
    if (sub_type == __SSORPrecon) {
      precon.param.BlkDiag.precon[i].param.__SSOR.omega  = va_arg(ap, REAL);
      precon.param.BlkDiag.precon[i].param.__SSOR.n_iter = va_arg(ap, int);
    }
  }

  return _AI_get_block_precon(A, mask, info, &precon);
}

const PRECON *_AI_get_block_diag_precon(const DOF_MATRIX *A,
					const DOF_SCHAR_VEC *mask,
					int info,
					...)
{
  const PRECON *precon;
  va_list ap;
  
  va_start(ap, info);
  precon = _AI_vget_block_diag_precon(A, mask, info, ap);
  va_end(ap);

  return precon;
}

const PRECON *_AI_get_block_precon(const DOF_MATRIX *A,
				   const DOF_SCHAR_VEC *mask,
				   int info,
				   const PRECON_TYPE *prec_type)
{
  FUNCNAME("_AI_get_block_precon");
  BLOCK_PRECON_DATA *data;
  SCRATCH_MEM scrmem;
  bool repeat = false;
  OEM_PRECON sub_type = NoPrecon;
  int i = 0;
  const FE_SPACE *col_fe_space;
  
  TEST_EXIT(ROW_CHAIN_LENGTH(A) == COL_CHAIN_LENGTH(A),
	    "Makes sense for quadratic block-matrices only.\n");
  TEST_EXIT(ROW_CHAIN_LENGTH(A) < N_BLOCK_PRECON_MAX,
	    "Only implemented for up to %d x %d blocks.\n",
	    N_BLOCK_PRECON_MAX, N_BLOCK_PRECON_MAX);

  SCRATCH_MEM_INIT(scrmem);
  data = SCRATCH_MEM_ALLOC(scrmem, 1, BLOCK_PRECON_DATA);
  memset(data, 0, sizeof(*data));
  SCRATCH_MEM_CPY(data->scratch, scrmem);

  CHAIN_INIT(data);

  data->matrix = A;
  data->mask   = mask;
  data->dim    = dof_real_vec_d_length(A->row_fe_space);

  data->precon.precon_data = data;

  col_fe_space = A->col_fe_space ? A->col_fe_space : A->row_fe_space;

  switch (prec_type->type) {
  case BlkDiagPrecon:
    data->block_precon       = DiagPrecon;
    data->precon.precon      = block_diag_precon;
    data->precon.init_precon = block_diag_init_precon;
    data->precon.exit_precon = block_diag_exit_precon;
    break;
  case BlkSSORPrecon:
    data->block_precon       = SSORPrecon;
    data->precon.precon      = block_SSOR_precon;
    data->precon.init_precon = block_SSOR_init_precon;
    data->precon.exit_precon = block_SSOR_exit_precon;
    data->omega              = prec_type->param.BlkSSOR.omega;
    data->n_iter             = prec_type->param.BlkSSOR.n_iter;
    
    data->rhs = get_dof_real_vec_d("SSOR rhs", col_fe_space);
    data->r_skel = init_dof_real_vec_d_skel(
      SCRATCH_MEM_ALLOC(scrmem, CHAIN_LENGTH(col_fe_space), DOF_REAL_VEC_D),
      "SSOR r skeleton",
      col_fe_space);

    break;
  default:
    ERROR_EXIT("Precon type %d is not implemented.\n", prec_type->type);
    data->block_precon = NoPrecon;
  }

  COL_CHAIN_DO(A, const DOF_MATRIX) {
    REAL omega = 1.0;
    int  n_iter = 2;
    int  ilu_level = 0;
    PRECON_NODE *node = SCRATCH_MEM_ALLOC(scrmem, 1, PRECON_NODE);
    memset(node, 0, sizeof(*node));
    CHAIN_INIT(node);

    CHAIN_ADD_TAIL(data, node);

    /* Generate flat, unchained copies from the diagonal block and
     * boundary masks. This way one can use the ordinary
     * preconditioners with those flat copies.
     */
    node->A    = dof_matrix_sub_chain(scrmem, A, 0x1, 0x1);
    node->mask = mask ? dof_schar_vec_sub_chain(scrmem, mask, 0x1) : NULL;

    if (data->block_precon == SSORPrecon) {
      /* Extract a sub-chain with all elements of the row safe the diagonal */
      node->A_row = dof_matrix_sub_chain(scrmem, A, 0x1, ~0x1);
      /* We need some scratch memory for the rhs and so on. */
      node->accu = get_dof_real_vec_d("SSOR accu", col_fe_space->unchained);
    }

    /* Compute the size of our fraction of the index set. */
    node->dim  = dof_real_vec_d_length(node->A->row_fe_space);

    if (!repeat && prec_type->param.BlkDiag.precon[i].type == PreconRepeat) {
      repeat = true;
    }
    if (repeat) {
      node->type = sub_type;
    } else {
      node->type = sub_type = prec_type->param.BlkDiag.precon[i].type;
    }
    switch (node->type) {
    case NoPrecon:
      break;
    case DiagPrecon:
      node->precon = get_diag_precon(node->A, node->mask);
      break;
    case HBPrecon:
      node->precon =
	get_HB_precon(node->A, node->mask, info);
      break;
    case BPXPrecon:
      node->precon =
	get_BPX_precon(node->A, node->mask, info);
      break;
    case __SSORPrecon:
      if (!repeat) {
	omega = prec_type->param.BlkDiag.precon[i].param.__SSOR.omega;
	n_iter = prec_type->param.BlkDiag.precon[i].param.__SSOR.n_iter;
      
	TEST(0.0 <= omega && omega <= 2.0,
	     "SSORPrecon: omega = %e???\n", omega);
	TEST(0 <= n_iter && n_iter < 1000,
	     "SSORPrecon: #iter = %d???\n", n_iter);
      }
    case SSORPrecon:
      node->precon = get_SSOR_precon(node->A, node->mask, omega, n_iter);
      break;
    case ILUkPrecon:
      if (!repeat) {
	ilu_level = prec_type->param.BlkDiag.precon[i].param.ILUk.level;
      }
      node->precon = get_ILUk_precon(node->A, node->mask, ilu_level, info);
      break;
    default:
      ERROR("Unknow precon-type %d, ignoring it.\n");
      node->type = NoPrecon;
      break;
    }
    
    i++;
    
    A = ROW_CHAIN_NEXT(A, const DOF_MATRIX); /* walk over the diagonal */
    col_fe_space = CHAIN_NEXT(col_fe_space, const FE_SPACE);
  } COL_CHAIN_WHILE(A, const DOF_MATRIX);
  
  return &data->precon;
}
