#include "../include/Charge.h"

/**
  Constructor
    * Store references and copies of the data
    * Setup DoFs
    * Init the Linear Algebra Data
**/
template<int dim> Charge<dim>::Charge( const Triangulation<dim> &tria, Material_Parameters &params,
                                       const unsigned deg, const unsigned u_prec, const double ee, const double thres,
                                       const unsigned sweeps ) :
    Problem<dim>( tria, 1, ee, true, u_prec ), conductivity( params.K ), lambda( params.lambda ),
    get_charge( this->dh, this->sol, "charge" ){
  prec_data.elliptic = true;
  prec_data.higher_order_elements = deg > 1;
  prec_data.aggregation_threshold = thres;
  prec_data.smoother_sweeps = sweeps;
  this->fe.push_back( FE_Q<dim> ( deg ) );
  this->fe.push_back( FE_Nothing<dim>() );
}

/**
  Destructor
    * Clear data
**/
template<int dim> Charge<dim>::~Charge() {}

/**
  ReinitPrec
    * Reinitialize the preconditioner
**/
template<int dim> void Charge<dim>::ReinitPrec(){
  prec.initialize( this->K, prec_data );
}

/**
  DoSolve
    * Solve the system to the give accuracy
**/
template<int dim> void Charge<dim>::DoSolve( SolverControl &control ){
  control.set_tolerance( std::max( this->eps, control.tolerance() ) );
  TrilinosWrappers::SolverGMRES gmres( control );
  gmres.solve( this->K, this->sol, this->rhs, prec );
}

/**
  SetInitialData
    * q = 0
**/
template<int dim> void Charge<dim>::SetInitialData(){
  this->sol = 0.;
}

/**
  AssembleSystem
    * Clear global matrix and rhs
    * Create a sample of the scratch and local data for the particular problem that we want to have
    * If running on multithreaded mode create a list of tasks to assemble each cell
    * Else do a for loop on cells
**/
template<int dim> void Charge<dim>::AssembleSystem( std::vector< AsFunction<dim> *> &data, const double,
                                                    const bool m_threaded ) {
  Assert( data.size() == 3, ExcDimensionMismatch( data.size(), 3 ) );
  this->K = 0.;
  this->rhs = 0.;
  hp::QCollection<dim> quad( QGauss<dim>( this->fe[0].degree + 1 ) );
  quad.push_back( QGauss<dim>( this->fe[1].degree + 1 ) );
  Iterator start, end, cc;
  start(0) = this->dh.begin_active();
  end(0) = this->dh.end();
  for( unsigned i=0; i<3; ++i ){
    start( i+1 ) = data[i]->get_begin();
    end( i+1 ) = data[i]->get_end();
  }
  data[0]->reset( quad, update_gradients ); // voltage
  data[1]->reset( quad, update_values ); // phase
  data[2]->reset( quad, update_values | update_gradients ); // velocity
  typename Problem<dim>::PerTaskData ptdata( this->fe.max_dofs_per_cell() );
  ScratchData scratch( this->fe, data, quad, update_values | update_gradients | update_JxW_values );
  if( m_threaded )
    WorkStream::run( start, end,
                     std_cxx1x::bind( &Charge<dim>::AssembleCell, this, std_cxx1x::_1, std_cxx1x::_2, std_cxx1x::_3 ),
                     std_cxx1x::bind( &Charge<dim>::CopyToGlob, this, std_cxx1x::_1 ),
                     scratch, ptdata );
  else{
    for( cc = start; cc not_eq end; ++cc ){
      AssembleCell( cc, scratch, ptdata );
      this->CopyToGlob ( ptdata );
    }
  }
  for( unsigned i=0; i<3; ++i )
    data[i]->reset();
}

/**
  AssembleCell
    * reinit fe_values and fe_functions
    * assemble the local matrix and rhs
    * m_ij = int_K p_i p_j /dt + \lambda K(\phi) \grad p_i \grad p_j
    * f_i = int_K  q^n p_i/dt + q^n v^n \grad p_i - K(\phi)\grad V \grad p_i
**/
template<int dim> void Charge<dim>::AssembleCell( const Iterator Its, ScratchData &scratch,
                                                   typename Problem<dim>::PerTaskData &data ){
  scratch.fe_val.reinit( Its(0) );
  const FEValues<dim> &fe_val = scratch.fe_val.get_present_fe_values();
  const unsigned nqp = fe_val.n_quadrature_points;
  data.dpc = fe_val.get_fe().dofs_per_cell;
  scratch.loc_charge.resize( nqp );
  scratch.loc_phase.resize( nqp );
  scratch.loc_vel.resize( nqp );
  scratch.loc_div_vel.resize( nqp );
  scratch.loc_grad_volt.resize( nqp );
  scratch.voltage.reinit( Its(1) );
  scratch.phase.reinit( Its(2) );
  scratch.velocity.reinit( Its(3) );
  fe_val.get_function_values( this->sol, scratch.loc_charge );
  scratch.voltage.get_function_gradients( scratch.loc_grad_volt );
  scratch.velocity.get_function_values( scratch.loc_vel );
  scratch.velocity.get_function_divergences( scratch.loc_div_vel );
  scratch.phase.get_function_values( scratch.loc_phase );
  const double dd = 0.1*Its(0)->diameter()/double( this->fe[0].degree );
  data.loc_rhs.reinit( data.dpc );
  data.loc_m.reinit( data.dpc, data.dpc );
  for( unsigned q=0; q<nqp; ++q )
    for( unsigned i=0; i<data.dpc; ++i ){
      data.loc_rhs(i) += fe_val.JxW(q)*(
                            scratch.loc_charge[q]*fe_val.shape_value( i, q )/this->dt
                            +(
                                scratch.loc_charge[q]*scratch.loc_vel[q]
                                - conductivity( scratch.loc_phase[q] )*scratch.loc_grad_volt[q]
                            )*fe_val.shape_grad( i, q )
                        );
      for( unsigned j=0; j<data.dpc; ++j )
        data.loc_m( i, j ) += fe_val.JxW(q)*(
                                  fe_val.shape_value( j, q )*fe_val.shape_value( i, q )/this->dt
                                  + lambda*conductivity( scratch.loc_phase[q] )
                                    *fe_val.shape_grad( j, q )*fe_val.shape_grad( i, q )
                              );
    }
  data.ldi.resize( data.dpc );
  Its(0)->get_dof_indices( data.ldi );
}

/**
  PreRefinementPreffix
    * Copy the solution to the transfer vector
**/
template<int dim> void Charge<dim>::PreRefinementPreffix(){
  this->x_sol[0] = this->sol;
}

/**
  PostRefinementSuffix
    * Copy the solution from the transfer vector
**/
template<int dim> void Charge<dim>::PostRefinementSuffix( const std::vector<TrilinosWrappers::Vector> &sol_tmp ){
  this->sol = sol_tmp[0];
}


template class Charge<DIM>;
