/**
 * @file   solver_periodic.cpp
 * @author Julian Iseringhausen <isering@ins.uni-bonn.de>
 * @date   Mon Apr 18 13:12:02 2011
 *
 * @brief  VMG::SolverPeriodic
 *
 */

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

#include <cmath>
#include <cassert>
#include <iostream>
#include <limits>

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

using namespace VMG;

//TODO: Implement global MPI communication here

void SolverPeriodic::AssembleMatrix(const Grid& rhs)
{
  int index, index2, ind_x, ind_y, ind_z;
  vmg_float row_sum;

  const vmg_float prefactor = MG::GetDiscretization()->OperatorPrefactor(rhs);
  const Stencil& A = MG::GetDiscretization()->GetStencil();

  // Make sure that arrays are big enough to hold expanded system of equations
  this->Realloc(rhs.Global().Size().Product() + 1);

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

      	// Compute 1-dimensional index from 3-dimensional grid
      	index = rhs.GlobalLinearIndex(rhs.Global().Begin().X()+i, rhs.Global().Begin().Y()+j, rhs.Global().Begin().Z()+k);

      	// Check if we computed the index correctly
      	assert(index >= 0 && index < this->Size()-1);

      	// Since loop goes from 0 to Local().Size, we have to add Local().Begin
      	ind_x = rhs.Local().Begin().X() + i;
      	ind_y = rhs.Local().Begin().Y() + j;
      	ind_z = rhs.Local().Begin().Z() + k;

      	// Set solution and right hand side vectors
      	this->Sol(index) = 0.0;
      	this->Rhs(index) = rhs.GetVal(ind_x, ind_y, ind_z);

      	// Initialize matrix with zeros and then set entries according to the stencil
      	for (int l=0; l<this->Size(); l++)
	  this->Mat(index,l) = 0.0;

      	this->Mat(index,index) = prefactor * A.GetDiag();

      	for (Stencil::iterator iter = A.begin(); iter != A.end(); iter++) {

	  // Compute global 3-dimensional indices
	  ind_x = rhs.Global().Begin().X() + i + iter->m;
	  ind_y = rhs.Global().Begin().Y() + j + iter->n;
	  ind_z = rhs.Global().Begin().Z() + k + iter->o;

	  // Wrap indices around
	  if (MG::GetComm()->BoundaryConditions() == Periodic) {
	    if (ind_x < 0) ind_x += rhs.Global().Size().X();
	    else if (ind_x >= rhs.Global().Size().X()) ind_x -= rhs.Global().Size().X();

	    if (ind_y < 0) ind_y += rhs.Global().Size().Y();
	    else if (ind_y >= rhs.Global().Size().Y()) ind_y -= rhs.Global().Size().Y();

	    if (ind_z < 0) ind_z += rhs.Global().Size().Z();
	    else if (ind_z >= rhs.Global().Size().Z()) ind_z -= rhs.Global().Size().Z();
	  }

	  // Compute global 1-dimensional index
	  index2 = rhs.GlobalLinearIndex(ind_x, ind_y, ind_z);

	  // Set matrix entry
	  this->Mat(index,index2) += prefactor * iter->val;
      	}
      }

  // Check if matrix has zero row sum (i.e. (1,1,...,1) is an Eigenvector to the Eigenvalue 0)
  row_sum = A.GetDiag();
  for (Stencil::iterator iter=A.begin(); iter!=A.end(); iter++)
    row_sum += iter->val;

  if (fabs(row_sum) <= std::numeric_limits<vmg_float>::epsilon()) {

    // Expand equation system in order to make the system regular.
    // The last entry of the solution vector will hold a correction to the right hand side,
    // ensuring that the discrete compatibility condition holds. (Compare Trottenberg)
    for (int i=0; i<this->Size()-1; i++)
      this->Mat(this->Size()-1, i) = this->Mat(i, this->Size()-1) = 1.0;

    this->Mat(this->Size()-1, this->Size()-1) = 0.0;
    this->Sol(this->Size()-1) = 0.0;
    this->Rhs(this->Size()-1) = 0.0;

  }else {
    //TODO: Implement this
    assert(0 == "At the first glance your stencil does not seem to be singular. Try SolverRegular instead.");
  }
}

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

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

      	// Compute global 1-dimensional index
      	index = sol.GlobalLinearIndex(sol.Global().Begin().X()+i, sol.Global().Begin().Y()+j, sol.Global().Begin().Z()+k);

      	// Set solution
      	sol(sol.Local().Begin().X()+i, sol.Local().Begin().Y()+j, sol.Local().Begin().Z()+k) = this->Sol(index);

      }

  // Set correction
  Grid::Correction() = this->Sol(this->Size()-1);
}