//
// 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/PetscSolverNavierStokes.h"

namespace AMDiS {

  using namespace std;


  int petscMultNavierStokesSchur(Mat mat, Vec x, Vec y)
  {
    void *ctx;
    MatShellGetContext(mat, &ctx);
    MatShellNavierStokesSchur* data = static_cast<MatShellNavierStokesSchur*>(ctx);

    switch (data->solverMode) {
    case 0:
      {
	Vec vec0, vec1;
	VecDuplicate(x, &vec0);
	MatGetVecs(data->A00, &vec1, PETSC_NULL);
	
	MatMult(data->A11, x, y);
	MatMult(data->A01, x, vec1);
	KSPSolve(data->kspVelocity, vec1, vec1);
	MatMult(data->A10, vec1, y);
	VecAYPX(y, -1.0, vec0);
	
	VecDestroy(&vec0);
	VecDestroy(&vec1);
      }
      break;
    case 1:
      {
	Vec vec0, vec1;
	VecDuplicate(x, &vec0);
	VecDuplicate(x, &vec1);

	KSPSolve(data->kspLaplace, x, vec0);
	MatMult(data->matConDif, vec0, vec1);
	KSPSolve(data->kspMass, vec1, y);

	VecDestroy(&vec0);
	VecDestroy(&vec1);
      }
      break;
    default:
      ERROR_EXIT("Wrong solver mode %d\n", data->solverMode);
    }
  }
  

  PetscSolverNavierStokes::PetscSolverNavierStokes(string name)
    : PetscSolverGlobalMatrix(name),
      pressureComponent(-1),
      massMatrixSolver(NULL),
      laplaceMatrixSolver(NULL)
  {
    Parameters::get(initFileStr + "->stokes->pressure component", 
		    pressureComponent);

    TEST_EXIT(pressureComponent >= 0)
      ("For using PETSc stokes solver you must define a pressure component!\n");

    TEST_EXIT(pressureComponent == 2)("Fixed for pressure component = 2\n");
  }


  void PetscSolverNavierStokes::initSolver(KSP &ksp)
  {
    FUNCNAME("PetscSolverNavierStokes::initSolver()");

    KSPCreate(mpiCommGlobal, &ksp);
    KSPSetOperators(ksp, getMatInterior(), getMatInterior(), 
		    SAME_NONZERO_PATTERN); 
    KSPSetTolerances(ksp, 0.0, 1e-8, PETSC_DEFAULT, PETSC_DEFAULT);
    KSPSetType(ksp, KSPGMRES);
    KSPSetOptionsPrefix(ksp, "ns_");
    KSPSetFromOptions(ksp);
    KSPMonitorSet(ksp, KSPMonitorTrueResidualNorm, PETSC_NULL, PETSC_NULL);

    
    // === Create null space information. ===

    Vec nullSpaceBasis;
    VecDuplicate(getVecSolInterior(), &nullSpaceBasis);

    SystemVector basisVec("tmp", componentSpaces, componentSpaces.size(), true);
    basisVec.set(0.0);
    basisVec.getDOFVector(pressureComponent)->set(1.0);

    setDofVector(nullSpaceBasis, basisVec, true);
    VecAssemblyBegin(nullSpaceBasis);
    VecAssemblyEnd(nullSpaceBasis);
    VecNormalize(nullSpaceBasis, PETSC_NULL);

    MatNullSpace matNullSpace;
    MatNullSpaceCreate(mpiCommGlobal, PETSC_FALSE, 1, &nullSpaceBasis, &matNullSpace);
  }


