/*--------------------------------------------------------------------------*/
/* ALBERTA:  an Adaptive multi Level finite element toolbox using           */
/*           Bisectioning refinement and Error control by Residual          */
/*           Techniques for scientific Applications                         */
/*                                                                          */
/* file:     MG_s1.c                                                        */
/*                                                                          */
/* description:  multigrid method for a scalar elliptic equation            */
/*                                                                          */
/*--------------------------------------------------------------------------*/
/*                                                                          */
/*  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                                    */
/*                                                                          */
/*  http://www.mathematik.uni-freiburg.de/IAM/ALBERTA                       */
/*                                                                          */
/*  (c) by A. Schmidt and K.G. Siebert (1996-2003)                          */
/*                                                                          */
/*--------------------------------------------------------------------------*/

#include "alberta.h"
#include <time.h>
#include <string.h>

/****************************************************************************/
typedef struct mg_traverse_data {
  int v_n0;
  int max_level, max_dof_level;
  U_CHAR *dof_level;
  DOF *dof_parent0, *dof_parent1;
} MG_TRAVERSE_DATA;

static void max_level_fct(const EL_INFO *el_info, void *data)
{
  MG_TRAVERSE_DATA *ud = (MG_TRAVERSE_DATA *)data;
  int dof, dof0, dof1, level = (int)(el_info->level);
  int dim = el_info->mesh->dim;
  EL  *el = el_info->el;
  int n0  = ud->v_n0;

  ud->max_level = MAX(level, ud->max_level);

  if (!(IS_LEAF_EL(el))) {
    dof   = el->child[0]->dof[N_VERTICES(dim)-1][n0+0];
    dof0  = ud->dof_parent0[dof] = el->dof[0][n0+0];
    dof1  = ud->dof_parent1[dof] = el->dof[1][n0+0];
    level = ud->dof_level[dof]   = 1 + MAX(ud->dof_level[dof0],
					   ud->dof_level[dof1]);
    ud->max_dof_level = MAX(level, ud->max_dof_level);
  }
}

