#include "../include/Problem.h"

/**
  Constructor
    * Copy all the data
**/
template<int dim> Problem<dim>::Problem( const Triangulation<dim> &tria, const unsigned _n_transfers, const double _eps,
                                         bool _prec_needs_reinit, const unsigned _update_prec ) :
    tri( tria ), mesh_has_changed( false ), dh( tri ), steps(0), eps( _eps ), update_prec( _update_prec ),
    prec_needs_reinit( _prec_needs_reinit ), n_of_transfers( _n_transfers ), x_sol( n_of_transfers ), transfer( dh ) {}

/**
  Destructor
    * Clear data
**/
template<int dim> Problem<dim>::~Problem(){
  dh.clear();
}

/**
  set_dt
    * dt = _dt
**/
template<int dim> void Problem<dim>::set_dt( const double _dt ){
  dt = _dt;
}

/**
  get_fe
    * returns a constant reference to the used Finite Element object
**/
template<int dim> const hp::FECollection<dim> &Problem<dim>::get_fe() const{
  return fe;
}

/**
  size
    * return number of degrees of freedom
**/
template<int dim> unsigned Problem<dim>::size() const{
  return n_dofs;
}

/**
  init
    * SetupDoFs
    * InitLAData
    * SetInitialData
    * Connect <PreRefinement> and <PostRefinement> to the signals emmited by <Triangulation> before
      and after refinement, respectively.
**/
template<int dim> void Problem<dim>::init(){
  SetupDoFs();
  InitLAData();
  SetInitialData();
  tri.signals.pre_refinement.connect( std_cxx1x::bind( &Problem<dim>::PreRefinement,
                                                       std_cxx1x::ref( *this ) ) );
  tri.signals.post_refinement.connect( std_cxx1x::bind( &Problem<dim>::PostRefinement,
                                                        std_cxx1x::ref( *this ) ) );
}

/**
  SetupDoFs
    * Distribute degrees of freedom
    * Set number of degrees of freedom
    * make hanging node constraints
**/
template<int dim> void Problem<dim>::SetupDoFs(){
  typename hp::DoFHandler<dim>::active_cell_iterator cell = dh.begin_active(), end = dh.end();
  for( ; cell not_eq end; ++cell ){
    const unsigned char material_id = cell->material_id();
    cell->set_active_fe_index( ( material_id == 'f' )?0:1 );
  }
  dh.distribute_dofs( fe );
  n_dofs = dh.n_dofs();
  constraints.clear();
  DoFTools::make_hanging_node_constraints( dh, constraints );
  SetupDoFsSuffix();
  constraints.close();
}

/**
  SetupDoFsSuffix
    * By default do nothing
**/
template<int dim> void Problem<dim>::SetupDoFsSuffix() {}

/**
  InitLAData
    * Create sparsity pattern
    * Reinit vectors to the correct size
**/
template<int dim> void Problem<dim>::InitLAData(){
  {
    CompressedSparsityPattern sp( n_dofs, n_dofs );
    DoFTools::make_sparsity_pattern( dh, sp, constraints );
    K.reinit( sp );
  }
  sol.reinit( n_dofs );
  rhs.reinit( n_dofs );
  InitLADataSuffix();
}

/**
  InitLADataSuffix
    * By default do nothing
**/
template<int dim> void Problem<dim>::InitLADataSuffix() {}

/**
  solve
    * AssembleSystem
    * SolveSystem
**/
template<int dim> void Problem<dim>::solve( std::vector< AsFunction<dim> *> &data, const double time ){
  AssembleSystem( data, time );
  SolveSystem();
}

/**
  SolveSystem
    * Do the Preffix
    * Check, and if needed, call ReinitPrec
    * Solve System
    * Suffix
**/
template<int dim> void Problem<dim>::SolveSystem(){
  SolveSystemPreffix();
  if( steps == 0 )
    ReinitPrec();
  else
    if( ( prec_needs_reinit and ( steps % update_prec == 0 ) ) or mesh_has_changed ){
      ReinitPrec();
      mesh_has_changed = false;
    }
  steps++;
  SolverControl control( n_dofs, eps*rhs.l2_norm() );
  DoSolve( control );
  constraints.distribute( sol );
  SolveSystemSuffix();
}

/**
  SolveSystemPreffix
    * By default do nothing
**/
template<int dim> void Problem<dim>::SolveSystemPreffix() {}

/**
  SolveSystemSuffix
    * By default do nothing
**/
template<int dim> void Problem<dim>::SolveSystemSuffix() {}

/**
  CopyToGlob
    * Copy the local contributions to the system matrix and rhs
**/
template<int dim> void Problem<dim>::CopyToGlob( const PerTaskData &data ){
  constraints.distribute_local_to_global( data.loc_m, data.loc_rhs, data.ldi, K, rhs );
}

/**
  PreRefinement
    * Call the preffix, which should copy
**/
template<int dim> void Problem<dim>::PreRefinement(){
  if( steps ){
    PreRefinementPreffix();
    transfer.prepare_for_coarsening_and_refinement( x_sol );
  }
}

template<int dim> void Problem<dim>::PostRefinement(){
  SetupDoFs();
  InitLAData();
  if( not steps )
    SetInitialData();
  else{
    std::vector<TrilinosWrappers::Vector> x_sol_tmp( n_of_transfers, TrilinosWrappers::Vector( n_dofs ) );
    transfer.interpolate( x_sol, x_sol_tmp );
    for( unsigned i=0; i<n_of_transfers; ++i )
      constraints.distribute( x_sol_tmp[i] );
    PostRefinementSuffix( x_sol_tmp );
  }
  mesh_has_changed = true;
}

template<int dim> bool Problem<dim>::isNeeded( const typename hp::DoFHandler<dim>::active_cell_iterator &c,
                                               const unsigned face ) const{
  bool needed = false;
  if( c->active_fe_index() not_eq 0 )
    needed = false;
  else{
    if( c->face(face)->at_boundary() )
      needed = false;
    else
      needed = c->neighbor(face)->active_fe_index() not_eq 0;
  }
  return needed;
}

template class Problem<DIM>;
