/*
 *      authors: Claus-Justus Heine
 *               Abteilung fuer Angewandte Mathematik                     
 *               Albert-Ludwigs-Universit"at Freiburg
 *               Hermann-Herder-Str. 10
 *               79104 Freiburg
 *               Germany
 *               claus@mathematik.uni-freiburg.de
 *
 *   Copyright (c) by C.-J. Heine (2005-2007)
 *
 * Generic support for assembling linear systems for time dependent problems.
 */

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

/** Theta-splitting scheme, given the element mass matrix and the part
 * of the stiffness matrix which is subject to the \theta-splitting.
 *
 * Parameters:
 *
 * @param[in]     el_info Element descriptor, filled as specified by
 *                    ::EL_MATRIX_INFO::fill_flag.
 * @param[in,out] dof_matrix The (global) system matrix.
 * @param[in,out] f_h The righ-hand-side vector.
 * @param[in]     u_h The interpolation of the old solution onto the new
 *                    mesh.
 * @param[in]     tau Time-step size.
 * @param[in]     mass_info The fill-info for the element mass matrix.
 * @param[in]     theta The splitting coefficient for the elliptic part.
 * @param[in]     stiff_info The fill-info for the element stiffness matrix,
 *                      the elliptic part of the PDE.
 *
 */

typedef struct system_fill_info_instat
{
  EL_SYS_INFO_INSTAT   elsii;
  const DOF_REAL_VEC   *u_h;
  const EL_MATRIX_INFO *stiff_info;
  const EL_MATRIX_INFO *mass_info;
  int                  n_row;
  int                  n_col;
  int                  n_row_max;
  int                  n_col_max;
} SYSTEM_FILL_INFO_INSTAT;

typedef struct system_fill_info_dow_instat
{
  EL_SYS_INFO_DOW_INSTAT elsii;
  const DOF_REAL_VEC_D   *u_h_d;
  const EL_MATRIX_INFO   *stiff_info;
  const EL_MATRIX_INFO   *mass_info;
  int                    n_row;
  int                    n_col;
  int                    n_row_max;
  int                    n_col_max;
} SYSTEM_FILL_INFO_DOW_INSTAT;

static INIT_EL_TAG element_system_instat(const EL_INFO *el_info,
					 REAL tau,
					 REAL theta,
					 EL_SYS_INFO_INSTAT *thisptr)
{
  SYSTEM_FILL_INFO_INSTAT *info = (SYSTEM_FILL_INFO_INSTAT *)thisptr;
  REAL           theta1 = 1.0 - theta;
  REAL           tau1 = 1.0/tau;
  int            n_row, n_col;

  /* First compute the element matrices, el_matrix_fct() will call all
   * necessary INIT_ELEMENT() methods for us.
   */
  thisptr->el_stiff =
    info->stiff_info->el_matrix_fct(el_info, info->stiff_info->fill_info);
  if (thisptr->el_stiff == NULL) {
    return INIT_EL_TAG_NULL;
  }
  thisptr->el_mass =
    info->mass_info->el_matrix_fct(el_info, info->mass_info->fill_info);
  if (thisptr->el_mass == NULL) {
    return INIT_EL_TAG_NULL;
  }

  thisptr->el_matrix->n_row =
    n_row = thisptr->row_fe_space->bas_fcts->n_bas_fcts;
  thisptr->el_matrix->n_col =
    n_col = thisptr->col_fe_space->bas_fcts->n_bas_fcts;

  info->n_row = n_row;
  info->n_col = n_col;
  
  fill_el_real_vec((EL_REAL_VEC *)thisptr->u_h_loc, el_info->el, info->u_h);

  /*  f += -(1-theta)*a(u_old,psi_i) + 1/tau*m(u_old,psi_i) */
  if (theta1) {
    el_bi_mat_vec(-theta1, thisptr->el_stiff, tau1, thisptr->el_mass,
		  thisptr->u_h_loc, 1.0, thisptr->el_load);
  } else {
    el_mat_vec(tau1, thisptr->el_mass, thisptr->u_h_loc, thisptr->el_load);
  }

  /* initialize with 1/tau*m(psi_i,psi_j) */
  el_mat_axey(tau1, thisptr->el_mass, thisptr->el_matrix);

  /*  add theta*a(psi_i,psi_j) */
  if (theta) {
    el_mat_axpy(theta, thisptr->el_stiff, thisptr->el_matrix);
  }
  return INIT_EL_TAG_DFLT;
}