/****************************************************************************/
/* int MG_s_setup_levels(MG_S_INFO *mg_s_info):                             */
/*  build  levels, dofs_per_level[], dof_level[], dof_parent[][],           */
/*  sort_dof[ sorted dof ]          = unsorted dof,                         */
/*  sort_dof_invers[ unsorted dof ] = sorted dof                            */
/* following info must be present in *mg_s_info:                            */
/*  mg_info, fe_space                                                       */
/****************************************************************************/
void MG_s_setup_levels(MG_S_INFO *mg_s_info)
{
  FUNCNAME("MG_s_setup_levels");
  static MG_TRAVERSE_DATA td[1];
  int     levels, *tmp_per_level, *dofs_per_level;
  DOF     *sort_dof, *sort_dof_invers, *parent;
  int     i, j, k, sort_size, dim;
  char    name[128];
  const DOF_ADMIN  *admin = NULL;
  MULTI_GRID_INFO  *mg_info = NULL;
  DOF_MATRIX       *mat;

  TEST_EXIT(mg_s_info && (mg_info = mg_s_info->mg_info),
    "no mg_s_info or mg_info\n");
  TEST_EXIT(mg_s_info->fe_space && (admin = mg_s_info->fe_space->admin),
    "no fe_space or admin\n");
  dim = mg_s_info->fe_space->mesh->dim;

  TEST_EXIT(mg_s_info->fe_space->bas_fcts,"no bas_fcts\n");
  if (!strstr(mg_s_info->fe_space->bas_fcts->name, "lagrange1")) {
    ERROR_EXIT("not for bas_fcts <%s>, only for <lagrange1>\n",
	       mg_s_info->fe_space->bas_fcts->name);
  }

  sort_size = admin->used_count;
  if (mg_s_info->sort_size < sort_size) {
    mg_s_info->sort_dof = MEM_REALLOC(mg_s_info->sort_dof, 
				      mg_s_info->sort_size, sort_size, int);
    mg_s_info->dof_parent[0] = MEM_REALLOC(mg_s_info->dof_parent[0], 
				      mg_s_info->sort_size, sort_size, DOF);
    mg_s_info->dof_parent[1] = MEM_REALLOC(mg_s_info->dof_parent[1], 
				      mg_s_info->sort_size, sort_size, DOF);
    mg_s_info->dof_level = MEM_REALLOC(mg_s_info->dof_level, 
				      mg_s_info->sort_size, sort_size, U_CHAR);
    mg_s_info->sort_bound = MEM_REALLOC(mg_s_info->sort_bound, 
				      mg_s_info->sort_size, sort_size, S_CHAR);
    mg_s_info->sort_size = sort_size;
  }

  if (mg_s_info->sort_invers_size < admin->size_used) {
    mg_s_info->sort_dof_invers = MEM_REALLOC(mg_s_info->sort_dof_invers,
					     mg_s_info->sort_invers_size,
					     admin->size_used, int);
    mg_s_info->sort_invers_size = admin->size_used;
  }

  td->v_n0        = mg_s_info->vertex_admin->n0_dof[VERTEX];
  td->dof_level   = mg_s_info->dof_level;
  td->dof_parent0 = mg_s_info->dof_parent[0];
  td->dof_parent1 = mg_s_info->dof_parent[1];
  for (i=0; i<sort_size; i++) {
    td->dof_level[i] = 0;
    td->dof_parent0[i] = td->dof_parent1[i] = 0;
  }
  td->max_level = td->max_dof_level = 0;
  mesh_traverse(mg_s_info->fe_space->mesh, -1, CALL_EVERY_EL_PREORDER,
		max_level_fct, td);
  levels = (td->max_level + dim - 1) / dim + 1;
  TEST_EXIT(levels == (td->max_dof_level+1),
    "levels %d != max_dof_level %d + 1\n", levels, td->max_dof_level);

  if (levels > mg_s_info->size)
  {
    mg_s_info->matrix = MEM_REALLOC(mg_s_info->matrix,
				    mg_s_info->size, levels, DOF_MATRIX *);
    mg_s_info->f_h = MEM_REALLOC(mg_s_info->f_h,
				 mg_s_info->size, levels, REAL *);
    mg_s_info->u_h = MEM_REALLOC(mg_s_info->u_h,
				 mg_s_info->size, levels, REAL *);
    mg_s_info->r_h = MEM_REALLOC(mg_s_info->r_h,
				 mg_s_info->size, levels, REAL *);
    mg_s_info->dofs_per_level = MEM_REALLOC(mg_s_info->dofs_per_level,
					    mg_s_info->size, levels, int);

    for (i = mg_s_info->size; i < levels; i++)
    {
      sprintf(name, "mg matrix level %d", i);
      mg_s_info->matrix[i] = get_dof_matrix(name, NULL, NULL);
      mg_s_info->f_h[i] = NULL;
      mg_s_info->u_h[i] = NULL;
      mg_s_info->r_h[i] = NULL;
    }

    mg_s_info->size = levels;
  }

  mg_info->mg_levels = levels;


/****************************************************************************/
/* count dofs per multigrid level                                           */

  sort_dof        = mg_s_info->sort_dof;
  sort_dof_invers = mg_s_info->sort_dof_invers;
  dofs_per_level  = mg_s_info->dofs_per_level;
  tmp_per_level   = MEM_ALLOC(levels, int);

  for (i = 0; i < levels; i++) dofs_per_level[i] = 0;

  FOR_ALL_DOFS(admin, dofs_per_level[td->dof_level[dof]]++);

  if (mg_s_info->mg_info->info > 2) {
    MSG("dofs_per_level:");
    for (i = 0; i < levels; i++) print_msg(" %d", dofs_per_level[i]);
    print_msg("\n");
  }

  for (i = 1; i < levels; i++) {
    tmp_per_level[i]   = dofs_per_level[i-1];
    dofs_per_level[i] += dofs_per_level[i-1];
  }
  tmp_per_level[0] = 0;            /* pointers for filling the sort vectors */


  if (mg_s_info->mg_info->info > 2) {
    MSG("dofs_per_level accumulated:");
    for (i = 0; i < levels; i++) print_msg(" %d", dofs_per_level[i]);
    print_msg("\n");
  }

  if (mg_s_info->mg_info->info > 9) {
    for (i = 0; i < dofs_per_level[levels-1]; i++)
      MSG("dof_parent[%3d] = (%3d,%3d), lev=%2d (%2d,%2d)\n",
	  i, td->dof_parent0[i], td->dof_parent1[i],
	  td->dof_level[i], td->dof_level[td->dof_parent0[i]],
	  td->dof_level[td->dof_parent1[i]]);
  }

  DEBUG_TEST(dofs_per_level[levels-1] == sort_size,
    "dofs_per_level[levels-1] = %d != %d = sort_size\n",
     dofs_per_level[levels-1], sort_size);

/****************************************************************************/
/* build sort_dof[] and sort_dof_invers[] vectors                           */

  FOR_ALL_DOFS(admin,
	       j = td->dof_level[dof];
	       k = tmp_per_level[j]++;
	       sort_dof[k] = dof;
	       sort_dof_invers[dof] = k;
    );

  if (mg_s_info->mg_info->info > 9) {
    for (i = 0; i < dofs_per_level[levels-1]; i++)
    {
      j = sort_dof[i];
      MSG("sort[%3d]: dof=%3d, lev=%2d; invers[%3d]=%3d\n",
	  i, j, td->dof_level[j], j, sort_dof_invers[j]);
    }
  }

  MEM_FREE(tmp_per_level, levels, int);

/****************************************************************************/
/* enlarge vectors and matrices, if necessary                               */
/****************************************************************************/

  for (i = 0; i < levels; i++)
  {
    mat = mg_s_info->matrix[i];
    mat->matrix_row = MEM_ALLOC(dofs_per_level[i], MATRIX_ROW *);
    for (j = 0; j < dofs_per_level[i]; j++)
      mat->matrix_row[j] = NULL;
    mat->size = dofs_per_level[i];

    mg_s_info->f_h[i] = MEM_ALLOC(dofs_per_level[i], REAL);
    mg_s_info->u_h[i] = MEM_ALLOC(dofs_per_level[i], REAL);
    mg_s_info->r_h[i] = MEM_ALLOC(dofs_per_level[i], REAL);
  }

/****************************************************************************/
/* transform dof_parent vectors to sorted dofs                              */
/****************************************************************************/
  parent = MEM_ALLOC(sort_size, DOF);

  for (i=0; i<sort_size; i++)
    parent[i] = sort_dof_invers[mg_s_info->dof_parent[0][sort_dof[i]]];
  for (i=0; i<sort_size; i++)
    mg_s_info->dof_parent[0][i] = parent[i];

  for (i=0; i<sort_size; i++)
    parent[i] = sort_dof_invers[mg_s_info->dof_parent[1][sort_dof[i]]];
  for (i=0; i<sort_size; i++)
    mg_s_info->dof_parent[1][i] = parent[i];

  MEM_FREE(parent, sort_size, DOF);

  return;
}

