/*******************************************************************************
 *
 * ALBERTA:  an Adaptive multi Level finite element toolbox  using
 *           Bisectioning refinement and Error control by Residual
 *           Techniques
 *
 * www.alberta-fem.de
 *
 * file:     heat-easy.c
 *
 * description:  solver for parabolic model problem
 *                                                                          
 *                   u,t - \Delta u = f  in \Omega                          
 *                                u = g  on \partial \Omega                 
 *                                                                          
 * This version demonstrates the use of EL_SYS_INFO_INSTAT together with  
 * update_system_instat()                                                  
 *
 ******************************************************************************
 *
 * this file's author(s):
 *            Claus-Justus Heine
 *            Abteilung fuer Angewandte Mathematik
 *            Albert-Ludwigs-Universitaet Freiburg
 *            Hermann-Herder-Str. 10
 *            79104 Freiburg
 *            Germany
 *            Claus.Heine@Mathematik.Uni-Freiburg.DE
 *
 *  (c) by C.-J. Heine (2005-2007)
 *                                                           
 * Based on the heat.c demo-program by K.G. Siebert and A. Schmidt.
 *              
 ******************************************************************************/

#include <alberta/alberta.h>

#include "alberta-demo.h"

static bool do_graphics = true; /* global graphics "kill-switch" */

/*******************************************************************************
 * global variables: finite element space, discrete solution, discrete  /
 *                   solution from previous time step, load vector,     
 *                   system matrix, and structure for adaptive procedure
 *
 * These variables are kept global because they are shared across build(),
 * solve() and estimate().
 ******************************************************************************/
static const FE_SPACE *fe_space;      /* initialized by main() */
static DOF_REAL_VEC   *u_h = NULL;    /* initialized by main() */
static DOF_REAL_VEC   *u_old = NULL;  /* initialized by main() */
static DOF_REAL_VEC   *f_h = NULL;    /* initialized by main() */
static DOF_MATRIX     *matrix = NULL; /* initialized by main() */
static ADAPT_INSTAT   *adapt_instat;  /* initialized by main() */
static EL_SYS_INFO_INSTAT *elsii;     /* initialized by main() */
static const QUAD     *load_quad;     /* initialized by main() */
static BNDRY_FLAGS    dirichlet_mask; /* initialized by main() */

static REAL theta = 0.5;   /* parameter of the time discretization */
static REAL err_L2  = 0.0; /* spatial error in a single time step  */

/*******************************************************************************
 * struct heat_leaf_data: structure for storing one REAL value on each      
 *                          leaf element as LEAF_DATA                       
 * rw_el_est():  return a pointer to the memory for storing the element     
 *               estimate (stored as LEAF_DATA), called by heat_est()       
 * get_el_est(): return the value of the element estimates (from LEAF_DATA),
 *               called by adapt_method_stat() and graphics()               
 ******************************************************************************/

struct heat_leaf_data
{
  REAL estimate;            /*  one real for the estimate */
  REAL est_c;               /*  one real for the coarsening estimate */
};

static REAL *rw_el_est(EL *el)
{
  if (IS_LEAF_EL(el))
    return(&((struct heat_leaf_data *)LEAF_DATA(el))->estimate);
  else
    return(NULL);
}

static REAL get_el_est(EL *el)
{
  if (IS_LEAF_EL(el))
    return(((struct heat_leaf_data *)LEAF_DATA(el))->estimate);
  else
    return(0.0);
}

static REAL *rw_el_estc(EL *el)
{
  if (IS_LEAF_EL(el))
    return(&((struct heat_leaf_data *)LEAF_DATA(el))->est_c);
  else
    return(NULL);
}

static REAL get_el_estc(EL *el)
{
  if (IS_LEAF_EL(el))
    return(((struct heat_leaf_data *)LEAF_DATA(el))->est_c);
  else
    return(0.0);
}

static REAL time_est = 0.0;

static REAL get_time_est(MESH *mesh, ADAPT_INSTAT *adapt)
{
  return(time_est);
}

/*******************************************************************************
 * For test purposes: exact solution and its gradient (optional)
 ******************************************************************************/

static REAL eval_time_u = 0.0;
static REAL u(const REAL_D x)
{
  return(sin(M_PI*eval_time_u)*exp(-10.0*SCP_DOW(x,x)));
}

