/*******************************************************************************
 * ALBERTA:  an Adaptive multi Level finite element toolbox using
 *           Bisectioning refinement and Error control by Residual
 *           Techniques for scientific Applications
 *
 * file:     oem_solve.c
 *
 * description: Part of the OEM interface which does not depend on the
 *              DIM_OF_WORLD dimensionality of a DOF_REAL[_D]_VEC[_D]
 *
 *******************************************************************************
 *
 *  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 mv_data
{
  const DOF_MATRIX    *matrix;
  MatrixTranspose     transpose;
  const FE_SPACE      *domain_space;
  const FE_SPACE      *range_space;
  int                 dim;

  const DOF_SCHAR_VEC *mask;

  DOF_REAL_VEC_D      *x_skel;
  DOF_REAL_VEC_D      *y_skel;

  SCRATCH_MEM         scratch;
};

/*---8<---------------------------------------------------------------------*/
/*---   y <- A x  or y <- A^T x                                          ---*/
/*--------------------------------------------------------------------->8---*/

int oem_mat_vec(void *ud, int dim, const REAL *x, REAL *y)
{
  FUNCNAME("mat_vec_s");
  struct mv_data *data = (struct mv_data *)ud;
  DOF_REAL_VEC_D *dof_x = data->x_skel;
  DOF_REAL_VEC_D *dof_y = data->y_skel;

  TEST_EXIT(dim == data->dim, "argument dim != FE_SPACE dim\n");

  distribute_to_dof_real_vec_d_skel(data->x_skel, x);
  distribute_to_dof_real_vec_d_skel(data->y_skel, y);

  dof_mv_dow(data->transpose, data->matrix, data->mask, dof_x, dof_y);

  return 0; /* ??? cH asks: what the heck is the meaning of the return value? */
}

OEM_MV_FCT init_oem_mat_vec(void **datap,
			    MatrixTranspose transpose, const DOF_MATRIX *A, 
			    const DOF_SCHAR_VEC *mask)
{
  /*FUNCNAME("init_oem_mat_vec");*/
  struct mv_data *data;
  SCRATCH_MEM scrmem;

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

  data->matrix    = A;
  data->transpose = transpose;

  data->mask     = mask;

  if (transpose == NoTranspose) {
    data->range_space  = A->row_fe_space;
    data->domain_space = A->col_fe_space;
  } else {
    data->range_space  = A->col_fe_space;
    data->domain_space = A->row_fe_space;
  }

  data->x_skel = init_dof_real_vec_d_skel(
    SCRATCH_MEM_ALLOC(scrmem, CHAIN_LENGTH(data->domain_space), DOF_REAL_VEC_D),
    "x skel",
    data->domain_space);
  data->y_skel = init_dof_real_vec_d_skel(
    SCRATCH_MEM_ALLOC(scrmem, CHAIN_LENGTH(data->range_space), DOF_REAL_VEC_D),
    "y skel",
    data->range_space);

  data->dim = dof_real_vec_d_length(data->range_space);

  *datap = data;
  
  return oem_mat_vec;
}

void exit_oem_mat_vec(void *vdata)
{
  struct mv_data *data = (struct mv_data *)vdata;
  
  SCRATCH_MEM_ZAP(data->scratch);
}

