//
// Software License for AMDiS
//
// Copyright (c) 2010 Dresden University of Technology 
// All rights reserved.
// Authors: Simon Vey, Thomas Witkowski et al.
//
// This file is part of AMDiS
//
// See also license.opensource.txt in the distribution.


#include "parallel/SubDomainSolver.h"
#include "SystemVector.h"

namespace AMDiS {
  
  using namespace std;


  void SubDomainSolver::setDofMapping(ParallelDofMapping *coarseDofs,
				      ParallelDofMapping *interiorDofs)
  {
    coarseSpaceMap = coarseDofs;
    interiorMap = interiorDofs;

    if (mpiCommInterior.Get_size() == 1) {
      rStartInterior = 0;
      nGlobalOverallInterior = interiorMap->getOverallDofs();
    } else {
      int groupRowsInterior = 0;
      if (mpiCommInterior.Get_rank() == 0)
	groupRowsInterior = interiorMap->getOverallDofs();

      mpi::getDofNumbering(mpiCommCoarseSpace, groupRowsInterior,
			   rStartInterior, nGlobalOverallInterior);

      int tmp = 0;
      if (mpiCommInterior.Get_rank() == 0)
	tmp = rStartInterior;

      mpiCommInterior.Allreduce(&tmp, &rStartInterior, 1, MPI_INT, MPI_SUM);
                 
      MSG("COMM TEST: %d %d %d %d %d\n",
	  mpiCommInterior.Get_size(), 
	  interiorMap->getRankDofs(),
	  interiorMap->getOverallDofs(),
	  nGlobalOverallInterior, rStartInterior);
    }
  }