static REAL eval_time_u0 = 0.0;
static REAL u0(const REAL_D x)
{
  eval_time_u = eval_time_u0;
  return(u(x));
}


/*******************************************************************************
 * problem data: right hand side, boundary values
 ******************************************************************************/

static REAL eval_time_g = 0.0;
static REAL g(const REAL_D x) /* boundary values, not optional */
{
  eval_time_u = eval_time_g;
  return(u(x));
}

static REAL eval_time_f = 0.0;
static REAL f(const REAL_D x) /* -Delta u, not optional        */
{
  REAL  r2 = SCP_DOW(x,x), ux  = sin(M_PI*eval_time_f)*exp(-10.0*r2);
  REAL  ut = M_PI*cos(M_PI*eval_time_f)*exp(-10.0*r2);
  return(ut - (400.0*r2 - 20.0*DIM_OF_WORLD)*ux);
}


/*******************************************************************************
 * write error and estimator data to files
 ******************************************************************************/

static void write_statistics(const char *path, ADAPT_INSTAT *adapt, int n_dof,
			     REAL space_est, REAL time_est, REAL err_L2)
{
  static FILE *file_ndof = NULL, *file_tau = NULL;
  static FILE *file_space_est = NULL, *file_time_est = NULL, *file_L2_err = NULL;
  const char  *name = fe_space->bas_fcts->name;
  REAL        time = adapt->time;
  char        filename[1024];

  if (!file_ndof) {
    sprintf(filename, "%s/n_dof-%s.agr", path ? path : ".", name);
    file_ndof = fopen(filename, "w");
  }

  if (!file_tau) {
    sprintf(filename, "%s/tau-%s.agr", path ? path : ".", name);
    file_tau       = fopen(filename, "w");
  }

  if (!file_space_est) {
    sprintf(filename, "%s/space_est-%s.agr", path ? path : ".", name);
    file_space_est = fopen(filename, "w");
  }

  if (!file_time_est) {
    sprintf(filename, "%s/time_est-%s.agr", path ? path : ".", name);
    file_time_est  = fopen(filename, "w");
  }

  if (!file_L2_err){
    sprintf(filename, "%s/L2_err-%s.agr", path ? path : ".", name);
    file_L2_err    = fopen(filename, "w");
  }

  if (file_ndof) {
    fprintf(file_ndof, "%.6le %d\n", time, n_dof);
  }

  /* don't print zeros, zeros do not allow log display of estimate */
  if (file_space_est) {
    fprintf(file_space_est, "%.6le %.6le\n", time, MAX(space_est,1.e-20));
  }

  if (time > adapt->start_time) {
    if (file_tau) {
      fprintf(file_tau, "%.6le %.6le\n", time, adapt->timestep);
    }
    /* don't print zeros, zeros do not allow log display of estimate */
    if (file_time_est) {
      fprintf(file_time_est, "%.6le %.6le\n", time, MAX(time_est,1.e-20));
    }
  }

  /* don't print zeros, zeros do not allow log display of error */
  if (file_L2_err) {
    fprintf(file_L2_err, "%.6le %.6le\n", time, MAX(err_L2,1.e-20));
  }

  if (time >= adapt->end_time) {
    if (file_ndof)      fclose(file_ndof);
    if (file_tau)       fclose(file_tau);
    if (file_space_est) fclose(file_space_est);
    if (file_time_est)  fclose(file_time_est);
    if (file_L2_err)    fclose(file_L2_err);
  } else {
    fflush(NULL);
  }
}

/*******************************************************************************
 * Use interpoaltion as "solve on the initial grid". In Principle this
 * should better be a Ritz projection ...
 ******************************************************************************/

static void interpol_u0(MESH *mesh)
{
  dof_compress(mesh);
  interpol(u0, u_h);
}

static void init_timestep(MESH *mesh, ADAPT_INSTAT *adapt)
{
  FUNCNAME("init_timestep");

  INFO(adapt_instat->info,1,
    "---8<---------------------------------------------------\n");
  INFO(adapt_instat->info, 1,"starting new timestep\n");

  dof_copy(u_h, u_old);
}

