/**
 * @file   solver_dirichlet.cpp
 * @author Julian Iseringhausen <isering@ins.uni-bonn.de>
 * @date   Mon Apr 18 13:11:32 2011
 *
 * @brief  VMG::SolverDirichlet
 *
 */

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

#include <iostream>

#include "base/discretization.hpp"
#include "base/stencil.hpp"
#include "grid/multigrid.hpp"
#include "solver/solver_dirichlet.hpp"
#include "mg.hpp"

using namespace VMG;

// TODO: Implement global communication here
// TODO: Implement this more efficiently
// TODO: This has to be reviewed for parallelization

void SolverDirichlet::AssembleMatrix(const Grid& rhs)
{
  Index i;
  int mat_index, mat_index2;
  vmg_float prefactor_inv = 1.0 / MG::GetDiscretization()->OperatorPrefactor(rhs);
  const Stencil& A = MG::GetDiscretization()->GetStencil();

#ifdef DEBUG
  rhs.IsConsistent();
#endif

  this->Realloc(rhs.Global().Size().Product());

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

	mat_index = rhs.GlobalLinearIndex(i + rhs.Global().Begin());

	assert(mat_index >= 0 && mat_index<this->Size());

	this->Sol(mat_index) = 0.0;
	this->Rhs(mat_index) = prefactor_inv * rhs.GetVal(i);

	for (int l=0; l<this->Size(); l++)
	  this->Mat(mat_index, l) = 0.0;

	this->Mat(mat_index, mat_index) = A.GetDiag();

	for (Stencil::iterator iter = A.begin(); iter != A.end(); iter++) {
	  mat_index2 = rhs.GlobalLinearIndex(i + rhs.Global().Begin() + iter);
	  assert(mat_index2 >= 0 && mat_index2<this->Size());
	  this->Mat(mat_index, mat_index2) += iter->val;
	}

      }

  for (i.X()=rhs.Local().BoundaryBegin1().X(); i.X()<rhs.Local().BoundaryEnd1().X(); i.X()++)
    for (i.Y()=rhs.Local().Begin().Y(); i.Y()<rhs.Local().End().Y(); i.Y()++)
      for (i.Z()=rhs.Local().Begin().Z(); i.Z()<rhs.Local().End().Z(); i.Z()++) {

	mat_index = rhs.GlobalLinearIndex(i + rhs.Global().Begin());

	assert(mat_index >= 0 && mat_index<this->Size());

	this->Sol(mat_index) = this->Rhs(mat_index) = rhs.GetVal(i);

	for (int l=0; l<this->Size(); l++)
	  this->Mat(mat_index, l) = 0.0;

	this->Mat(mat_index, mat_index) = 1.0;

      }

  for (i.X()=rhs.Local().BoundaryBegin2().X(); i.X()<rhs.Local().BoundaryEnd2().X(); i.X()++)
    for (i.Y()=rhs.Local().Begin().Y(); i.Y()<rhs.Local().End().Y(); i.Y()++)
      for (i.Z()=rhs.Local().Begin().Z(); i.Z()<rhs.Local().End().Z(); i.Z()++) {

	mat_index = rhs.GlobalLinearIndex(i + rhs.Global().Begin());

	assert(mat_index >= 0 && mat_index<this->Size());

	this->Sol(mat_index) = this->Rhs(mat_index) = rhs.GetVal(i);

	for (int l=0; l<this->Size(); l++)
	  this->Mat(mat_index, l) = 0.0;

	this->Mat(mat_index, mat_index) = 1.0;

      }

  for (i.X()=0; i.X()<rhs.Local().SizeTotal().X(); i.X()++)
    for (i.Y()=rhs.Local().BoundaryBegin1().Y(); i.Y()<rhs.Local().BoundaryEnd1().Y(); i.Y()++)
      for (i.Z()=rhs.Local().Begin().Z(); i.Z()<rhs.Local().End().Z(); i.Z()++) {

	mat_index = rhs.GlobalLinearIndex(i + rhs.Global().Begin());

	assert(mat_index >= 0 && mat_index<this->Size());

	this->Sol(mat_index) = this->Rhs(mat_index) = rhs.GetVal(i);

	for (int l=0; l<this->Size(); l++)
	  this->Mat(mat_index, l) = 0.0;

	this->Mat(mat_index, mat_index) = 1.0;

      }

  for (i.X()=0; i.X()<rhs.Local().SizeTotal().X(); i.X()++)
    for (i.Y()=rhs.Local().BoundaryBegin2().Y(); i.Y()<rhs.Local().BoundaryEnd2().Y(); i.Y()++)
      for (i.Z()=rhs.Local().Begin().Z(); i.Z()<rhs.Local().End().Z(); i.Z()++) {

	mat_index = rhs.GlobalLinearIndex(i + rhs.Global().Begin());

	assert(mat_index >= 0 && mat_index<this->Size());

	this->Sol(mat_index) = this->Rhs(mat_index) = rhs.GetVal(i);

	for (int l=0; l<this->Size(); l++)
	  this->Mat(mat_index, l) = 0.0;

	this->Mat(mat_index, mat_index) = 1.0;

      }

  for (i.X()=0; i.X()<rhs.Local().SizeTotal().X(); i.X()++)
    for (i.Y()=0; i.Y()<rhs.Local().SizeTotal().Y(); i.Y()++)
      for (i.Z()=rhs.Local().BoundaryBegin1().Z(); i.Z()<rhs.Local().BoundaryEnd1().Z(); i.Z()++) {

	mat_index = rhs.GlobalLinearIndex(i + rhs.Global().Begin());

	assert(mat_index >= 0 && mat_index<this->Size());

	this->Sol(mat_index) = this->Rhs(mat_index) = rhs.GetVal(i);

	for (int l=0; l<this->Size(); l++)
	  this->Mat(mat_index, l) = 0.0;

	this->Mat(mat_index, mat_index) = 1.0;

      }

  for (i.X()=0; i.X()<rhs.Local().SizeTotal().X(); i.X()++)
    for (i.Y()=0; i.Y()<rhs.Local().SizeTotal().Y(); i.Y()++)
      for (i.Z()=rhs.Local().BoundaryBegin2().Z(); i.Z()<rhs.Local().BoundaryEnd2().Z(); i.Z()++) {

	mat_index = rhs.GlobalLinearIndex(i + rhs.Global().Begin());

	assert(mat_index >= 0 && mat_index<this->Size());

	this->Sol(mat_index) = this->Rhs(mat_index) = rhs.GetVal(i);

	for (int l=0; l<this->Size(); l++)
	  this->Mat(mat_index, l) = 0.0;

	this->Mat(mat_index, mat_index) = 1.0;

      }
}

void SolverDirichlet::ExportSol(Grid& sol, Grid& rhs)
{
  int index;

  for (int i=0; i<sol.Local().SizeTotal().X(); i++)
    for (int j=0; j<sol.Local().SizeTotal().Y(); j++)
      for (int k=0; k<sol.Local().SizeTotal().Z(); k++) {

      	index = sol.GlobalLinearIndex(sol.Global().Begin().X() + i,
				      sol.Global().Begin().Y() + j,
				      sol.Global().Begin().Z() + k);
      	sol(i,j,k) = this->Sol(index);

      }

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