/**
 * @file   level_operator_fas.cpp
 * @author Julian Iseringhausen <isering@ins.uni-bonn.de>
 * @date   Mon Apr 18 13:00:23 2011
 *
 * @brief  VMG::LevelOperatorFAS
 *
 */

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

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

using namespace VMG;

void LevelOperatorFAS::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 val;

  Grid& sol_c = MG::GetComm()->GetCoarserGrid(sol);
  Grid& rhs_c = MG::GetComm()->GetCoarserGrid(rhs);

  Index begin_f = rhs().Local().Begin();
  const Index end_f = rhs().Local().End();
  const Index begin_c = rhs_c.Local().AlignmentBegin();

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

  const Stencil& op_res = OperatorRestrict();
  const Stencil& op_sol = OperatorSol();
  const Stencil& op_pde = MG::GetDiscretization()->GetStencil();

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

  MG::GetDiscretization()->SetInnerBoundary(sol(), rhs(), sol_c);

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

  VMG::TempGrid* sol_old = this->GetTempGrid(sol.Level()-1);
  sol_old->SetProperties(sol_c);
  sol_old->Clear();

  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 = op_sol.GetDiag() * sol().GetVal(i_f);

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

        sol_c(i_c) = (*sol_old)(i_c) = val;

      }

  MG::GetComm()->CommToGhosts(sol_c);

  const vmg_float prefactor = MG::GetDiscretization()->OperatorPrefactor(sol_c);

  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)
        rhs_c(i_c) = op_res.Apply(*res_grid, i_f) + prefactor * op_pde.Apply(sol_c, i_c);

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

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

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

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

  Grid& sol_f = MG::GetComm()->GetFinerGrid(sol);

  const Stencil& op = OperatorProlongate();

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

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

  VMG::TempGrid *sol_old = this->GetTempGrid(sol.Level());

  if (sol().Global().BoundaryType() == LocallyRefined)
    MG::GetDiscretization()->SetInnerBoundary(sol(), rhs(), MG::GetComm()->GetCoarserGrid(sol));

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

    sol_f.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().GetVal(i_c) - sol_old->GetVal(i_c);

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

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

	}

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

  }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().GetVal(i_c) - sol_old->GetVal(i_c);

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

          //TODO: This does not work in parallel
          for (iter=op.begin(); iter!=op.end(); ++iter)
            if (i_f.X()+iter->m >= sol_f.Local().Begin().X() && i_f.X()+iter->m < sol_f.Local().End().X() &&
                i_f.Y()+iter->n >= sol_f.Local().Begin().Y() && i_f.Y()+iter->n < sol_f.Local().End().Y() &&
                i_f.Z()+iter->o >= sol_f.Local().Begin().Z() && i_f.Z()+iter->o < sol_f.Local().End().Z()) {
              sol_f(i_f + iter) += iter->val * val;
            }
	}
  }

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

#ifdef DEBUG
  sol_f.IsConsistent();
#endif
}