static INIT_EL_TAG element_system_dow_instat(const EL_INFO *el_info,
					     REAL tau,
					     REAL theta,
					     EL_SYS_INFO_DOW_INSTAT *thisptr)
{
  SYSTEM_FILL_INFO_DOW_INSTAT *info = (SYSTEM_FILL_INFO_DOW_INSTAT *)thisptr;
  REAL theta1 = 1.0 - theta;
  REAL tau1 = 1.0/tau;
  int  n_row, n_col;
  
  /* First compute the element matrices, el_matrix_fct() will call all
   * necessary INIT_ELEMENT() methods for us.
   */
  thisptr->el_stiff = info->stiff_info->el_matrix_fct(el_info,
						   info->stiff_info->fill_info);
  if (thisptr->el_stiff == NULL) {
    return INIT_EL_TAG_NULL;
  }

  thisptr->el_mass =
    info->mass_info->el_matrix_fct(el_info, info->mass_info->fill_info);
  if (thisptr->el_mass == NULL) {
    return INIT_EL_TAG_NULL;
  }

  thisptr->el_matrix->n_row =
    n_row = thisptr->row_fe_space->bas_fcts->n_bas_fcts;
  thisptr->el_matrix->n_col =
    n_col = thisptr->col_fe_space->bas_fcts->n_bas_fcts;

  info->n_row = n_row;
  info->n_col = n_col;
  
  fill_el_real_vec_d(
    (EL_REAL_VEC_D *)thisptr->u_h_loc, el_info->el, info->u_h_d);
  
  /*  f += -(1-theta)*a(u_old,psi_i) + 1/tau*m(u_old,psi_i) */
  if (theta1) {
    el_bi_mat_vec_dow(-theta1, thisptr->el_stiff, tau1, thisptr->el_mass,
		      thisptr->u_h_loc, 1.0, thisptr->el_load);
  } else {
    el_mat_vec_dow(tau1, thisptr->el_mass, thisptr->u_h_loc, thisptr->el_load);
  }

  /* initialize with 1/tau*m(psi_i,psi_j) */
  el_mat_axey(tau1, thisptr->el_mass, thisptr->el_matrix);

  /*  add theta*a(psi_i,psi_j) */
  if (theta) {
    el_mat_axpy(theta, thisptr->el_stiff, thisptr->el_matrix);
  }
  return INIT_EL_TAG_DFLT;
}

EL_SYS_INFO_INSTAT *fill_sys_info_instat(const OPERATOR_INFO *stiff_info,
					 const OPERATOR_INFO *mass_info,
					 const DOF_REAL_VEC *u_h)
{
  SYSTEM_FILL_INFO_INSTAT *info;

  info = MEM_CALLOC(1, SYSTEM_FILL_INFO_INSTAT);
 
  info->stiff_info = fill_matrix_info(stiff_info, NULL);
  info->mass_info  = fill_matrix_info(mass_info, NULL);

  info->elsii.row_fe_space = info->mass_info->row_fe_space;
  info->elsii.col_fe_space = info->mass_info->col_fe_space;
  if (info->elsii.col_fe_space == NULL) {
    info->elsii.col_fe_space = info->elsii.row_fe_space;
  }

  info->elsii.el_update_fct = element_system_instat;

  info->n_row     = info->elsii.row_fe_space->bas_fcts->n_bas_fcts;
  info->n_row_max = info->elsii.row_fe_space->bas_fcts->n_bas_fcts_max;
  info->n_col     = info->elsii.col_fe_space->bas_fcts->n_bas_fcts;
  info->n_col_max = info->elsii.col_fe_space->bas_fcts->n_bas_fcts_max;

  info->elsii.el_matrix = get_el_matrix(info->elsii.row_fe_space,
					info->elsii.col_fe_space,
					MATENT_REAL);
  info->elsii.el_load = get_el_real_vec(info->elsii.row_fe_space->bas_fcts);
  info->elsii.u_h_loc = get_el_real_vec(info->elsii.col_fe_space->bas_fcts);
  
  info->elsii.fill_flag =
    info->stiff_info->fill_flag | info->mass_info->fill_flag;
  BNDRY_FLAGS_CPY(info->elsii.dirichlet_bndry,
		  info->mass_info->dirichlet_bndry);
  BNDRY_FLAGS_OR(info->elsii.dirichlet_bndry,
		 info->stiff_info->dirichlet_bndry);  
  if (!BNDRY_FLAGS_IS_INTERIOR(info->elsii.dirichlet_bndry)) {
    info->elsii.fill_flag |= FILL_BOUND;
    if (info->elsii.row_fe_space->mesh->is_periodic &&
	!(info->elsii.row_fe_space->admin->flags & ADM_PERIODIC))
      info->elsii.fill_flag |= FILL_NON_PERIODIC;
  }

  info->u_h = u_h;
  
  return &info->elsii;
}

