/**
 * @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/grid_iterator.hpp"
#include "grid/multigrid.hpp"
#include "grid/tempgrid.hpp"
#include "level/level_operator_cs.hpp"
#include "mg.hpp"

using namespace VMG;

void LevelOperatorCS::Restrict(Grid& sol_f, Grid& rhs_f, Grid& sol_c, Grid& rhs_c)
{
#ifdef DEBUG_MATRIX_CHECKS
  sol_f.IsConsistent();
  rhs_f.IsConsistent();
#endif

#ifdef DEBUG
  sol_f.IsCompatible(rhs_f);
#endif

  Grid::iterator iter_f, iter_c;
  Stencil::iterator stencil_iter;
  vmg_float res;

  const Stencil& op = OperatorRestrict();

  Index begin_f = rhs_f.Local().Begin() + 2*rhs_c.Global().BeginLocal() - rhs_f.Global().BeginLocal();
  Index end_f = rhs_f.Local().End();

  /* Modify fine begin/end to align the points on both levels correctly */
  if (rhs_c.Global().BoundaryType() == GlobalCoarsened)
    for (int i=0; i<3; ++i) {

      if (rhs_c.Local().HasBoundary1()[i])
	++begin_f[i];

      if (rhs_c.Local().HasBoundary2()[i])
	--end_f[i];

    }

  const GridIteratorSet bounds_f(begin_f, end_f);
  const GridIteratorSet bounds_c(rhs_c.Local().AlignmentBegin(), rhs_c.Local().AlignmentEnd());

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

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

  for (iter_f=bounds_f.Begin(), iter_c=bounds_c.Begin(); iter_c!=bounds_c.End(); iter_f+=2, ++iter_c) {

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

    for (stencil_iter=op.begin(); stencil_iter!=op.end(); ++stencil_iter)
      res += stencil_iter->Val() * temp->GetVal(*iter_f + stencil_iter->Disp());

    rhs_c(*iter_c) = res;

  }

#ifdef DEBUG_MATRIX_CHECKS
  rhs_c.IsConsistent();
#endif
}

void LevelOperatorCS::Prolongate(Grid& sol_f, Grid& rhs_f, Grid& sol_c, Grid& rhs_c)
{
#ifdef DEBUG_MATRIX_CHECKS
  sol_c.IsConsistent();
  rhs_c.IsConsistent();
  sol_f.IsConsistent();
  rhs_f.IsConsistent();
#endif

#ifdef DEBUG
  sol_c.IsCompatible(rhs_c);
  sol_f.IsCompatible(rhs_f);
#endif

  Grid::iterator iter_f, iter_c;
  Stencil::iterator stencil_iter;
  vmg_float val;

  const Stencil& op = OperatorProlongate();
  Comm* comm = MG::GetComm();

  Index begin_f = sol_f.Local().Begin() + 2*sol_c.Global().BeginLocal() - sol_f.Global().BeginLocal();
  Index end_f = sol_f.Local().End();

  if (sol_c.Global().BoundaryType() == GlobalCoarsened)
    for (int i=0; i<3; ++i) {

      if (rhs_f.Local().HasBoundary1()[i])
	++begin_f[i];

      if (rhs_f.Local().HasBoundary2()[i])
	--end_f[i];

    }

  const GridIteratorSet bounds_f(begin_f, end_f);
  const GridIteratorSet bounds_c(sol_c.Local().AlignmentBegin(), sol_c.Local().AlignmentEnd());

  Index apply_stencil_begin, apply_stencil_end;

  for (int i=0; i<3; ++i) {
    apply_stencil_begin[i] = (sol_f.Local().HasBoundary1()[i] ? 1 : 0);
    apply_stencil_end[i] = (sol_f.Local().HasBoundary2()[i] ? sol_f.Local().End()[i] : sol_f.Local().SizeTotal()[i]);
  }

  sol_f.ClearHalo();

  for (iter_f=bounds_f.Begin(), iter_c=bounds_c.Begin(); iter_c!=bounds_c.End(); iter_f+=2, ++iter_c) {

    val = sol_c.GetVal(*iter_c);

    sol_f(*iter_f) += op.GetDiag() * val;

    for (stencil_iter = op.begin(); stencil_iter != op.end(); ++stencil_iter) {

      sol_f(*iter_f + stencil_iter->Disp()) += stencil_iter->Val() * val;

    }

  }

  comm->CommFromGhosts(sol_f);
}