static void set_time(MESH *mesh, ADAPT_INSTAT *adapt)
{
  FUNCNAME("set_time");

  INFO(adapt->info,1,
       "---8<---------------------------------------------------\n");
  if (adapt->time == adapt->start_time) {
    INFO(adapt->info, 1,"start time: %.4le\n", adapt->time);
  } else {
    INFO(adapt->info, 1,"timestep for (%.4le %.4le), tau = %.4le\n",
	 adapt->time-adapt->timestep, adapt->time,
	 adapt->timestep);
  }

  eval_time_f = adapt->time - (1 - theta)*adapt->timestep;
  eval_time_g = adapt->time;
}

static void close_timestep(MESH *mesh, ADAPT_INSTAT *adapt)
{
  FUNCNAME("close_timestep");
  static REAL err_max = 0.0;                     /* max space-time error    */
  static REAL est_max = 0.0;                     /* max space-time estimate */
  static int  write_fe_data = 0, write_stat_data = 0;
  static int  step = 0;
  static char path[256] = "./";

  REAL        space_est = adapt->adapt_space->err_sum;
  REAL        tolerance = adapt->rel_time_error*adapt->tolerance;

  err_max = MAX(err_max, err_L2);
  est_max = MAX(est_max, space_est + time_est);

  INFO(adapt->info,1,
       "---8<---------------------------------------------------\n");

  if (adapt->time == adapt->start_time) {
    tolerance = adapt->adapt_initial->tolerance;
    INFO(adapt->info,1,"start time: %.4le\n", adapt->time);
  } else {
    tolerance += adapt->adapt_space->tolerance;
    INFO(adapt->info,1,"timestep for (%.4le %.4le), tau = %.4le\n",
			adapt->time-adapt->timestep, adapt->time, 
			adapt->timestep);
  }
  INFO(adapt->info,2,"max. est.  = %.4le, tolerance = %.4le\n", 
		      est_max, tolerance);
  INFO(adapt->info,2,"max. error = %.4le, ratio = %.2lf\n", 
		      err_max, err_max/MAX(est_max,1.0e-20));

  if (!step) {
    GET_PARAMETER(1, "write finite element data", "%d", &write_fe_data);
    GET_PARAMETER(1, "write statistical data", "%d", &write_stat_data);
    GET_PARAMETER(1, "data path", "%s", path);
  }

  /* write mesh and discrete solution to file for post-processing */
  if (write_fe_data) {
    const char *fn;
    
    fn= generate_filename(path, "mesh", step);
    write_mesh_xdr(mesh, fn, adapt->time);
    fn= generate_filename(path, "u_h", step);
    write_dof_real_vec_xdr(u_h, fn);
  }

  step++;

  /* write data about estimate, error, time step size, etc. */
  if (write_stat_data) {
    int n_dof = fe_space->admin->size_used;
    write_statistics(path, adapt, n_dof, space_est, time_est, err_L2);
  }

  if (do_graphics) {
    graphics(mesh, u_h, get_el_est, u, adapt->time);
  }
}

/*******************************************************************************
 * build(): assemblage of the linear system: matrix, load vector,
 *          boundary values, called by adapt_method_stat()
 *          on the first call initialize u_h, f_h, matrix and information
 *          for assembling the system matrix
 *
 * struct op_data: structure for passing information from init_element() to
 *                 LALt() ("operator data").
 * init_element(): initialization on the element; calculates the
 *                 coordinates and |det DF_S| used by LALt; passes these
 *                 values to LALt via user_data,
 *                 called on each element by update_matrix()
 * LALt():         implementation of Lambda id Lambda^t for -Delta u,
 *                 called by update_matrix() after init_element()
 * c():            implementation of 1/tau*m(,)
 ******************************************************************************/

struct op_data
{
  REAL_BD Lambda; /*  the gradient of the barycentric coordinates */
  REAL    det;    /*  |det D F_S|                                 */
};

static bool init_element(const EL_INFO *el_info, const QUAD *quad[3], void *ud)
{
  struct op_data *info = (struct op_data *)ud;

  info->det = el_grd_lambda_0cd(el_info, info->Lambda);

  return 0;
}