  void SubDomainSolver::fillPetscMatrix(Matrix<DOFMatrix*> *mat)
  {
    FUNCNAME("SubDomainSolver::fillPetscMatrix()");

    vector<const FiniteElemSpace*> feSpaces = getFeSpaces(mat);

    int nRowsRankInterior = interiorMap->getRankDofs();
    int nRowsOverallInterior = interiorMap->getOverallDofs();
    int nRowsRankCoarse = coarseSpaceMap->getRankDofs();
    int nRowsOverallCoarse = coarseSpaceMap->getOverallDofs();

    bool multilevel = false;
    if (mpiCommInterior.Get_size() == 1) {
      nGlobalOverallInterior = nRowsOverallInterior;

      MatCreateSeqAIJ(mpiCommInterior, nRowsRankInterior, nRowsRankInterior,
		      60, PETSC_NULL, &matIntInt);
    } else {
      multilevel = true;

      MatCreateMPIAIJ(mpiCommInterior, 
		      nRowsRankInterior, nRowsRankInterior,
		      nRowsOverallInterior, nRowsOverallInterior,
		      60, PETSC_NULL, 60, PETSC_NULL, &matIntInt);
    }

    MatCreateMPIAIJ(mpiCommCoarseSpace,
		    nRowsRankCoarse, nRowsRankCoarse,
		    nRowsOverallCoarse, nRowsOverallCoarse,
		    60, PETSC_NULL, 60, PETSC_NULL, &matCoarseCoarse);

    MatCreateMPIAIJ(mpiCommCoarseSpace,
		    nRowsRankCoarse, nRowsRankInterior,
		    nRowsOverallCoarse, nGlobalOverallInterior,
		    60, PETSC_NULL, 60, PETSC_NULL, &matCoarseInt);

    MatCreateMPIAIJ(mpiCommCoarseSpace,
		    nRowsRankInterior, nRowsRankCoarse,
		    nGlobalOverallInterior, nRowsOverallCoarse,
		    60, PETSC_NULL, 60, PETSC_NULL, &matIntCoarse);

    // === Prepare traverse of sequentially created matrices. ===

    using mtl::tag::row; using mtl::tag::nz; using mtl::begin; using mtl::end;
    namespace traits = mtl::traits;
    typedef DOFMatrix::base_matrix_type Matrix;

    typedef traits::range_generator<row, Matrix>::type cursor_type;
    typedef traits::range_generator<nz, cursor_type>::type icursor_type;

    vector<int> cols, colsOther;
    vector<double> values, valuesOther;
    cols.reserve(300);
    colsOther.reserve(300);
    values.reserve(300);
    valuesOther.reserve(300);

    // === Traverse all sequentially created matrices and add the values to ===
    // === the global PETSc matrices.                                       ===

    int nComponents = mat->getSize();
    for (int i = 0; i < nComponents; i++) {
      for (int j = 0; j < nComponents; j++) {
	if (!(*mat)[i][j])
	  continue;

	traits::col<Matrix>::type col((*mat)[i][j]->getBaseMatrix());
	traits::const_value<Matrix>::type value((*mat)[i][j]->getBaseMatrix());
	
	// Traverse all rows.
	for (cursor_type cursor = begin<row>((*mat)[i][j]->getBaseMatrix()), 
	       cend = end<row>((*mat)[i][j]->getBaseMatrix()); cursor != cend; ++cursor) {

	  bool rowPrimal = isCoarseSpace(feSpaces[i], *cursor);
  
	  cols.clear();
	  colsOther.clear();
	  values.clear();	  
	  valuesOther.clear();

	  // Traverse all columns.
	  for (icursor_type icursor = begin<nz>(cursor), icend = end<nz>(cursor); 
	       icursor != icend; ++icursor) {

	    bool colPrimal = isCoarseSpace(feSpaces[j], col(*icursor));

	    if (colPrimal) {
	      if (rowPrimal) {
		cols.push_back(col(*icursor));
		values.push_back(value(*icursor));
	      } else {
		colsOther.push_back(col(*icursor));
		valuesOther.push_back(value(*icursor));
	      }
	    } else {
	      if (rowPrimal) {
		colsOther.push_back(col(*icursor));
		valuesOther.push_back(value(*icursor));
	      } else {
		cols.push_back(col(*icursor));
		values.push_back(value(*icursor));
	      }
	    }
	  }  // for each nnz in row


	  // === Set matrix values. ===

	  if (rowPrimal) {
	    int rowIndex = coarseSpaceMap->getMatIndex(i, *cursor);
	    for (unsigned int k = 0; k < cols.size(); k++)
	      cols[k] = coarseSpaceMap->getMatIndex(j, cols[k]);

	    MatSetValues(matCoarseCoarse, 1, &rowIndex, cols.size(),
			 &(cols[0]), &(values[0]), ADD_VALUES);

	    if (colsOther.size()) {
	      if (multilevel == false) {
		for (unsigned int k = 0; k < colsOther.size(); k++)
		  colsOther[k] = interiorMap->getMatIndex(j, colsOther[k]);
	      } else {
		for (unsigned int k = 0; k < colsOther.size(); k++) {
		  colsOther[k] = 
		    interiorMap->getMatIndex(j, colsOther[k]) + rStartInterior;
		  if (colsOther[k] == 50723) {
		    MSG("RRTEST A: %d %d %d\n",
			nRowsOverallInterior,
			nGlobalOverallInterior,
			rStartInterior);
			
		  }
		}
	      }
 	      
	      MatSetValues(matCoarseInt, 1, &rowIndex, colsOther.size(),
 			   &(colsOther[0]), &(valuesOther[0]), ADD_VALUES);
	    }
	  } else {
	    int localRowIndex = 
	      (multilevel == false ? interiorMap->getLocalMatIndex(i, *cursor) :
	       interiorMap->getMatIndex(i, *cursor));

	    for (unsigned int k = 0; k < cols.size(); k++) {
	      if (multilevel == false)
		cols[k] = interiorMap->getLocalMatIndex(j, cols[k]);
	      else  {
		cols[k] = interiorMap->getMatIndex(j, cols[k]);
		  if (colsOther[k] == 50723) {
		    MSG("RRTEST B: %d %d %d\n",
			nRowsOverallInterior,
			nGlobalOverallInterior,
			rStartInterior);
			
		  }

	      }
	    }
	    
  	    MatSetValues(matIntInt, 1, &localRowIndex, cols.size(),
  			 &(cols[0]), &(values[0]), ADD_VALUES);

	    if (colsOther.size()) {
	      int globalRowIndex = interiorMap->getMatIndex(i, *cursor);

	      if (multilevel == false)
		globalRowIndex += rStartInterior;

	      for (unsigned int k = 0; k < colsOther.size(); k++)
		colsOther[k] = coarseSpaceMap->getMatIndex(j, colsOther[k]);

  	      MatSetValues(matIntCoarse, 1, &globalRowIndex, colsOther.size(),
  			   &(colsOther[0]), &(valuesOther[0]), ADD_VALUES);
	    }
	  }
	} 
      }
    }

    // === Start global assembly procedure. ===

    MatAssemblyBegin(matIntInt, MAT_FINAL_ASSEMBLY);
    MatAssemblyEnd(matIntInt, MAT_FINAL_ASSEMBLY);

    MatAssemblyBegin(matCoarseCoarse, MAT_FINAL_ASSEMBLY);
    MatAssemblyEnd(matCoarseCoarse, MAT_FINAL_ASSEMBLY);

    MatAssemblyBegin(matIntCoarse, MAT_FINAL_ASSEMBLY);
    MatAssemblyEnd(matIntCoarse, MAT_FINAL_ASSEMBLY);

    MatAssemblyBegin(matCoarseInt, MAT_FINAL_ASSEMBLY);
    MatAssemblyEnd(matCoarseInt, MAT_FINAL_ASSEMBLY);


    // === Create solver for the non primal (thus local) variables. ===

    KSPCreate(mpiCommInterior, &kspInterior);
    KSPSetOperators(kspInterior, matIntInt, matIntInt, SAME_NONZERO_PATTERN);
    KSPSetOptionsPrefix(kspInterior, "interior_");
    KSPSetType(kspInterior, KSPPREONLY);
    PC pcInterior;
    KSPGetPC(kspInterior, &pcInterior);
    PCSetType(pcInterior, PCLU);
    if (multilevel == false)
      PCFactorSetMatSolverPackage(pcInterior, MATSOLVERUMFPACK);
    else
      PCFactorSetMatSolverPackage(pcInterior, MATSOLVERMUMPS);
    KSPSetFromOptions(kspInterior);
  }


