#include <getfem/getfem_mesh.h>
#include <getfem/getfem_mesh_im.h>
#include <getfem/bgeot_mesh_structure.h>
#include <getfem/getfem_export.h>
#include <getfem/getfem_regular_meshes.h>
#include <getfem/getfem_assembling.h>
#include <getfem/getfem_mesh_region.h>
#include <getfem/getfem_mesh_fem_global_function.h>
#include <iostream>
#include <vector>

/* This program calculates global conduction matrix for heat conduction
 * in a cylinder using axisymmetrical setup
 *
 *
 * Compare J.E. Akin "Finite Element Analysis with Error Estimators"
 * Chapter 8, "Cylindrical analysis problems", page 218.
 */

int main(int argc, char *argv[]) {
  try {

  /* Create mesh by building it from scratch 
   *
   * Inner radious r=1
   * Outer radious r=2
   * 4 linear elements of equal length
   */
  getfem::mesh mesh;

  int nnodes = 5;
  std::vector<bgeot::size_type> ind(nnodes);
  double dr = 0.25;
  double r0 = 1.0;
  for (int i = 0; i<nnodes; i++) {
    ind[i] = mesh.add_point(bgeot::base_node(r0 + i*dr, 0.0));
  }  
  for (int i = 0; i<nnodes-1; i++) {
     mesh.add_segment(ind[i], ind[i+1]);
  }
  /* Create FEM description of the problem */
   
  getfem::mesh_fem femT(mesh);

  getfem::pfem pf = getfem::fem_descriptor("FEM_PK(1,1)");
  femT.set_finite_element(pf);
  
  getfem::mesh_im mim(mesh);
  mim.set_integration_method(
        getfem::int_method_descriptor("IM_GAUSS1D(2)"));

  getfem::size_type Ndof =  femT.nb_dof();

  /* Create matrix. Since the problem is small
   * dense matrices are used here which also helps in printing them.
   *
   * The equation: M*T = P
   */ 
  gmm::dense_matrix<double> M(Ndof, Ndof); /* characteristic matrix */

  struct global_function_r : public getfem::global_function {
    getfem::scalar_type 
       val(const getfem::fem_interpolation_context& c) const {
      getfem::base_node x = c.xreal();
      return gmm::vect_norm2(x);
    }
    /* one should also provide methods for gradient and hesian 
     * calculations if they are used. Here we can skip this
     */
  };
 
  global_function_r r_gf;
  getfem::pglobal_function pgf_r = new global_function_r;
  getfem::mesh_fem_global_function r_fem(mesh);
  r_fem.set_functions(pgf_r);

  getfem::generic_assembly assem;
  assem.push_mi(mim);
  assem.push_mf(femT);
  assem.push_mf(r_fem);
  assem.push_mat(M);
  assem.set(
            "t=2*comp(Grad(#1).Grad(#1).Base(#2))(:,k,:,k,1);"
            "M$1(#1,#1)+=t;"
           );
  /* CAUTION : the calculated matrix should be scaled by pi * K
   * where K is the conductivity constant to make a true conductivity
   * matrix
   */
  assem.assembly();
  std::cout << M << "\n";

  } GMM_STANDARD_CATCH_ERROR;
}