/*
 *    vmg - a versatile multigrid solver
 *    Copyright (C) 2012 Institute for Numerical Simulation, University of Bonn
 *
 *  vmg is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  vmg is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

/**
 * @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 <libvmg_config.h>
#endif

#include <limits>

#include "base/discretization.hpp"
#include "base/helper.hpp"
#include "base/index.hpp"
#include "comm/comm.hpp"
#include "grid/grid_index_translations.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)
{
  Grid::iterator iter_f, iter_c;
  Stencil::iterator stencil_iter;

  Comm& comm = *MG::GetComm();

  Grid& sol_f = sol(sol.Level());
  Grid& rhs_f = rhs(rhs.Level());

  Grid& sol_c_dist = sol(sol.Level()-1);
  Grid& rhs_c_dist = rhs(rhs.Level()-1);

  Grid& sol_c_undist = comm.GetCoarserGrid(sol);
  Grid& rhs_c_undist = comm.GetCoarserGrid(rhs);

  GridIteratorSet bounds_c, bounds_f;
  GridIndexTranslations::GetGridAlignment(rhs_c_undist, bounds_c, rhs_f, bounds_f);

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

  // Communication step
  comm.CommToGhosts(sol_f);
  comm.CommSubgrid(sol_c_dist, sol_c_undist, 0);
  comm.CommSubgrid(rhs_c_dist, rhs_c_undist, 0);

  MG::GetDiscretization()->SetInnerBoundary(sol_f, rhs_f, sol_c_undist);

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

  for (iter_f=bounds_f.Begin(), iter_c=bounds_c.Begin(); iter_c!=bounds_c.End(); iter_f+=2, ++iter_c) {
    Helper::AssertVectorsEqual(sol_c_undist.GetSpatialPos(*iter_c), sol_f.GetSpatialPos(*iter_f), 1e6 * std::numeric_limits<vmg_float>::epsilon());
    sol_c_undist(*iter_c) = op_sol.Apply(sol_f, *iter_f);
  }

  comm.CommToGhosts(sol_c_undist);

  for (iter_f=bounds_f.Begin(), iter_c=bounds_c.Begin(); iter_c!=bounds_c.End(); iter_f+=2, ++iter_c) {
    Helper::AssertVectorsEqual(rhs_c_undist.GetSpatialPos(*iter_c), (*res_grid).GetSpatialPos(*iter_f), 1e6 * std::numeric_limits<vmg_float>::epsilon());
    rhs_c_undist(*iter_c) = op_res.Apply(*res_grid, *iter_f) + prefactor * op_pde.Apply(sol_c_undist, *iter_c);
  }

  /*
  if (rhs_c_undist.Global().BoundaryType() == GlobalCoarsened)
    rhs_c_undist.SetBoundary(sol_c_undist);
  */

  comm.CommSubgrid(sol_c_undist, sol_c_dist, 1);
  comm.CommSubgrid(rhs_c_undist, rhs_c_dist, 1);

  VMG::TempGrid* sol_old = this->GetTempGrid(sol_c_dist.Level());
  sol_old->SetProperties(sol_c_dist);
  sol_old->SetGrid(sol_c_dist);

  //  comm.CommToGhosts(rhs_c_dist);

  sol.ToCoarserLevel();
  rhs.ToCoarserLevel();
}

void LevelOperatorFAS::Prolongate(Multigrid& sol, Multigrid& rhs)
{
  Grid::iterator iter_f, iter_c;
  Stencil::iterator stencil_iter;
  vmg_float val;

  Comm& comm = *MG::GetComm();

  Grid& sol_c = sol(sol.Level());
  Grid& sol_f_dist = sol(sol.Level()+1);
  Grid& rhs_f_dist = rhs(rhs.Level()+1);
  Grid& sol_f_undist = comm.GetFinerGrid(sol);

  const Stencil& op = OperatorProlongate();

  TempGrid *sol_old = this->GetTempGrid(sol_c.Level());

  GridIteratorSet bounds_c, bounds_f;
  GridIndexTranslations::GetGridAlignment(sol_c, bounds_c, sol_f_undist, bounds_f);

  comm.CommSubgrid(sol_f_dist, sol_f_undist, 0);
  sol_f_undist.ClearHalo();

  for (iter_f=bounds_f.Begin(), iter_c=bounds_c.Begin(); iter_c!=bounds_c.End(); iter_f+=2, ++iter_c) {
    Helper::AssertVectorsEqual(sol_c.GetSpatialPos(*iter_c), sol_f_undist.GetSpatialPos(*iter_f), 1e6 * std::numeric_limits<vmg_float>::epsilon());
    val = sol_c.GetVal(*iter_c) - sol_old->GetVal(*iter_c);
    sol_f_undist(*iter_f) += op.GetDiag() * val;
    for (stencil_iter = op.begin(); stencil_iter != op.end(); ++stencil_iter)
	sol_f_undist(*iter_f + stencil_iter->Disp()) += stencil_iter->Val() * val;
  }

  comm.CommFromGhosts(sol_f_undist);
  comm.CommSubgrid(sol_f_undist, sol_f_dist, 1);

  if (sol_f_dist.Global().BoundaryType() == LocallyRefined)
    MG::GetDiscretization()->SetInnerBoundary(sol_f_dist, rhs_f_dist, sol_c);

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