const PRECON *vinit_oem_precon(const DOF_MATRIX *A,
			       const DOF_SCHAR_VEC *mask,
			       int info, OEM_PRECON prec_type,
			       va_list ap)
{
  const PRECON *precon;
  REAL omega = 1.0;
  int n_iter = 2;
  int ilu_level = 0;

  if ((!ROW_CHAIN_SINGLE(A) || !COL_CHAIN_SINGLE(A))
      &&
      prec_type < BlkDiagPrecon) {
    precon = _AI_get_block_diag_precon(A, mask, info, prec_type, PreconRepeat);
  } else {
    switch(prec_type) {
    case NoPrecon: 
      precon = NULL;
      break;
    case DiagPrecon:
      precon = get_diag_precon(A, mask);
      break;
    case HBPrecon:
      TEST_EXIT(ROW_CHAIN_SINGLE(A) && COL_CHAIN_SINGLE(A),
		"HB-preonditioner does not make sense for "
		"horizontal direct sums\n");
      precon = get_HB_precon(A, mask, info);
      break;
    case BPXPrecon:
      TEST_EXIT(ROW_CHAIN_SINGLE(A) && COL_CHAIN_SINGLE(A),
		"BPX-preonditioner does not make sense for "
		"horizontal direct sums\n");
      precon = get_BPX_precon(A, mask, info);
      break;
    case __SSORPrecon:
      omega  = va_arg(ap, REAL);
      n_iter = va_arg(ap, int);

      TEST(0.0 <= omega && omega <= 2.0,
	   "SSORPrecon: omega = %e???\n", omega);
      TEST(0 <= n_iter && n_iter < 10,
	   "SSORPrecon: #iter = %d???\n", n_iter);
    case SSORPrecon:
      TEST_EXIT(ROW_CHAIN_SINGLE(A) && COL_CHAIN_SINGLE(A),
		"SSOR-preconditioner not implemented for horizontal "
		"direct sums. Very sorry.\n");
      precon = get_SSOR_precon(A, mask, omega, n_iter);
      break;
    case ILUkPrecon:
      ilu_level = va_arg(ap, int);
      precon = get_ILUk_precon(A, mask, ilu_level, info);
      break;
    case BlkDiagPrecon:
      precon = _AI_vget_block_diag_precon(A, mask, info, ap);
      break;
    case BlkSSORPrecon:
      precon = _AI_vget_block_SSOR_precon(A, mask, info, ap);
      break;
    default:
      ERROR_EXIT("Unknown precon type: %d\n", prec_type);
      break;
    }
  }

  return precon;
}

const PRECON *init_oem_precon(const DOF_MATRIX *A,
			      const DOF_SCHAR_VEC *mask,
			      int info, OEM_PRECON prec_type,
			      ... /* ssor_omega, ssor_n_iter etc. */)
{
  const PRECON *precon;
  va_list ap;
  
  va_start(ap, prec_type);
  precon = vinit_oem_precon(A, mask, info, prec_type, ap);
  va_end(ap);

  return precon;
}

const PRECON *init_precon_from_type(const DOF_MATRIX *A,
				    const DOF_SCHAR_VEC *mask,
				    int info,
				    const PRECON_TYPE *prec_type)
{
  FUNCNAME("init_precon_from_type");
  
  const PRECON *precon;
  
  switch (prec_type->type) {
  case __SSORPrecon:
    precon = init_oem_precon(A, mask, info,
			     __SSORPrecon,
			     prec_type->param.__SSOR.omega,
			     prec_type->param.__SSOR.n_iter);
    break;
  case ILUkPrecon:
    precon =
      init_oem_precon(A, mask, info, ILUkPrecon, prec_type->param.ILUk.level);
    break;
  case BlkSSORPrecon:
  case BlkDiagPrecon:
    precon = _AI_get_block_precon(A, mask, info, prec_type);
    break;    
  default:
    precon = init_oem_precon(A, mask, info, prec_type->type, PreconEnd);
    break;
  }

  return precon;
}

OEM_DATA *init_oem_solve(const DOF_MATRIX *A,
			 const DOF_SCHAR_VEC *mask,
			 REAL tol, const PRECON *precon,
			 int restart, int max_iter, int info)
{
  FUNCNAME("init_oem_solve");
  OEM_DATA     *oem;  
  const MatrixTranspose transpose = NoTranspose;

  TEST_EXIT(FE_SPACE_EQ_P(A->row_fe_space, A->col_fe_space),
	    "Row and column FE_SPACEs don't match!\n");

  oem = MEM_CALLOC(1, OEM_DATA); 
  oem->mat_vec = init_oem_mat_vec(&oem->mat_vec_data, transpose, A, mask);

  if (precon) {
    if (precon->init_precon && !(*precon->init_precon)(precon->precon_data)) {
      precon = NULL;
      MSG("init_precon() failed, disabling preconditioner!\n");
    } else {
      oem->left_precon_data = precon->precon_data;
      oem->left_precon = precon->precon;
    }
  }

  oem->ws        = NULL; /* Let the solvers handle this point themselves. */
  oem->tolerance = tol;
  oem->restart   = restart;
  oem->max_iter  = max_iter;
  oem->info      = MAX(0, info);

  return oem;
}

