/**
 * @file   grid.cpp
 * @author Julian Iseringhausen <isering@ins.uni-bonn.de>
 * @date   Mon Apr 18 12:53:27 2011
 *
 * @brief  VMG::Grid
 *
 */

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

#include <cstdio>
#include <limits>

#include "base/helper.hpp"
#include "comm/comm.hpp"
#include "grid/grid.hpp"
#include "grid/tempgrid.hpp"
#include "mg.hpp"

using namespace VMG;

vmg_float Grid::correction;

Grid::Grid(const Comm* comm, const Index& size, const vmg_float& width)
{
  grid = new vmg_float[size.Product()];

  global.Size() = size;
  global.Begin() = 0;
  global.End() = size;

  local.SizeTotal() = size;
  local.HaloBegin1() = 0;
  local.HaloBegin2() = 0;
  local.HaloEnd1() = 0;
  local.HaloEnd2() = 0;

  extent.MeshWidth() = width;

  if (comm->BoundaryConditions() == Dirichlet || comm->BoundaryConditions() == Quasiperiodic) {

    local.Size() = size - 2;
    local.Begin() = 1;
    local.End() = size-1;
    local.BoundaryBegin1() = 0;
    local.BoundaryEnd1() = 1;
    local.BoundaryBegin2() = size-1;
    local.BoundaryEnd2() = size;

  }else if (comm->BoundaryConditions() == Periodic) {

    local.Size() = size;
    local.Begin() = 0;
    local.End() = 0;
    local.BoundaryBegin1() = 0;
    local.BoundaryEnd1() = 0;
    local.BoundaryBegin2() = 0;
    local.BoundaryEnd2() = 0;

  }
}

void Grid::InitGrid()
{
  grid = new vmg_float[local.SizeTotal().Product()];
}

void Grid::Clear()
{
  for (int i=0; i<local.SizeTotal().Product(); ++i)
    grid[i] = 0.0;
}

void Grid::ClearInner()
{
  Index i;

  for (i.X()=local.Begin().X(); i.X()<local.End().X(); ++i.X())
    for (i.Y()=local.Begin().Y(); i.Y()<local.End().Y(); ++i.Y())
      for (i.Z()=local.Begin().Z(); i.Z()<local.End().Z(); ++i.Z())
	(*this)(i) = 0.0;
}

void Grid::ClearHalo()
{
  for (int i=Local().HaloBegin1().X(); i<Local().HaloEnd1().X(); i++)
    for (int j=0; j<Local().SizeTotal().Y(); j++)
      for (int k=0; k<Local().SizeTotal().Z(); k++)
	(*this)(i,j,k) = 0.0;

  for (int i=Local().HaloBegin2().X(); i<Local().HaloEnd2().X(); i++)
    for (int j=0; j<Local().SizeTotal().Y(); j++)
      for (int k=0; k<Local().SizeTotal().Z(); k++)
	(*this)(i,j,k) = 0.0;

  for (int i=0; i<Local().SizeTotal().X(); i++)
    for (int j=Local().HaloBegin1().Y(); j<Local().HaloEnd1().Y(); j++)
      for (int k=0; k<Local().SizeTotal().Z(); k++)
	(*this)(i,j,k) = 0.0;

  for (int i=0; i<Local().SizeTotal().X(); i++)
    for (int j=Local().HaloBegin2().Y(); j<Local().HaloEnd2().Y(); j++)
      for (int k=0; k<Local().SizeTotal().Z(); k++)
	(*this)(i,j,k) = 0.0;

  for (int i=0; i<Local().SizeTotal().X(); i++)
    for (int j=0; j<Local().SizeTotal().Y(); j++)
      for (int k=Local().HaloBegin1().Z(); k<Local().HaloEnd1().Z(); k++)
	(*this)(i,j,k) = 0.0;

  for (int i=0; i<Local().SizeTotal().X(); i++)
    for (int j=0; j<Local().SizeTotal().Y(); j++)
      for (int k=Local().HaloBegin2().Z(); k<Local().HaloEnd2().Z(); k++)
	(*this)(i,j,k) = 0.0;
}

