/**
 * @file   level_operator_cs.cpp
 * @author Julian Iseringhausen <isering@ins.uni-bonn.de>
 * @date   Mon Apr 18 12:59:46 2011
 *
 * @brief  VMG::LevelOperatorCS
 *
 */

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include "base/index.hpp"
#include "comm/comm.hpp"
#include "grid/multigrid.hpp"
#include "grid/tempgrid.hpp"
#include "level/level_operator_cs.hpp"
#include "mg.hpp"

using namespace VMG;

void LevelOperatorCS::Restrict(Multigrid& sol, Multigrid& rhs)
{
#ifdef DEBUG
  sol().IsConsistent();
  rhs().IsConsistent();
  sol().IsCompatible(rhs());
#endif

  Index i_c, i_f;
  Stencil::iterator iter;
  vmg_float res;

  const int l_f = rhs.Level();
  const int l_c = l_f - 1;

  Index begin_f = rhs().Local().Begin();
  const Index end_f = rhs().Local().End() - 1;
  const Index begin_c = rhs(l_c).Local().AlignmentBegin();
  const Stencil& op = OperatorRestrict();

  if (sol(l_c).Global().BoundaryType() == GlobalCoarsened &&
     (MG::GetComm()->BoundaryConditions() == Dirichlet || MG::GetComm()->BoundaryConditions() == Quasiperiodic))
    begin_f += 1;

  // Communication step
  MG::GetComm()->CommToGhosts(sol());

  // Compute residual
  TempGrid *temp = MG::GetTempGrid();
  temp->SetProperties(rhs());
  temp->ImportFromResidual(sol(), rhs());

  // Bring grids to the next coarser level
  rhs.ToCoarserLevel();
  sol.ToCoarserLevel();

#ifdef DEBUG
  sol().IsConsistent();
  rhs().IsConsistent();
  sol().IsCompatible(rhs());
#endif

  if (rhs().Global().BoundaryType() == GlobalCoarsened)
    rhs().ClearBoundary();

  // Set coarse grid values according to the given stencil
  for (i_c.X()=begin_c.X(), i_f.X()=begin_f.X(); i_f.X()<end_f.X(); ++i_c.X(), i_f.X()+=2)
    for (i_c.Y()=begin_c.Y(), i_f.Y()=begin_f.Y(); i_f.Y()<end_f.Y(); ++i_c.Y(), i_f.Y()+=2)
      for (i_c.Z()=begin_c.Z(), i_f.Z()=begin_f.Z(); i_f.Z()<end_f.Z(); ++i_c.Z(), i_f.Z()+=2) {

	res = op.GetDiag() * temp->GetVal(i_f);

	for (iter=op.begin(); iter!=op.end(); ++iter)
	  res += iter->val * temp->GetVal(i_f + iter);

	rhs()(i_c) = res;

      }
}

void LevelOperatorCS::Prolongate(Multigrid& sol, Multigrid& rhs)
{
#ifdef DEBUG
  sol().IsConsistent();
  rhs().IsConsistent();
  sol().IsCompatible(rhs());
#endif

  Stencil::iterator iter;
  Index i_c, i_f;
  vmg_float val;

  const Stencil& op = OperatorProlongate();

  const int l_c = sol.Level();
  const int l_f = l_c + 1;

  const Index begin_c = sol().Local().AlignmentBegin();
  Index begin_f = sol(l_f).Local().Begin();
  const Index end_f = sol(l_f).Local().End();

  if (sol(l_c).Global().BoundaryType() == GlobalCoarsened &&
      (MG::GetComm()->BoundaryConditions() == Dirichlet || MG::GetComm()->BoundaryConditions() == Quasiperiodic))
    begin_f += 1;

  sol.ToFinerLevel();
  rhs.ToFinerLevel();

#ifdef DEBUG
  sol().IsConsistent();
  rhs().IsConsistent();
  sol().IsCompatible(rhs());
#endif

  if (MG::GetComm()->BoundaryConditions() == Periodic) {

    sol().ClearHalo();

    for (i_c.X()=begin_c.X(), i_f.X()=begin_f.X(); i_f.X()<end_f.X(); i_c.X()++, i_f.X()+=2)
      for (i_c.Y()=begin_c.Y(), i_f.Y()=begin_f.Y(); i_f.Y()<end_f.Y(); i_c.Y()++, i_f.Y()+=2)
	for (i_c.Z()=begin_c.Z(), i_f.Z()=begin_f.Z(); i_f.Z()<end_f.Z(); i_c.Z()++, i_f.Z()+=2) {

	  val = sol(l_c).GetVal(i_c);

	  sol()(i_f) += op.GetDiag() * val;

	  for (iter=op.begin(); iter!=op.end(); ++iter)
	    sol()(i_f + iter) += iter->val * val;

	}

    MG::GetComm()->CommFromGhosts(sol());

  }else if (MG::GetComm()->BoundaryConditions() == Dirichlet ||
	    MG::GetComm()->BoundaryConditions() == Quasiperiodic) {

    for (i_c.X()=begin_c.X(), i_f.X()=begin_f.X(); i_f.X()<end_f.X(); ++i_c.X(), i_f.X()+=2)
      for (i_c.Y()=begin_c.Y(), i_f.Y()=begin_f.Y(); i_f.Y()<end_f.Y(); ++i_c.Y(), i_f.Y()+=2)
	for (i_c.Z()=begin_c.Z(), i_f.Z()=begin_f.Z(); i_f.Z()<end_f.Z(); ++i_c.Z(), i_f.Z()+=2) {

	  val = sol(l_c).GetVal(i_c);

	  sol()(i_f) += op.GetDiag() * val;

	  for (iter=op.begin(); iter!=op.end(); ++iter)
	    if (i_f.X()+iter->m >= sol().Local().Begin().X() && i_f.X()+iter->m < sol().Local().End().X() &&
		i_f.Y()+iter->n >= sol().Local().Begin().Y() && i_f.Y()+iter->n < sol().Local().End().Y() &&
		i_f.Z()+iter->o >= sol().Local().Begin().Z() && i_f.Z()+iter->o < sol().Local().End().Z()) {
	      sol()(i_f + iter) += iter->val * val;
	    }
	}
  }
}