void release_oem_solve(const OEM_DATA *_oem)
{
  OEM_DATA *oem = (OEM_DATA *)_oem;
  const PRECON *precon;
  
  exit_oem_mat_vec(oem->mat_vec_data);
  if ((precon = (PRECON *)oem->left_precon_data) != NULL &&
      precon && precon->exit_precon) {
    precon->exit_precon((void *)precon);
  }
  MEM_FREE(oem, 1, OEM_DATA);
}

OEM_MV_FCT get_oem_solver(OEM_SOLVER solver)
{

  switch (solver) {
  case BiCGStab:
    return (OEM_MV_FCT)oem_bicgstab;
  case CG:
    return (OEM_MV_FCT)oem_cg;
  case TfQMR:
    return (OEM_MV_FCT)oem_tfqmr;
  case GMRes:
    return (OEM_MV_FCT)oem_gmres;
  case GMRes_k:
    return (OEM_MV_FCT)oem_gmres_k;
  case ODir:
    return (OEM_MV_FCT)oem_odir;
  case ORes:
    return (OEM_MV_FCT)oem_ores;
  case SymmLQ:
    return (OEM_MV_FCT)oem_symmlq;
  default:
    ERROR_EXIT("unknown OEM solver %d\n", (int) solver);
    return NULL;
  }
}

/* scalar problems */

int call_oem_solve_s(const OEM_DATA *_oem, OEM_SOLVER solver, 
		     const DOF_REAL_VEC *f, DOF_REAL_VEC *u)
{
  int iter, restart, dim;
  OEM_DATA *oem = (OEM_DATA *)_oem;
  REAL *uvec, *fvec;

  TEST_EXIT(FE_SPACE_EQ_P(f->fe_space, u->fe_space),
	    "Row and column FE_SPACEs don't match!\n");

  dim = dof_real_vec_length(f->fe_space);

  if (!CHAIN_SINGLE(u)) {
    uvec = MEM_ALLOC(dim, REAL);
    fvec = MEM_ALLOC(dim, REAL);
    copy_from_dof_real_vec(uvec, u);
    copy_from_dof_real_vec(fvec, f);
  } else {
    FOR_ALL_FREE_DOFS(u->fe_space->admin,
		      if (dof < dim) u->vec[dof] = f->vec[dof] = 0.0);
    fvec = f->vec;
    uvec = u->vec;
  }

  switch (solver) {
  case BiCGStab:
    iter = oem_bicgstab(oem, dim, fvec, uvec);
    break;
  case CG:
    iter = oem_cg(oem, dim, fvec, uvec);
    break;
  case TfQMR:
    iter = oem_tfqmr(oem, dim, fvec, uvec);
    break;
  case GMRes:
    restart = oem->restart;
    oem->restart = MAX(0, MIN(oem->restart, dim));
    iter = oem_gmres(oem, dim, fvec, uvec);
    oem->restart = restart;
    break;
  case GMRes_k:
    restart = oem->restart;
    oem->restart = MAX(0, MIN(oem->restart, dim));
    iter = oem_gmres_k(oem, dim, fvec, uvec);
    oem->restart = restart;
    break;
  case ODir:
    iter = oem_odir(oem, dim, fvec, uvec);
    break;
  case ORes:
    iter = oem_ores(oem, dim, fvec, uvec);
    break;
  case SymmLQ:
    iter = oem_symmlq(oem, dim, fvec, uvec);
    break;
  default:
    iter = -1; /* make the compiler happy ... -> no-return attribute? */
    ERROR_EXIT("unknown OEM solver %d\n", (int) solver);
  }
  
  if (!CHAIN_SINGLE(u)) {
    copy_to_dof_real_vec(u, uvec);
    MEM_FREE(uvec, dim, REAL);
    MEM_FREE(fvec, dim, REAL);
  }

  return iter;
}