const REAL_B *LALt(const EL_INFO *el_info, const QUAD *quad,
		   int iq, void *ud)
{
  static REAL_BB LALt;
  struct op_data *info = (struct op_data *)ud;
  int            i, j, dim = el_info->mesh->dim;

  for (i = 0; i < N_VERTICES(dim); i++) {
    LALt[i][i] = info->det*SCP_DOW(info->Lambda[i], info->Lambda[i]);
    for (j = i+1; j < N_VERTICES(dim); j++) {
      LALt[i][j] = SCP_DOW(info->Lambda[i], info->Lambda[j]);
      LALt[i][j] *= info->det;
      LALt[j][i] = LALt[i][j];
    }
  }

  return (const REAL_B *)LALt;
}


static REAL c(const EL_INFO *el_info, const QUAD *quad, int iq, void *ud)
{
  struct op_data *info = (struct op_data *)ud;

  return info->det;
}

static void build(MESH *mesh, U_CHAR flag)
{
  FUNCNAME("build");

  dof_compress(mesh);
  
  INFO(adapt_instat->adapt_space->info, 2,
       "%d DOFs for %s\n", fe_space->admin->size_used, fe_space->name);

  clear_dof_matrix(matrix);
  dof_set(0.0, f_h);

  update_system_instat(matrix, f_h, adapt_instat->timestep, theta, elsii);

  L2scp_fct_bas(f, load_quad, f_h);
  dirichlet_bound(f_h, u_h, NULL, dirichlet_mask, g);
}

/*******************************************************************************
 * solve(): solve the linear system, called by adapt_method_stat()
 ******************************************************************************/

static void solve(MESH *mesh)
{
  FUNCNAME("solve");
  static REAL tol = 1.e-8, ssor_omega = 1.0;
  static int  max_iter = 1000, info = 2, restart = 0, ssor_iter = 1, ilu_k = 8;
  static OEM_PRECON icon = DiagPrecon;
  static OEM_SOLVER solver = NoSolver;
  const PRECON *precon;

  if (solver == NoSolver) {
    GET_PARAMETER(1, "solver", "%d", &solver);
    GET_PARAMETER(1, "solver tolerance", "%f", &tol);
    GET_PARAMETER(1, "solver precon", "%d", &icon);
    GET_PARAMETER(1, "solver max iteration", "%d", &max_iter);
    GET_PARAMETER(1, "solver info", "%d", &info);
    if (icon == __SSORPrecon) {
    	GET_PARAMETER(1, "precon ssor omega", "%f", &ssor_omega);
    	GET_PARAMETER(1, "precon ssor iter", "%d", &ssor_iter);
    }
    if (icon == ILUkPrecon)
        GET_PARAMETER(1, "precon ilu(k)", "%d", &ilu_k);
    if (solver == GMRes) {
      GET_PARAMETER(1, "solver restart", "%d", &restart);
    }
  }
  
  if (icon == ILUkPrecon)
	  precon = init_oem_precon(matrix, NULL, info, ILUkPrecon, ilu_k);
  else
	  precon = init_oem_precon(matrix, NULL, info, icon, ssor_omega, ssor_iter);
  oem_solve_s(matrix, NULL, f_h, u_h,
	      solver, tol, precon, restart, max_iter, info);
}

/*******************************************************************************
 * Functions for error estimate:
 * estimate():   calculates error estimate via heat_est()
 *               calculates exact error also (only for test purpose),
 *               called by adapt_method_stat()
 * r():          calculates the lower order terms of the element residual
 *               on each element at the quadrature node iq of quad
 *               argument to ellipt_est() and called by ellipt_est()
 ******************************************************************************/

static REAL r(const EL_INFO *el_info, const QUAD *quad, int iq,
	      REAL uh_at_qp, const REAL_D grd_uh_at_qp, REAL t)
{
  /* we simply return gn(), expoiting the fact that the geometry cache
   * of the quadrature already contains the world-coordinates of the
   * quadrature points.
   */
  const QUAD_EL_CACHE *qelc =
    fill_quad_el_cache(el_info, quad, FILL_EL_QUAD_WORLD);

  eval_time_f = t;
  return -f(qelc->world[iq]);
}