  void PetscSolverNavierStokes::initPreconditioner(PC pc)
  {
    FUNCNAME("PetscSolverNavierStokes::initPreconditioner()");

    vector<int> velocityComponents;
    velocityComponents.push_back(0);
    velocityComponents.push_back(1);

    PCSetType(pc, PCFIELDSPLIT);
    PCFieldSplitSetType(pc, PC_COMPOSITE_SCHUR);
    PCFieldSplitSetSchurFactType(pc, PC_FIELDSPLIT_SCHUR_FACT_FULL);
    createFieldSplit(pc, "velocity", velocityComponents);
    createFieldSplit(pc, "pressure", pressureComponent);

    KSPSetUp(kspInterior);

    KSP *subKsp;
    int nSubKsp;
    PCFieldSplitGetSubKSP(pc, &nSubKsp, &subKsp);

    TEST_EXIT(nSubKsp == 2)
      ("Wrong numer of KSPs inside of the fieldsplit preconditioner!\n");

    KSP kspVelocity = subKsp[0];
    KSP kspSchur = subKsp[1];
    PetscFree(subKsp);

    Mat A00, A01, A10, A11;
    PCFieldSplitGetSchurBlocks(pc, &A00, &A01, &A10, &A11);
    matShellContext.A00 = A00;
    matShellContext.A01 = A01;
    matShellContext.A10 = A10;
    matShellContext.A11 = A11;
    matShellContext.kspVelocity = kspVelocity;

    Mat matShell;
    MatDuplicate(A11, MAT_DO_NOT_COPY_VALUES, &matShell);
    MatSetType(matShell, MATSHELL);
    MatShellSetContext(matShell, &matShellContext);
    MatShellSetOperation(matShell, MATOP_MULT, (void(*)(void))petscMultNavierStokesSchur);
    MatSetUp(matShell);

    KSPSetType(kspVelocity, KSPPREONLY);
    PC pcSub;
    KSPGetPC(kspVelocity, &pcSub);
    PCSetType(pcSub, PCLU);
    PCFactorSetMatSolverPackage(pcSub, MATSOLVERMUMPS);


    KSPSetType(kspSchur, KSPGMRES);
    KSPSetTolerances(kspSchur, 0, 1e-14, 1e+3, 1000);
    KSPSetOperators(kspSchur, matShell, matShell, SAME_NONZERO_PATTERN);
    KSPGetPC(kspSchur, &pcSub);
    PCSetType(pcSub, PCNONE);


    // === Mass matrix solver ===

    MSG("Initialize mass matrix solver ...\n");

    const FiniteElemSpace *pressureFeSpace = componentSpaces[pressureComponent];
    DOFMatrix massMatrix(pressureFeSpace, pressureFeSpace);
    Operator massOp(pressureFeSpace, pressureFeSpace);
    Simple_ZOT zot;
    massOp.addTerm(&zot);
    massMatrix.assembleOperator(massOp);

    massMatrixSolver = createSubSolver(pressureComponent, "mass_");
    massMatrixSolver->fillPetscMatrix(&massMatrix);

    MSG("... OK!\n");


    // === Laplace matrix solver ===

    MSG("Initialize laplace matrix solver ...\n");

    DOFMatrix laplaceMatrix(pressureFeSpace, pressureFeSpace);
    Operator laplaceOp(pressureFeSpace, pressureFeSpace);
    Simple_SOT sot;
    laplaceOp.addTerm(&sot);
    laplaceMatrix.assembleOperator(laplaceOp);

    laplaceMatrixSolver = createSubSolver(pressureComponent, "laplace_");
    laplaceMatrixSolver->fillPetscMatrix(&laplaceMatrix);

    MSG("... OK!\n");


    // === Create convection-diffusion operator ===

    MSG("Initialize pressure convection-diffusion operator ...\n");

    double timestep = 1.0;
    Parameters::get("navierstokes->adapt->timestep", timestep);
    timestep = 1.0 / timestep;

    double nu = 0.01;
    Parameters::get("nu", nu);

    DOFMatrix conDifMatrix(pressureFeSpace, pressureFeSpace);
    Operator conDifOp(pressureFeSpace, pressureFeSpace);
    Simple_ZOT conDif0(timestep);
    conDifOp.addTerm(&conDif0);
    Simple_SOT conDif1(nu);
    conDifOp.addTerm(&conDif1);
    Vector_FOT conDif2(0);
    //    conDifOp.addTerm(&conDif2, GRD_PHI);
    Vector_FOT conDif3(1);
    //    conDifOp.addTerm(&conDif3, GRD_PHI);

    conDifMatrix.assembleOperator(conDifOp);

    conDifMatrixSolver = createSubSolver(pressureComponent, "condif_");
    conDifMatrixSolver->fillPetscMatrix(&conDifMatrix);

    MSG("... OK!\n");

    matShellContext.kspMass = massMatrixSolver->getSolver();
    matShellContext.kspLaplace = laplaceMatrixSolver->getSolver();
    matShellContext.matConDif = conDifMatrixSolver->getMatInterior();    

    KSPSetType(matShellContext.kspMass, KSPRICHARDSON);
    KSPSetTolerances(matShellContext.kspMass, 0, 1e-13, 1e+3, 1);
    KSPGetPC(matShellContext.kspMass, &pcSub);
    PCSetType(pcSub, PCLU);
    PCFactorSetMatSolverPackage(pcSub, MATSOLVERMUMPS);
    KSPSetFromOptions(matShellContext.kspMass);

    KSPSetType(matShellContext.kspLaplace, KSPCG);
    KSPSetTolerances(matShellContext.kspLaplace, 0, 1e-13, 1e+3, 10000);
    KSPGetPC(matShellContext.kspLaplace, &pcSub);
    PCSetType(pcSub, PCJACOBI);
    KSPSetFromOptions(matShellContext.kspLaplace);

    matShellContext.solverMode = 1;


    KSPSetFromOptions(kspVelocity);
    KSPSetFromOptions(kspSchur);
  }

}