int oem_solve_s(const DOF_MATRIX *A, const DOF_SCHAR_VEC *mask,
		const DOF_REAL_VEC *f,  DOF_REAL_VEC *u,
		OEM_SOLVER solver, REAL tol, const PRECON *precon,
		int restart, int max_iter, int info)
{
  const OEM_DATA *oem;
  int iter;
  
  oem = init_oem_solve(A, mask, tol, precon, restart, max_iter, info);
  iter = call_oem_solve_s(oem, solver, f, u);
  release_oem_solve(oem);

  return iter;
}

/* DIM_OF_WORLD interface */

int call_oem_solve_dow(const OEM_DATA *_oem, OEM_SOLVER solver,
		       const DOF_REAL_VEC_D *f, DOF_REAL_VEC_D *u)
{
  int iter, restart, dim;
  OEM_DATA *oem = (OEM_DATA *)_oem;
  REAL *uvec, *fvec;

  TEST_EXIT(FE_SPACE_EQ_P(f->fe_space, u->fe_space),
	    "Row and column FE_SPACEs don't match!\n");

  dim = dof_real_vec_d_length(f->fe_space);

  if (!CHAIN_SINGLE(u)) {
    uvec = MEM_ALLOC(dim, REAL);
    fvec = MEM_ALLOC(dim, REAL);
    copy_from_dof_real_vec_d(uvec, u);
    copy_from_dof_real_vec_d(fvec, f);
  } else {
    const DOF_ADMIN *admin = u->fe_space->admin;

    fvec = f->vec;
    uvec = u->vec;
    
    FOR_ALL_FREE_DOFS(u->fe_space->admin,
		      if (dof >= admin->size_used) {
			break;
		      }
		      if (u->stride == 1) {
			u->vec[dof] = f->vec[dof] = 0.0;
		      } else {
			SET_DOW(0.0, ((REAL_D *)u->vec)[dof]);
			SET_DOW(0.0, ((REAL_D *)f->vec)[dof]);
		      });
  }

  switch (solver) {
  case BiCGStab:
    iter = oem_bicgstab(oem, dim, fvec, uvec);
    break;
  case CG:
    iter = oem_cg(oem, dim, fvec, uvec);
    break;
  case TfQMR:
    iter = oem_tfqmr(oem, dim, fvec, uvec);
    break;
  case GMRes:
    restart = oem->restart;
    oem->restart = MAX(0, MIN(oem->restart, dim));
    iter = oem_gmres(oem, dim, (REAL *) f->vec, (REAL *) u->vec);
    oem->restart = restart;
    break;
  case GMRes_k:
    restart = oem->restart;
    oem->restart = MAX(0, MIN(oem->restart, dim));
    iter = oem_gmres_k(oem, dim, fvec, uvec);
    oem->restart = restart;
    break;
  case ODir:
    iter = oem_odir(oem, dim, fvec, uvec);
    break;
  case ORes:
    iter = oem_ores(oem, dim, fvec, uvec);
    break;
  case SymmLQ:
    iter = oem_symmlq(oem, dim, fvec, uvec);
    break;
  default:
    iter = -1; /* make the compiler happy ... -> no-return attribute? */
    ERROR_EXIT("unknown OEM solver %d\n", (int) solver);
  }

  if (!CHAIN_SINGLE(u)) {
    copy_to_dof_real_vec_d(u, uvec);
    MEM_FREE(uvec, dim, REAL);
    MEM_FREE(fvec, dim, REAL);
  }

  return iter;
}

int oem_solve_dow(const DOF_MATRIX *A, const DOF_SCHAR_VEC *mask,
		  const DOF_REAL_VEC_D *f, DOF_REAL_VEC_D *u,
		  OEM_SOLVER solver, REAL tol, const PRECON *precon,
		  int restart, int max_iter, int info)
{
  const OEM_DATA *oem;
  int iter;
  
  oem = init_oem_solve(A, mask, tol, precon, restart, max_iter, info);
  iter = call_oem_solve_dow(oem, solver, f, u);
  release_oem_solve(oem);

  return iter;
}