static REAL estimate(MESH *mesh, ADAPT_STAT *adapt)
{
  FUNCNAME("estimate");
  static REAL    C[4] = {-1.0, 1.0, 1.0, 1.0};
  REAL_DD        A = {{0.0}};
  FLAGS          r_flag = 0;  /* = (INIT_UH|INIT_GRD_UH), if needed by r()  */
  int            n;
  REAL           space_est;

  for (n = 0; n < DIM_OF_WORLD; n++) {
    A[n][n] = 1.0; /* set diogonal of A; all other elements are zero */
  }

  eval_time_u = adapt_instat->time;

  if (C[0] < 0) {
    C[0] = 1.0;
    GET_PARAMETER(1, "estimator C0", "%f", &C[0]);
    GET_PARAMETER(1, "estimator C1", "%f", &C[1]);
    GET_PARAMETER(1, "estimator C2", "%f", &C[2]);
    GET_PARAMETER(1, "estimator C3", "%f", &C[3]);
  }
  
  time_est = heat_est(u_h, u_old, adapt_instat, rw_el_est, rw_el_estc,
		      -1 /* quad_degree */,
		      C, (const REAL_D *)A, dirichlet_mask,
		      r, r_flag, NULL /* gn() */, 0 /* gn_flag */);
  
  space_est = adapt_instat->adapt_space->err_sum;
  err_L2 = L2_err(u, u_h, NULL, false, false, NULL, NULL);
  
  INFO(adapt_instat->info,2,
       "---8<---------------------------------------------------\n");
  INFO(adapt_instat->info, 2,"time = %.4le with timestep = %.4le\n",
       adapt_instat->time, adapt_instat->timestep);
  INFO(adapt_instat->info, 2,"estimate   = %.4le, max = %.4le\n", space_est,
       sqrt(adapt_instat->adapt_space->err_max));
  INFO(adapt_instat->info, 2,"||u-uh||L2 = %.4le, ratio = %.2lf\n", err_L2,
       err_L2/MAX(space_est,1.e-20));
  
  return adapt_instat->adapt_space->err_sum;
}

static REAL est_initial(MESH *mesh, ADAPT_STAT *adapt)
{
  err_L2 = adapt->err_sum = L2_err(u0, u_h, NULL, false, false, rw_el_est, &adapt->err_max);
  return(adapt->err_sum);
}


/*******************************************************************************
 * main program
 ******************************************************************************/