void Grid::ClearBoundary()
{
  for (int i=Local().BoundaryBegin1().X(); i<Local().BoundaryEnd1().X(); i++)
    for (int j=Local().Begin().Y(); j<Local().End().Y(); j++)
      for (int k=Local().Begin().Z(); k<Local().End().Z(); k++)
	(*this)(i,j,k) = 0.0;

  for (int i=Local().BoundaryBegin2().X(); i<Local().BoundaryEnd2().X(); i++)
    for (int j=Local().Begin().Y(); j<Local().End().Y(); j++)
      for (int k=Local().Begin().Z(); k<Local().End().Z(); k++)
	(*this)(i,j,k) = 0.0;

  for (int i=0; i<Local().SizeTotal().X(); i++)
    for (int j=Local().BoundaryBegin1().Y(); j<Local().BoundaryEnd1().Y(); j++)
      for (int k=Local().Begin().Z(); k<Local().End().Z(); k++)
	(*this)(i,j,k) = 0.0;

  for (int i=0; i<Local().SizeTotal().X(); i++)
    for (int j=Local().BoundaryBegin2().Y(); j<Local().BoundaryEnd2().Y(); j++)
      for (int k=Local().Begin().Z(); k<Local().End().Z(); k++)
	(*this)(i,j,k) = 0.0;

  for (int i=0; i<Local().SizeTotal().X(); i++)
    for (int j=0; j<Local().SizeTotal().Y(); j++)
      for (int k=Local().BoundaryBegin1().Z(); k<Local().BoundaryEnd1().Z(); k++)
	(*this)(i,j,k) = 0.0;

  for (int i=0; i<Local().SizeTotal().X(); i++)
    for (int j=0; j<Local().SizeTotal().Y(); j++)
      for (int k=Local().BoundaryBegin2().Z(); k<Local().BoundaryEnd2().Z(); k++)
	(*this)(i,j,k) = 0.0;
}

void Grid::SetAverageToZero()
{
  vmg_float avg = 0.0;

  for (int i=Local().Begin().X(); i<Local().End().X(); i++)
    for (int j=Local().Begin().Y(); j<Local().End().Y(); j++)
      for (int k=Local().Begin().Z(); k<Local().End().Z(); k++)
	avg += GetVal(i,j,k);

  avg /= Local().Size().Product();

  for (int i=Local().Begin().X(); i<Local().End().X(); i++)
    for (int j=Local().Begin().Y(); j<Local().End().Y(); j++)
      for (int k=Local().Begin().Z(); k<Local().End().Z(); k++)
	(*this)(i,j,k) -= avg;
}

//TODO: Put a global communication step her
void Grid::ForceDiscreteCompatibilityCondition()
{
  vmg_float val = 0.0;

  for (int i=Local().Begin().X(); i<Local().End().X(); i++)
    for (int j=Local().Begin().Y(); j<Local().End().Y(); j++)
      for (int k=Local().Begin().Z(); k<Local().End().Z(); k++)
	val += GetVal(i,j,k);

  val = MG::GetComm()->AllReduce(val);

  if (fabs(val) <= Global().Size().Product() * std::numeric_limits<vmg_float>::epsilon()) {
#ifdef DEBUGVERBOSE
    printf("Multigrid: Right hand side satisfies compatibility condition.\n");
#endif
  }else {
    printf("Multigrid: Right hand side does NOT satisfy the compatibility condition.\n\tTotal sum: %e\nForcing compatibility condition now\n", val);

    val *= MeshWidth() * MeshWidth() * MeshWidth();

    for (int i=Local().Begin().X(); i<Local().End().X(); i++)
      for (int j=Local().Begin().Y(); j<Local().End().Y(); j++)
	for (int k=Local().Begin().Z(); k<Local().End().Z(); k++)
	  (*this)(i,j,k) -= val;
  }
}

void Grid::SetGrid(const Grid& rhs)
{
#ifdef DEBUG
  IsCompatible(rhs);
#endif

  for (int i=0; i<Local().SizeTotal().X(); i++)
    for (int j=0; j<Local().SizeTotal().Y(); j++)
      for (int k=0; k<Local().SizeTotal().Z(); k++)
	(*this)(i,j,k) = rhs.GetVal(i,j,k);
}

void Grid::SetBoundary(const Grid& rhs)
{
#ifdef DEBUG
  IsCompatible(rhs);
#endif

  for (int i=Local().BoundaryBegin1().X(); i<Local().BoundaryEnd1().X(); i++)
    for (int j=Local().Begin().Y(); j<Local().End().Y(); j++)
      for (int k=Local().Begin().Z(); k<Local().End().Z(); k++)
	(*this)(i,j,k) = rhs.GetVal(i,j,k);

  for (int i=Local().BoundaryBegin2().X(); i<Local().BoundaryEnd2().X(); i++)
    for (int j=Local().Begin().Y(); j<Local().End().Y(); j++)
      for (int k=Local().Begin().Z(); k<Local().End().Z(); k++)
	(*this)(i,j,k) = rhs.GetVal(i,j,k);

  for (int i=0; i<Local().SizeTotal().X(); i++)
    for (int j=Local().BoundaryBegin1().Y(); j<Local().BoundaryEnd1().Y(); j++)
      for (int k=Local().Begin().Z(); k<Local().End().Z(); k++)
	(*this)(i,j,k) = rhs.GetVal(i,j,k);

  for (int i=0; i<Local().SizeTotal().X(); i++)
    for (int j=Local().BoundaryBegin2().Y(); j<Local().BoundaryEnd2().Y(); j++)
      for (int k=Local().Begin().Z(); k<Local().End().Z(); k++)
	(*this)(i,j,k) = rhs.GetVal(i,j,k);

  for (int i=0; i<Local().SizeTotal().X(); i++)
    for (int j=0; j<Local().SizeTotal().Y(); j++)
      for (int k=Local().BoundaryBegin1().Z(); k<Local().BoundaryEnd1().Z(); k++)
	(*this)(i,j,k) = rhs.GetVal(i,j,k);

  for (int i=0; i<Local().SizeTotal().X(); i++)
    for (int j=0; j<Local().SizeTotal().Y(); j++)
      for (int k=Local().BoundaryBegin2().Z(); k<Local().BoundaryEnd2().Z(); k++)
	(*this)(i,j,k) = rhs.GetVal(i,j,k);
}