EL_SYS_INFO_DOW_INSTAT *
fill_sys_info_instat_dow(const OPERATOR_INFO *stiff_info,
			 const OPERATOR_INFO *mass_info,
			 const DOF_REAL_VEC_D *u_h_d)
{
  SYSTEM_FILL_INFO_DOW_INSTAT *info;
  
  info = MEM_CALLOC(1, SYSTEM_FILL_INFO_DOW_INSTAT);
 
  info->stiff_info = fill_matrix_info(stiff_info, NULL);
  info->mass_info  = fill_matrix_info(mass_info, NULL);

  info->elsii.krn_blk_type = info->stiff_info->krn_blk_type;

  info->elsii.row_fe_space = info->mass_info->row_fe_space;
  info->elsii.col_fe_space = info->mass_info->col_fe_space;
  if (info->elsii.col_fe_space == NULL) {
    info->elsii.col_fe_space = info->elsii.row_fe_space;
  }

  info->elsii.el_update_fct = element_system_dow_instat;

  info->n_row     = info->elsii.row_fe_space->bas_fcts->n_bas_fcts;
  info->n_row_max = info->elsii.row_fe_space->bas_fcts->n_bas_fcts_max;
  info->n_col     = info->elsii.col_fe_space->bas_fcts->n_bas_fcts;
  info->n_col_max = info->elsii.col_fe_space->bas_fcts->n_bas_fcts_max;

  info->elsii.el_matrix = get_el_matrix(info->elsii.row_fe_space,
					info->elsii.col_fe_space,
					info->elsii.krn_blk_type);
  info->elsii.el_load = get_el_real_vec_d(info->elsii.row_fe_space->bas_fcts);
  info->elsii.u_h_loc = get_el_real_vec_d(info->elsii.col_fe_space->bas_fcts);

  info->elsii.fill_flag =
    info->stiff_info->fill_flag | info->mass_info->fill_flag;
  BNDRY_FLAGS_CPY(info->elsii.dirichlet_bndry,
		  info->mass_info->dirichlet_bndry);
  BNDRY_FLAGS_OR(info->elsii.dirichlet_bndry,
		 info->stiff_info->dirichlet_bndry);  
  if (!BNDRY_FLAGS_IS_INTERIOR(info->elsii.dirichlet_bndry)) {
    info->elsii.fill_flag |= FILL_BOUND;
    if (info->elsii.row_fe_space->mesh->is_periodic &&
	!(info->elsii.row_fe_space->admin->flags & ADM_PERIODIC))
      info->elsii.fill_flag |= FILL_NON_PERIODIC;
  }

  info->u_h_d = u_h_d;
  
  return &info->elsii;
}

void free_sys_info_instat(EL_SYS_INFO_INSTAT *elsii)
{
  SYSTEM_FILL_INFO_INSTAT *info = (SYSTEM_FILL_INFO_INSTAT *)elsii;

  free_el_matrix(elsii->el_matrix);
  free_el_real_vec(elsii->el_load);
  free_el_real_vec((EL_REAL_VEC *)elsii->u_h_loc);
  MEM_FREE(info->stiff_info, 1, EL_MATRIX_INFO);
  MEM_FREE(info->mass_info, 1, EL_MATRIX_INFO);
  MEM_FREE(info, 1, SYSTEM_FILL_INFO_INSTAT);
}

void free_sys_info_dow_instat(EL_SYS_INFO_DOW_INSTAT *elsii)
{
  SYSTEM_FILL_INFO_DOW_INSTAT *info = (SYSTEM_FILL_INFO_DOW_INSTAT *)elsii;

  free_el_matrix(elsii->el_matrix);
  free_el_real_vec_d(elsii->el_load);
  free_el_real_vec_d((EL_REAL_VEC_D *)elsii->u_h_loc);
  MEM_FREE(info->stiff_info, 1, EL_MATRIX_INFO);
  MEM_FREE(info->mass_info, 1, EL_MATRIX_INFO);
  MEM_FREE(info, 1, SYSTEM_FILL_INFO_DOW_INSTAT);
}