/****************************************************************************/
/* MG_s_setup_mat_b():                                                      */
/*  sort+restrict matrix, sort bound                                        */
/*  parameter bound may be NULL.                                             */
/****************************************************************************/
void MG_s_setup_mat_b(MG_S_INFO *mg_s_info,
		      DOF_MATRIX *mat, const DOF_SCHAR_VEC *bound)
{
  FUNCNAME("MG_s_setup_mat_b");
  DOF             *sort_dof;
  int             *sort_dof_invers;
  S_CHAR          *b, *sb = NULL;
  int             size, level, i;
  const DOF_ADMIN *admin;
  DOF_MATRIX      *sort_mat;
  clock_t         first = 0, second;

  TEST_EXIT(mg_s_info && mg_s_info->fe_space,
    "no mg_s_info or fe_space\n");
  TEST_EXIT(admin=mg_s_info->fe_space->admin,
    "no admin\n");
  TEST_EXIT(sort_dof_invers=mg_s_info->sort_dof_invers,
    "no sort_dof_invers\n");
  TEST_EXIT(sort_dof=mg_s_info->sort_dof,
    "no sort_dof\n");

  TEST_EXIT(mat && mat->matrix_row,
    "no mat or matrix_row\n");

  mg_s_info->mat = mat;
  mg_s_info->bound = bound;



  if (mg_s_info->mg_info->info > 2) first = clock();

/****************************************************************************/
/* transform matrix to sorted dofs                                          */
/* change column indices in place!!!!! (i.e. matrix is changed!!!)          */
/****************************************************************************/

  level    = mg_s_info->mg_info->mg_levels - 1;
  size     = mg_s_info->dofs_per_level[level];
  sort_mat = mg_s_info->matrix[level];

  if (mg_s_info->mg_info->info > 7) {
    MSG("unsorted matrix:\n");
    print_dof_matrix(mat);
  }

  clear_dof_matrix(sort_mat);

  FOR_ALL_DOFS(
    admin,
    FOR_ALL_MAT_COLS(
      REAL,
      sort_mat->matrix_row[sort_dof_invers[dof]] = mat->matrix_row[dof],
      row->col[col_idx] = sort_dof_invers[col_dof]));

  for (i=0; i<size; i++) {
    DEBUG_TEST(sort_mat->matrix_row[i], "NULL sort_mat->matrix_row[%d]\n", i);
  }

  if (mg_s_info->mg_info->info > 2)
  {
    second = clock();
    MSG("sort needed %.5lf seconds\n",
	(double)(second-first)/(double)CLOCKS_PER_SEC);
  }


/****************************************************************************/
/* transform bound vector to sorted dofs                                    */
/****************************************************************************/

  TEST_EXIT(sb = mg_s_info->sort_bound,"no sort_bound\n");
  if (mg_s_info->bound) {
    TEST_EXIT(b  = mg_s_info->bound->vec,"no bound->vec\n");

    for (i=0; i<size; i++)
      sb[i] = b[sort_dof[i]];
  }
  else {
    for (i=0; i<size; i++)
      sb[i] = INTERIOR;
  }

/****************************************************************************/
/* restrict matrices to coarser levels                                      */
/****************************************************************************/

  MG_s_restrict_mg_matrices(mg_s_info);

  if (mg_s_info->mg_info->info > 7) {
    MSG("sorted matrix:\n");
    print_dof_matrix(sort_mat);
  }


  if (mg_s_info->mg_info->info > 2)
  {
    second = clock();
    MSG("matrix+bound sort+restrict needed %.5lf seconds\n",
	(double)(second-first)/(double)CLOCKS_PER_SEC);
  }

  return;
}