void Grid::AddGrid(const Grid& rhs)
{
#ifdef DEBUG
  IsCompatible(rhs);
#endif

  for (int i=0; i<Local().SizeTotal().X(); i++)
    for (int j=0; j<Local().SizeTotal().Y(); j++)
      for (int k=0; k<Local().SizeTotal().Z(); k++)
	(*this)(i,j,k) += rhs.GetVal(i,j,k);
}

void Grid::SubtractGrid(const Grid& rhs)
{
#ifdef DEBUG
  IsCompatible(rhs);
#endif

  for (int i=0; i<Local().SizeTotal().X(); i++)
    for (int j=0; j<Local().SizeTotal().Y(); j++)
      for (int k=0; k<Local().SizeTotal().Z(); k++)
	(*this)(i,j,k) -= rhs.GetVal(i,j,k);
}

void Grid::MultiplyScalar(const vmg_float& scalar)
{
  for (int i=0; i<Local().SizeTotal().X(); i++)
    for (int j=0; j<Local().SizeTotal().Y(); j++)
      for (int k=0; k<Local().SizeTotal().Z(); k++)
	(*this)(i,j,k) *= scalar;
}

void Grid::ApplyStencil(const Stencil& stencil)
{
  TempGrid *temp = MG::GetTempGrid();
  Stencil::iterator iter;

  temp->SetProperties(*this);
  temp->SetGrid(*this);
  MG::GetComm()->CommToGhosts(*temp);

  for (int i=Local().Begin().X(); i<Local().End().X(); i++)
    for (int j=Local().Begin().Y(); j<Local().End().Y(); j++)
      for (int k=Local().Begin().Z(); k<Local().End().Z(); k++) {

	(*this)(i,j,k) = stencil.GetDiag() * temp->GetVal(i,j,k);

	for (iter=stencil.begin(); iter!=stencil.end(); ++iter)
	  (*this)(i,j,k) += iter->val * temp->GetVal(i+iter->m, j+iter->n, k+iter->o);
      }
}

bool Grid::IsCompatible(const Grid& rhs) const
{
  bool eq = true;

  eq &= Helper::IsEq(Local().Begin(), rhs.Local().Begin(), "Local().Begin");
  eq &= Helper::IsEq(Local().End(), rhs.Local().End(), "Local().End");
  eq &= Helper::IsEq(Local().Size(), rhs.Local().Size(), "Local().Size");
  eq &= Helper::IsEq(Local().SizeTotal(), rhs.Local().SizeTotal(), "Local().SizeTotal");

  eq &= Helper::IsEq(Global().Begin(), rhs.Global().Begin(), "Global().Begin");
  eq &= Helper::IsEq(Global().End(), rhs.Global().End(), "Global().End");
  eq &= Helper::IsEq(Global().Size(), rhs.Global().Size(), "Global().Size");

  eq &= Helper::IsEq(Local().HaloBegin1(), rhs.Local().HaloBegin1(), "Local().HaloBegin1");
  eq &= Helper::IsEq(Local().HaloBegin2(), rhs.Local().HaloBegin2(), "Local().HaloBegin2");
  eq &= Helper::IsEq(Local().HaloEnd1(), rhs.Local().HaloEnd1(), "Local().HaloEnd1");
  eq &= Helper::IsEq(Local().HaloEnd2(), rhs.Local().HaloEnd2(), "Local().HaloEnd2");

  eq &= Helper::IsEq(MeshWidth(), rhs.MeshWidth(), "MeshWidth");

  return eq;
}

bool Grid::IsConsistent() const
{
  bool consistent = true;

  for (int i=0; i<Local().SizeTotal().X(); i++)
    for (int j=0; j<Local().SizeTotal().Y(); j++)
      for (int k=0; k<Local().SizeTotal().Z(); k++)
	consistent &= Helper::CheckNumber(GetVal(i,j,k));

  return consistent;
}

Grid::~Grid()
{
  delete [] grid;
}