int main(int argc, char **argv)
{
  FUNCNAME("main");
  MACRO_DATA     *data;
  MESH           *mesh;
  const BAS_FCTS *lagrange;
  int             n_refine = 0, p = 1, dim;
  char            filename[PATH_MAX];
  REAL            fac = 1.0;
  static OPERATOR_INFO  o_info2, o_info0;
  static struct op_data op_data;

  /*****************************************************************************
   * first of all, initialize the access to parameters of the init file
   ****************************************************************************/
  parse_parameters(argc, argv, "INIT/heat.dat");

  GET_PARAMETER(1, "mesh dimension", "%d", &dim);
  GET_PARAMETER(1, "macro file name", "%s", filename);
  GET_PARAMETER(1, "global refinements", "%d", &n_refine);
  GET_PARAMETER(1, "parameter theta", "%e", &theta);
  GET_PARAMETER(1, "polynomial degree", "%d", &p);
  GET_PARAMETER(1, "online graphics", "%d", &do_graphics);
  
  /*****************************************************************************
  *  get a mesh, and read the macro triangulation from file
  ****************************************************************************/
  data = read_macro(filename);
  mesh = GET_MESH(dim, "ALBERTA mesh", data,
		  NULL /* init_node_projection() */,
		  NULL /* init_wall_trafos() */);
  free_macro_data(data);


  init_leaf_data(mesh, sizeof(struct heat_leaf_data),
		 NULL /* refine_leaf_data() */,
		 NULL /* coarsen_leaf_data() */);

  /* Finite element spaces can be added at any time, but it is more
   * efficient to do so before refining the mesh a lot.
   */
  lagrange = get_lagrange(mesh->dim, p);
  TEST_EXIT(lagrange, "no lagrange BAS_FCTS\n");

  fe_space = get_fe_space(mesh, lagrange->name, lagrange, 1, ADM_FLAGS_DFLT);

  global_refine(mesh, n_refine * dim, FILL_NOTHING);

  /*****************************************************************************
   *  initialize the global variables shared across build(), solve()
   *  and estimate().
   ****************************************************************************/
  matrix = get_dof_matrix("A", fe_space, NULL);
  f_h    = get_dof_real_vec("f_h", fe_space);
  u_h    = get_dof_real_vec("u_h", fe_space);
  u_h->refine_interpol = fe_space->bas_fcts->real_refine_inter;
  u_h->coarse_restrict = fe_space->bas_fcts->real_coarse_inter;
  u_old  = get_dof_real_vec("u_old", fe_space);
  u_old->refine_interpol = fe_space->bas_fcts->real_refine_inter;
  u_old->coarse_restrict = fe_space->bas_fcts->real_coarse_inter;
  dof_set(0.0, u_h);      /*  initialize u_h  !                          */

  BNDRY_FLAGS_ALL(dirichlet_mask); /* Only Dirichlet b.c. supported here */

  /*****************************************************************************
   * aquire quadratures and fill-info.
   ****************************************************************************/

  /* Fill operator info for 2nd order term. */
  o_info2.row_fe_space   = o_info2.col_fe_space = u_h->fe_space;
  o_info2.quad[2]        = get_quadrature(dim, 2*(p-1));
  o_info2.LALt.real      = LALt;
  o_info2.LALt_pw_const  = true;
  o_info2.LALt_symmetric = true;
  BNDRY_FLAGS_CPY(o_info2.dirichlet_bndry, dirichlet_mask);
  o_info2.init_element   = init_element;
  o_info2.user_data      = &op_data;

  /* Fill operator info for mass matrix. */
  o_info0.row_fe_space   = o_info0.col_fe_space = u_h->fe_space;
  o_info0.quad[0]        = get_quadrature(dim, 2*p);
  o_info0.c.real         = c;
  o_info0.c_pw_const     = true;
  BNDRY_FLAGS_CPY(o_info0.dirichlet_bndry, dirichlet_mask);
  o_info0.init_element   = NULL; /* no need to call init_element twice */
  o_info0.user_data      = &op_data;

  /* Then get the assemble handle which implements the theta splitting scheme */
  elsii = fill_sys_info_instat(&o_info2, &o_info0, u_h);

  /* Get the load quadrature */
  load_quad = get_quadrature(dim, 2*p);

  /*****************************************************************************
   *  init adapt_instat structure
   ****************************************************************************/
  adapt_instat = get_adapt_instat(dim, "heat", "adapt", 2, adapt_instat);

  if (do_graphics) {
    graphics(mesh, NULL, NULL, NULL, adapt_instat->start_time);
  }

  /*****************************************************************************
   *  adapt time step size to refinement level and polynomial degree,
   *  based on the known L2-error error estimates.
   ****************************************************************************/
  if (theta < 0.5) { 
    WARNING("You are using the explicit Euler scheme\n");
    WARNING("Use a sufficiently small time step size!!!\n");
    fac = 1.0e-3;
  }

  if (theta == 0.5) {
    adapt_instat->timestep *= fac*pow(2, -(REAL)((p+1)*(n_refine))/2.0);
  } else {
    adapt_instat->timestep *= fac*pow(2, -(REAL)((p+1)*(n_refine)));
  }
  MSG("using initial timestep size = %.4le\n", adapt_instat->timestep);

  eval_time_u0 = adapt_instat->start_time;

  adapt_instat->adapt_initial->get_el_est = get_el_est;
  adapt_instat->adapt_initial->estimate = est_initial;
  adapt_instat->adapt_initial->solve = interpol_u0;

  adapt_instat->adapt_space->get_el_est   = get_el_est;
  adapt_instat->adapt_space->get_el_estc  = get_el_estc;
  adapt_instat->adapt_space->estimate = estimate;
  adapt_instat->adapt_space->build_after_coarsen = build;
  adapt_instat->adapt_space->solve = solve;

  adapt_instat->init_timestep  = init_timestep;
  adapt_instat->set_time       = set_time;
  adapt_instat->get_time_est   = get_time_est;
  adapt_instat->close_timestep = close_timestep;

  /*****************************************************************************
   * ... off we go ...
   ****************************************************************************/
  adapt_method_instat(mesh, adapt_instat);
  
  free_sys_info_instat(elsii); /* call the destructor, just for fun. */

  WAIT_REALLY;

  return 0;
}