  void SubDomainSolver::fillPetscRhs(SystemVector *vec)
  {
    FUNCNAME("SubDomainSolver::fillPetscRhs()");

    VecCreateMPI(mpiCommCoarseSpace, 
		 coarseSpaceMap->getRankDofs(), 
		 coarseSpaceMap->getOverallDofs(),
		 &rhsCoarseSpace);

    VecCreateMPI(mpiCommCoarseSpace, 
		 interiorMap->getRankDofs(), 
		 nGlobalOverallInterior,
		 &rhsInterior);

    for (int i = 0; i < vec->getSize(); i++) {
      const FiniteElemSpace *feSpace = vec->getDOFVector(i)->getFeSpace();
      DOFVector<double>::Iterator dofIt(vec->getDOFVector(i), USED_DOFS);
      for (dofIt.reset(); !dofIt.end(); ++dofIt) {
	int index = dofIt.getDOFIndex();
	if (isCoarseSpace(feSpace, index)) {
	  index = coarseSpaceMap->getMatIndex(i, index);
	  VecSetValue(rhsCoarseSpace, index, *dofIt, ADD_VALUES);
	} else {
	  index = interiorMap->getMatIndex(i, index) + rStartInterior;
	  VecSetValue(rhsInterior, index, *dofIt, INSERT_VALUES);
	}      
      }
    }

    VecAssemblyBegin(rhsCoarseSpace);
    VecAssemblyEnd(rhsCoarseSpace);

    VecAssemblyBegin(rhsInterior);
    VecAssemblyEnd(rhsInterior);
  }


  void SubDomainSolver::solvePetscMatrix(SystemVector &vec, AdaptInfo *adaptInfo)
  {
    FUNCNAME("SubDomainSolver::solvePetscMatrix()");
  }


  void SubDomainSolver::destroyVectorData()
  {
    FUNCNAME("SubDomainSolver::destroyVectorData()");

    VecDestroy(&rhsCoarseSpace);
    VecDestroy(&rhsInterior);
  }


  void SubDomainSolver::destroyMatrixData()
  {
    FUNCNAME("SubDomainSolver::destroyMatrixData()");

    MatDestroy(&matIntInt);
    MatDestroy(&matCoarseCoarse);
    MatDestroy(&matCoarseInt);
    MatDestroy(&matIntCoarse);

    KSPDestroy(&kspInterior);
  }


  void SubDomainSolver::solve(Vec &rhs, Vec &sol)
  {
    KSPSolve(kspInterior, rhs, sol);
  }


  void SubDomainSolver::solveGlobal(Vec &rhs, Vec &sol)
  {
    FUNCNAME("SubDomainSolver::solveGlobal()");

    int ml = 0;
    Parameters::get("parallel->multi level test", ml);

    Vec tmp;
    if (ml == 0)
      VecCreateSeq(PETSC_COMM_SELF, interiorMap->getRankDofs(), &tmp);
    else
      VecCreateMPI(mpiCommInterior,
		   interiorMap->getRankDofs(),
		   interiorMap->getOverallDofs(),
		   &tmp);

    PetscScalar *tmpValues, *rhsValues;
    VecGetArray(tmp, &tmpValues);
    VecGetArray(rhs, &rhsValues);

    for (int i = 0; i < interiorMap->getRankDofs(); i++)
      tmpValues[i] = rhsValues[i];

    VecRestoreArray(rhs, &rhsValues);
    VecRestoreArray(tmp, &tmpValues);

    KSPSolve(kspInterior, tmp, tmp);

    VecGetArray(tmp, &tmpValues);
    VecGetArray(sol, &rhsValues);

    for (int i = 0; i < interiorMap->getRankDofs(); i++) 
      rhsValues[i] = tmpValues[i];

    VecRestoreArray(sol, &rhsValues);
    VecRestoreArray(tmp, &tmpValues);

    VecDestroy(&tmp);
  }


  vector<const FiniteElemSpace*> SubDomainSolver::getFeSpaces(Matrix<DOFMatrix*> *mat)
  {
    FUNCNAME("SubDomainSolver::getFeSpaces()");

    int nComponents = mat->getNumRows();
    vector<const FiniteElemSpace*> result(nComponents);

    for (int i = 0; i < nComponents; i++) 
      for (int j = 0; j < nComponents; j++)
	if ((*mat)[i][j]) {
	  result[i] = (*mat)[i][j]->getRowFeSpace();
	  break;
	}

    return result;
  }


  vector<const FiniteElemSpace*> SubDomainSolver::getFeSpaces(SystemVector *vec)
  {
    FUNCNAME("SubDomainSolver::getFeSpaces()");

    int nComponents = vec->getSize();
    vector<const FiniteElemSpace*> result(nComponents);

    for (int i = 0; i < nComponents; i++) 
      result[i] = vec->getFeSpace(i);

    return result;
  }

}