/****************************************************************************/
/*  MG_s_dof_copy_to_sparse(DOF_ADMIN *admin, MULTI_GRID_INFO *mg_info,     */
/*       	            DOF_REAL_VEC *x, REAL *y)                       */
/*  MG_s_dof_copy_from_sparse(DOF_ADMIN *admin, MULTI_GRID_INFO *mg_info,   */
/*			      REAL *x, DOF_REAL_VEC *y)                     */
/****************************************************************************/
void MG_s_dof_copy_to_sparse(MG_S_INFO *mg_s_info,
			     const DOF_REAL_VEC *x, REAL *y)
{
  FUNCNAME("MG_s_dof_copy_to_sparse");
  int             j, jmax, *sort_dof_invers;
  REAL            *xvec;
  const DOF_ADMIN *admin = NULL;

  TEST_EXIT(mg_s_info && y,
    "no mg_s_info or y\n");
  TEST_EXIT(x && x->fe_space && (admin=x->fe_space->admin),
    "no x or x->fe_space or x->fe_space->admin\n");

  xvec = x->vec;
  jmax = mg_s_info->dofs_per_level[mg_s_info->mg_info->mg_levels-1];
  sort_dof_invers = mg_s_info->sort_dof_invers;

  FOR_ALL_DOFS(admin,
	       j = sort_dof_invers[dof];
	       TEST_EXIT(j < jmax,"j=%d too big; dof=%d, max+1=%d\n",
				   j, dof, jmax);
	       y[j] = xvec[dof];
    );
}


void MG_s_dof_copy_from_sparse(MG_S_INFO *mg_s_info, 
			       const REAL *x, DOF_REAL_VEC *y)
{
  FUNCNAME("MG_s_dof_copy_from_sparse");
  int             i, j, jmax, *sort_dof, ysize;
  REAL            *yvec;

  TEST_EXIT(mg_s_info && x,"no mg_info or x\n");
  TEST_EXIT(y,"no y\n");

  sort_dof = mg_s_info->sort_dof;
  jmax = mg_s_info->dofs_per_level[mg_s_info->mg_info->mg_levels-1];
  yvec = y->vec;
  ysize = y->size;

  for (j = 0; j < jmax; j++)
  {
    i = sort_dof[j];
    TEST_EXIT(i < ysize,"i=%d too big; j=%d, y->size=%d\n", i, j, ysize);
    yvec[i] = x[j];
  }
}

/****************************************************************************/
/* MG_s_reset_mat():                                                        */
/*  unsort matrix columns                                                   */
/****************************************************************************/
void MG_s_reset_mat(MG_S_INFO *mg_s_info)
{
  FUNCNAME("MG_s_reset_mat");
  DOF             *sort_dof;
  const DOF_ADMIN *admin;
  DOF_MATRIX      *mat;

  TEST_EXIT(mg_s_info && mg_s_info->fe_space,
    "no mg_s_info or fe_space\n");
  TEST_EXIT(admin=mg_s_info->fe_space->admin,
    "no admin\n");
  TEST_EXIT(sort_dof=mg_s_info->sort_dof,
    "no sort_dof\n");

  TEST_EXIT((mat = mg_s_info->mat) && mat->matrix_row,
    "no mat or matrix_row\n");

/****************************************************************************/
/* transform matrix to unsorted dofs                                        */
/* change column indices in place!!!!! (i.e. matrix is changed!!!)          */
/****************************************************************************/

  FOR_ALL_DOFS(admin,
	       FOR_ALL_MAT_COLS(REAL, mat->matrix_row[dof],
				row->col[col_idx] = sort_dof[col_dof]));
}