/* Assemble (1/tau M + theta S) and add (1/tau M uh + (1-theta) S uh)
 * to fh.
 */
void update_system_instat(DOF_MATRIX *dof_matrix,
			  DOF_REAL_VEC *f_h,
			  REAL tau,
			  REAL theta,
			  EL_SYS_INFO_INSTAT *elsii)			  
{
  const BAS_FCTS *row_bfcts = elsii->row_fe_space->bas_fcts;
  const EL_DOF_VEC *row_dof, *col_dof;
  EL_SCHAR_VEC       *bound = NULL;
  const EL_BNDRY_VEC *bndry_bits;
  bool           use_get_bound;

  BNDRY_FLAGS_CPY(dof_matrix->dirichlet_bndry, elsii->dirichlet_bndry);

  use_get_bound = !BNDRY_FLAGS_IS_INTERIOR(dof_matrix->dirichlet_bndry);

  if (use_get_bound) {
    bound = get_el_schar_vec(elsii->row_fe_space->bas_fcts);
  }

  TRAVERSE_FIRST(dof_matrix->row_fe_space->mesh,
		 -1,
		 elsii->fill_flag|CALL_LEAF_EL|FILL_COORDS) {

    if (elsii->el_update_fct(el_info, tau, theta, elsii) == INIT_EL_TAG_NULL) {
      continue;
    }

    row_dof = get_dof_indices(NULL, f_h->fe_space, el_info->el);
    if (elsii->row_fe_space != elsii->col_fe_space) {
      col_dof =
	get_dof_indices(NULL, elsii->col_fe_space, el_info->el);
    } else {
      col_dof = row_dof;
    }

    if (use_get_bound) {
      bndry_bits = get_bound(NULL, row_bfcts, el_info);
      dirichlet_map(bound, bndry_bits, dof_matrix->dirichlet_bndry);
    }

    add_element_matrix(dof_matrix, 1.0, elsii->el_matrix, NoTranspose,
		       row_dof, col_dof, bound);
    
    add_element_vec(f_h, 1.0, elsii->el_load, row_dof, bound);

  } TRAVERSE_NEXT();

  if (use_get_bound) {
    free_el_schar_vec(bound);
  }
}

void update_system_instat_dow(DOF_MATRIX *dof_matrix,
			      DOF_REAL_VEC_D *f_h,
			      REAL tau,
			      REAL theta,
			      EL_SYS_INFO_DOW_INSTAT *elsii)
{
  const BAS_FCTS *row_bfcts = elsii->row_fe_space->bas_fcts;
  const EL_DOF_VEC *row_dof, *col_dof;
  EL_SCHAR_VEC       *bound = NULL;
  const EL_BNDRY_VEC *bndry_bits;
  bool           use_get_bound;

  BNDRY_FLAGS_CPY(dof_matrix->dirichlet_bndry, elsii->dirichlet_bndry);

  use_get_bound = !BNDRY_FLAGS_IS_INTERIOR(dof_matrix->dirichlet_bndry);

  if (use_get_bound) {
    bound = get_el_schar_vec(elsii->row_fe_space->bas_fcts);
  }

  TRAVERSE_FIRST(dof_matrix->row_fe_space->mesh,
		 -1,
		 elsii->fill_flag|CALL_LEAF_EL|FILL_COORDS) {

    if (elsii->el_update_fct(el_info, tau, theta, elsii) == INIT_EL_TAG_NULL) {
      continue;
    }

    row_dof = get_dof_indices(NULL, f_h->fe_space, el_info->el);
    if (elsii->row_fe_space != elsii->col_fe_space) {
      col_dof =
	get_dof_indices(NULL, elsii->col_fe_space, el_info->el);
    } else {
      col_dof = row_dof;
    }

    if (use_get_bound) {
      bndry_bits = get_bound(NULL, row_bfcts, el_info);
      dirichlet_map(bound, bndry_bits, dof_matrix->dirichlet_bndry);
    }

    add_element_matrix(dof_matrix, 1.0, elsii->el_matrix, NoTranspose,
		       row_dof, col_dof, bound);
    
    add_element_vec_dow(f_h, 1.0, elsii->el_load, row_dof, bound);

  } TRAVERSE_NEXT();

  if (use_get_bound) {
    free_el_schar_vec(bound);
  }
}