#include "../include/NSE.h"

////////////////// Pressure
/**
  Constructor
    * Initialize data
**/
template<int dim> Pressure<dim>::Pressure( const Triangulation<dim> &tria, Material_Parameters &params,
                                           const unsigned deg, const double ee, const double thres,
                                           const unsigned sweeps ) :
    Problem<dim>( tria, 3, ee, false ), rho_min( std::min( params.rho(100.), params.rho(-100.) ) ),
    get_pressure( this->dh, p, "pres" ), get_extrapolated_pressure( this->dh, pres_extr, "psharp" ){
  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> Pressure<dim>::~Pressure() {}

/**
  SetupDoFsSuffix
    * Add a zero constraint to filter the constant modes
**/
template<int dim> void Pressure<dim>::SetupDoFsSuffix(){
  this->constraints.add_line ( 0 );
}

/**
  InitLADataSuffix
    * Initialize p and p_extr
**/
template<int dim> void Pressure<dim>::InitLADataSuffix(){
  p.reinit( this->n_dofs );
  pres_extr.reinit( this->n_dofs );
}

/**
  SetInitialData
    * p = phi = pres_extr = 0
**/
template<int dim> void Pressure<dim>::SetInitialData(){
  p = 0.;
  this->sol = 0.;
  pres_extr = 0.;
}

/**
  ReinitPrec
    * Reinit Preconditioner
**/
template<int dim> void Pressure<dim>::ReinitPrec(){
  prec.initialize( this->K, prec_data );
}

/**
  DoSolve
    * Solve to the given accuracy
**/
template<int dim> void Pressure<dim>::DoSolve( SolverControl &control ){
  TrilinosWrappers::SolverCG cg( control );
  cg.solve( this->K, this->sol, this->rhs, prec );
}

/**
  SolveSystemSuffix
    * p^{n+1} = p^n + \phi^{n+1}
    * p^\sharp = p^n + \phi^n
**/
template<int dim> void Pressure<dim>::SolveSystemSuffix(){
  p += this->sol;
  pres_extr.equ( 1., p, 1., this->sol );
}

/**
  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 Pressure<dim>::AssembleSystem( std::vector< AsFunction<dim> *> &data,
                                                      const double, const bool m_threaded ){
  Assert( data.size() == 1, ExcDimensionMismatch( data.size(), 1 ) );
  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 ) );
  data[0]->reset( quad, update_gradients );
  typename Problem<dim>::PerTaskData ptdata( this->fe.max_dofs_per_cell() );
  ScratchData scratch( this->fe, data, quad, update_values | update_gradients | update_JxW_values );
  Iterator start, end, cc;
  start(0) = this->dh.begin_active();
  start(1) = data[0]->get_begin();
  end(0) = this->dh.end();
  end(1) = data[0]->get_end();
  if( m_threaded )
    WorkStream::run( start, end,
                     std_cxx1x::bind( &Pressure<dim>::AssembleCell, this, std_cxx1x::_1, std_cxx1x::_2, std_cxx1x::_3 ),
                     std_cxx1x::bind( &Pressure<dim>::CopyToGlob, this, std_cxx1x::_1 ),
                     scratch, ptdata );
  else{
    for( cc = start; cc not_eq end; ++cc ){
      AssembleCell( cc, scratch, ptdata );
      CopyToGlob( ptdata );
    }
  }
  data[0]->reset();
}

/**
  AssembleCell
    * m_ij = grad p_i grad_pj
    * f_i = u grad p_i
**/
template<int dim> void Pressure<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_div_vel.resize( nqp );
  scratch.velocity.reinit( Its(1) );
  scratch.velocity.get_function_divergences( scratch.loc_div_vel );
  data.loc_m.reinit( data.dpc, data.dpc );
  data.loc_rhs.reinit( data.dpc );
  for( unsigned q=0; q<nqp; ++q )
    for( unsigned i=0; i<data.dpc; ++i ){
      data.loc_rhs(i) += -rho_min*fe_val.JxW(q)*scratch.loc_div_vel[q]*fe_val.shape_value( i, q )/( this->dt );
      for( unsigned j=0; j<data.dpc; ++j )
        data.loc_m( i, j ) += fe_val.JxW( 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 );
}

template<int dim> void Pressure<dim>::PreRefinementPreffix(){
  this->x_sol[0] = this->sol;
  this->x_sol[1] = p;
  this->x_sol[2] = pres_extr;
}

template<int dim> void Pressure<dim>::PostRefinementSuffix( const std::vector<TrilinosWrappers::Vector> &sol_tmp ){
  this->sol = sol_tmp[0];
  p = sol_tmp[1];
  pres_extr = sol_tmp[2];
}

///////////////////////////

///////////////////// Velocity
/**
  Constructor
    * Copy parameters
**/
template<int dim> Velocity<dim>::Velocity( const Triangulation<dim> &tria, Material_Parameters &params,
                                           const unsigned deg, const unsigned u_prec, const double ee,
                                           const double thres, const unsigned sweeps, const unsigned K_size ) :
    Problem<dim>( tria, 2, ee, true, u_prec ),
    density( params.rho ), viscosity( params.eta ), slip_coeff( params.beta ), gamma_fs( params.gamma_fs ),
    alpha( params.alpha ), lambda( params.lambda ), gamma( params.gamma ), delta( params.delta ),
    Krylov_size( K_size ), get_velocity( this->dh, this->sol, "vel", AsFunction<dim>::isVectorValued, 0 ){
  prec_data.elliptic = true;
  prec_data.higher_order_elements = true;
  prec_data.aggregation_threshold = thres;
  prec_data.smoother_sweeps = sweeps;
  this->fe.push_back( FESystem<dim>( FE_Q<dim>( deg+1 ), dim ) );
  this->fe.push_back( FESystem<dim>( FE_Nothing<dim>(), dim ) );
}

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

/**
  SetupDoFsSuffix
**/
template<int dim> void Velocity<dim>::SetupDoFsSuffix(){
  // Add no-slip to the constant density parts
  {
    ZeroFunction<dim> zero(dim);
    typename FunctionMap<dim>::type bb;
    bb['g'] = &zero;
    std::vector<bool> mask( dim, true );
    VectorTools::interpolate_boundary_values( this->dh, bb, this->constraints, mask );
  }
  // Add impermeability constraints
  {
    Assert( dim==2, ExcNotImplemented() );
    std::vector<unsigned> ldi( this->fe.max_dofs_per_face() );
    typename hp::DoFHandler<dim>::active_cell_iterator cell = this->dh.begin_active(), end = this->dh.end();
    for( ; cell not_eq end; ++cell ){
      for( unsigned f=0; f<GeometryInfo<dim>::faces_per_cell; ++f ){
        if( isNeeded( cell, f ) ){
          cell->face(f)->get_dof_indices( ldi, 0 );
          for( unsigned i=0; i<ldi.size(); ++i )
            if( this->fe[0].face_system_to_component_index(i).first == f/2 )
              this->constraints.add_line( ldi[i] );
        }
      }
    }
  }
}

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

/**
  ReinitPrec
    * Reinit the preconditioner
    * Notice that this is the only point where you need deal.II >= 7.2
**/
template<int dim> void Velocity<dim>::ReinitPrec(){
  std::vector< std::vector<bool> > constant_modes;
  DoFTools::extract_constant_modes( this->dh, std::vector<bool>( dim, true ), constant_modes );
  prec_data.constant_modes = constant_modes;
  prec.initialize( this->K, prec_data );
}

/**
  InitLADataSuffix
    * Set the sizes of oldsol and olditsol
**/
template<int dim> void Velocity<dim>::InitLADataSuffix(){
  oldsol.reinit( this->n_dofs );
  olditsol.reinit( this->n_dofs );
}

/**
  SolveSystemPreffix
    * Copy the solution to olditsol
**/
template<int dim> void Velocity<dim>::SolveSystemPreffix(){
  olditsol = this->sol;
}

/**
  advance
    * advance in time
**/
template<int dim> void Velocity<dim>::advance(){
  oldsol = this->sol;
}

/**
  it_error
    * compute the error of the inner iteration
**/
template<int dim> double Velocity<dim>::it_error(){
  olditsol.sadd(-1., this->sol );
  return olditsol.l2_norm();
}

/**
  DoSolve
    * Solve to the given accuracy
**/
template<int dim> void Velocity<dim>::DoSolve( SolverControl &control ){
  TrilinosWrappers::SolverGMRES::AdditionalData data( false, Krylov_size );
  TrilinosWrappers::SolverGMRES gmres( control, data );
  gmres.solve( this->K, this->sol, this->rhs, prec );
}

/**
  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 Velocity<dim>::AssembleSystem( std::vector< AsFunction<dim> *> &data, const double,
                                                      const bool m_threaded ){
  Assert( data.size() == 6, ExcDimensionMismatch( data.size(), 6 ) );
  this->K = 0.;
  this->rhs = 0.;
  hp::QCollection<dim> quad( QGauss<dim>( this->fe[0].degree + 2 ) );
  quad.push_back( QGauss<dim>( this->fe[1].degree + 2 ) );
  hp::QCollection<dim-1> fquad( QGauss<dim-1>( this->fe[0].degree + 2 ) );
  fquad.push_back( QGauss<dim-1>( this->fe[1].degree + 2 ) );
  Iterator start, end, cc;
  start(0) = this->dh.begin_active();
  end(0) = this->dh.end();
  for( unsigned i=0; i<6; ++i ){
    start( i+1 ) = data[i]->get_begin();
    end( i+1 ) = data[i]->get_end();
  }
  data[0]->reset( quad , update_values | update_gradients ); // phase
  data[1]->reset( quad, update_values ); // mu
  data[2]->reset( quad, update_values | update_gradients ); //q
  data[3]->reset( quad, update_gradients ); // V
  data[4]->reset( quad, update_values ); // p
  data[5]->reset( quad, update_values ); // phase_old
  data[0]->reset_face( fquad, update_values | update_gradients ); //phase
  data[5]->reset_face( fquad, update_values ); // oldphase
  typename Problem<dim>::PerTaskData ptdata( this->fe.max_dofs_per_cell() );
  ScratchData scratch( this->fe, data,
                       quad, update_values | update_gradients | update_JxW_values,
                       fquad, update_values | update_normal_vectors | update_JxW_values );
  if( m_threaded )
    WorkStream::run( start, end,
                     std_cxx1x::bind( &Velocity<dim>::AssembleCell, this, std_cxx1x::_1, std_cxx1x::_2, std_cxx1x::_3 ),
                     std_cxx1x::bind( &Velocity<dim>::CopyToGlob, this, std_cxx1x::_1 ),
                     scratch, ptdata );
  else{
    for( cc = start; cc not_eq end; ++cc ){
      AssembleCell( cc, scratch, ptdata );
      CopyToGlob ( ptdata );
    }
  }
  for( unsigned i=0; i<6; ++i )
    data[i]->reset();
  data[0]->reset_face();
  data[5]->reset_face();
}

/**
  AssembleCell
    * Assemble one cell for Navier Stokes
**/
template<int dim> void Velocity<dim>::AssembleCell( const Iterator Its, ScratchData &scratch,
                                                    typename Problem<dim>::PerTaskData &data ) const{
  FEValuesExtractors::Vector vv(0);
  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_phase.resize( nqp );
  scratch.loc_old_phase.resize( nqp );
  scratch.loc_q.resize( nqp );
  scratch.loc_psharp.resize( nqp );
  scratch.loc_divu.resize( nqp );
  scratch.loc_mu.resize( nqp );
  scratch.loc_grad_q.resize( nqp );
  scratch.loc_grad_V.resize( nqp );
  scratch.loc_grad_phase.resize( nqp );
  scratch.loc_u.resize( nqp );
  scratch.phase.reinit( Its(1) );
  scratch.mu.reinit( Its(2) );
  scratch.q.reinit( Its(3) );
  scratch.V.reinit( Its(4) );
  scratch.psharp.reinit( Its(5) );
  scratch.oldphase.reinit( Its(6) );
  scratch.phase.get_function_values( scratch.loc_phase );
  scratch.phase.get_function_gradients( scratch.loc_grad_phase );
  scratch.mu.get_function_values( scratch.loc_mu );
  scratch.q.get_function_values( scratch.loc_q );
  scratch.q.get_function_gradients( scratch.loc_grad_q );
  scratch.V.get_function_gradients( scratch.loc_grad_V );
  scratch.psharp.get_function_values( scratch.loc_psharp );
  scratch.oldphase.get_function_values( scratch.loc_old_phase );
  fe_val[vv].get_function_divergences( this->oldsol, scratch.loc_divu );
  fe_val[vv].get_function_values( this->oldsol, scratch.loc_u );
  data.loc_m.reinit( data.dpc, data.dpc );
  data.loc_rhs.reinit( data.dpc );
  for( unsigned q = 0; q < nqp; ++q ){
    const double rhoavg = density.average( scratch.loc_phase[q], scratch.loc_old_phase[q] ),
                 rho_q = density( scratch.loc_phase[q] ),
                 eta_q = viscosity( scratch.loc_phase[q] ),
                 rho_prime_q = density.deriv( scratch.loc_phase[q] ),
                 dt = this->dt;
    for( unsigned i=0; i<data.dpc; ++i ){
      data.loc_rhs(i) += fe_val.JxW( q )*(
                            (
                                (
                                    (
                                      rho_q
                                      + 0.5*rho_prime_q*( scratch.loc_phase[q] - scratch.loc_old_phase[q] )
                                    )/ dt
                                )*scratch.loc_u[q]
                                + scratch.loc_mu[q]*scratch.loc_grad_phase[q]
                                - scratch.loc_q[q]*( lambda*scratch.loc_grad_q[q] + scratch.loc_grad_V[q] )
                            )*fe_val[vv].value ( i, q )
                            + scratch.loc_psharp[q]*fe_val[vv].divergence( i, q )
                        );
      for( unsigned j=0; j<data.dpc; ++j )
        data.loc_m( i, j ) += fe_val.JxW( q )*(
                                  (
                                    ( rhoavg/dt )
                                    + 0.5*(
                                        rho_prime_q*scratch.loc_grad_phase[q]*scratch.loc_u[q]
                                        + rho_q*scratch.loc_divu[q]
                                    )
                                  )*fe_val[vv].value( j, q )*fe_val[vv].value( i, q )
                                  + rho_q*( scratch.loc_u[q]*fe_val[vv].gradient( j, q ) )*fe_val[vv].value ( i, q )
                                  + eta_q*fe_val[vv].symmetric_gradient( j, q )*fe_val[vv].symmetric_gradient( i, q )
                              );
    }
  }
  // loop over faces
  for( unsigned f=0; f<GeometryInfo<dim>::faces_per_cell; ++f ){
    if( isNeeded( Its(0), f ) ){
      scratch.fe_face_val.reinit( Its(0), f );
      scratch.phase.reinit( Its(1), f );
      scratch.oldphase.reinit( Its(6), f );
      const FEFaceValues<dim> &fe_face_val = scratch.fe_face_val.get_present_fe_values();
      const unsigned nfqp = fe_face_val.n_quadrature_points;
      scratch.face_phase.resize( nfqp );
      scratch.face_old_phase.resize( nfqp );
      scratch.face_grad_phase.resize( nfqp );
      scratch.face_base.resize( data.dpc );
      scratch.loc_face_u.resize( nfqp );
      scratch.phase.get_function_face_values( scratch.face_phase );
      scratch.phase.get_function_face_gradients( scratch.face_grad_phase );
      scratch.oldphase.get_function_face_values( scratch.face_old_phase );
      fe_face_val[vv].get_function_values( this->oldsol, scratch.loc_face_u );
      for( unsigned q = 0; q < nfqp; ++q ){
        scratch.normal = fe_face_val.normal_vector(q);
        const double beta_q = slip_coeff( scratch.face_phase[q] ),
                     gamma_fs_q = gamma_fs( scratch.face_phase[q] ),
                     n_deriv = scratch.face_grad_phase[q]*scratch.normal,
                     dt = this->dt;
        scratch.face_grad_phase[q] -= n_deriv*scratch.normal;
        scratch.loc_face_u[q] -= ( scratch.loc_face_u[q] * scratch.normal )*scratch.normal;
        for( unsigned i=0; i<data.dpc; ++i ){
          scratch.face_base[i] = fe_face_val[vv].value( i, q );
          scratch.face_base[i] -= ( scratch.face_base[i]*scratch.normal )*scratch.normal;
        }
        for( unsigned i=0; i<data.dpc; ++i ){
          data.loc_rhs(i) -= alpha*fe_face_val.JxW(q)*(
                                ( scratch.face_phase[q] - scratch.face_old_phase[q] )/dt
                                + scratch.loc_face_u[q] * scratch.face_grad_phase[q]
                            )*( scratch.face_grad_phase[q] * scratch.face_base[i] );
          for ( unsigned j = 0; j < data.dpc; ++j )
            data.loc_m( i, j ) += fe_face_val.JxW(q)*beta_q*scratch.face_base[j]*scratch.face_base[i];
        }
      }
    }
  }
  data.ldi.resize( data.dpc );
  Its(0)->get_dof_indices( data.ldi );
}

/**
  PreRefinementPreffix
    * Copy everything to the transfer vectors
**/
template<int dim> void Velocity<dim>::PreRefinementPreffix(){
  this->x_sol[0] = this->sol;
  this->x_sol[1] = oldsol;
}

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

//////////////////////////////

template class Pressure<DIM>;
template class Velocity<DIM>;