/****************************************************************************/
/* MG_s_sort_mat():                                                         */
/*  sort matrix columns                                                     */
/****************************************************************************/
void MG_s_sort_mat(MG_S_INFO *mg_s_info)
{
  FUNCNAME("MG_s_sort_mat");
  int             *sort_dof_invers;
  const DOF_ADMIN *admin;
  DOF_MATRIX      *mat;

  TEST_EXIT(mg_s_info && mg_s_info->fe_space,
    "no mg_s_info or fe_space\n");
  TEST_EXIT(admin=mg_s_info->fe_space->admin,
    "no admin\n");
  TEST_EXIT(sort_dof_invers=mg_s_info->sort_dof_invers,
    "no sort_dof_invers\n");

  TEST_EXIT((mat = mg_s_info->mat) && mat->matrix_row,
    "no mat or matrix_row\n");

/****************************************************************************/
/* transform matrix to unsorted dofs                                        */
/* change column indices in place!!!!! (i.e. matrix is changed!!!)          */
/****************************************************************************/

  FOR_ALL_DOFS(admin,
	       FOR_ALL_MAT_COLS(REAL, mat->matrix_row[dof],
				row->col[col_idx] = sort_dof_invers[col_dof]));
}

/****************************************************************************/
/*  mg_info, fe_space                                                       */
/****************************************************************************/
void MG_s_free_mem(MG_S_INFO *mg_s_info)
{
  FUNCNAME("MG_s_free_mem");
  MULTI_GRID_INFO  *mg_info = NULL;
  int levels, i, *dofs_per_level;

  TEST_EXIT(mg_s_info && (mg_info = mg_s_info->mg_info),
    "no mg_s_info or mg_info\n");
  TEST_EXIT(dofs_per_level = mg_s_info->dofs_per_level,
    "no dofs_per_level\n");
  levels = mg_info->mg_levels;

  /* fine grid matrix uses rows from mg_s_info->mat */
  for (i=0; i<mg_s_info->matrix[levels-1]->size; i++)
    mg_s_info->matrix[levels-1]->matrix_row[i] = NULL;

  for (i = levels-1; i >= 0; i--)
  {
    MEM_FREE(mg_s_info->r_h[i], dofs_per_level[i], REAL);
    MEM_FREE(mg_s_info->u_h[i], dofs_per_level[i], REAL);
    MEM_FREE(mg_s_info->f_h[i], dofs_per_level[i], REAL);
    free_dof_matrix(mg_s_info->matrix[i]);
  }

  MEM_FREE(mg_s_info->dofs_per_level, mg_s_info->size, int);
  MEM_FREE(mg_s_info->r_h, mg_s_info->size, REAL *);
  MEM_FREE(mg_s_info->u_h, mg_s_info->size, REAL *);
  MEM_FREE(mg_s_info->f_h, mg_s_info->size, REAL *);
  MEM_FREE(mg_s_info->matrix, mg_s_info->size, DOF_MATRIX *);
  MEM_FREE(mg_s_info->sort_dof_invers, mg_s_info->sort_invers_size, int);
  MEM_FREE(mg_s_info->sort_bound, mg_s_info->sort_size, S_CHAR);
  MEM_FREE(mg_s_info->dof_level, mg_s_info->sort_size, U_CHAR);
  MEM_FREE(mg_s_info->dof_parent[1], mg_s_info->sort_size, DOF);
  MEM_FREE(mg_s_info->dof_parent[0], mg_s_info->sort_size, DOF);
  MEM_FREE(mg_s_info->sort_dof, mg_s_info->sort_size, int);

  mg_s_info->sort_dof_invers = NULL;
  mg_s_info->dofs_per_level = NULL;
  mg_s_info->f_h = mg_s_info->u_h = mg_s_info->r_h = NULL;
  mg_s_info->matrix = NULL;
  mg_s_info->sort_dof_invers = NULL;
  mg_s_info->sort_bound = NULL;
  mg_s_info->dof_level = NULL;
  mg_s_info->dof_parent[0] = mg_s_info->dof_parent[1] = NULL;
  mg_s_info->sort_dof = NULL;

  mg_s_info->size = 0;
  mg_s_info->sort_size = 0;
  mg_s_info->sort_invers_size = 0;

  return;
}
