//
// 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/PetscSolverFetiOperators.h"
#include "parallel/PetscSolverFetiStructs.h"
#include "parallel/PetscSolverFetiTimings.h"

namespace AMDiS {

  int petscMultMatSchurPrimal(Mat mat, Vec x, Vec y)
  {
    // S_PiPi = K_PiPi - K_PiB inv(K_BB) K_BPi

    void *ctx;
    MatShellGetContext(mat, &ctx);
    SchurPrimalData* data = static_cast<SchurPrimalData*>(ctx);

    MatMult(data->subSolver->getMatInteriorCoarse(), x, data->tmp_vec_b);
    data->subSolver->solveGlobal(data->tmp_vec_b, data->tmp_vec_b);
    MatMult(data->subSolver->getMatCoarseInterior(), data->tmp_vec_b, 
	    data->tmp_vec_primal);
    MatMult(data->subSolver->getMatCoarse(), x, y);
    VecAXPBY(y, -1.0, 1.0, data->tmp_vec_primal);

    return 0;
  }


  // y = mat * x
  int petscMultMatFeti(Mat mat, Vec x, Vec y)
  {
    FUNCNAME("petscMultMatFeti()");

    //    F = L inv(K_BB) trans(L) + L inv(K_BB) K_BPi inv(S_PiPi) K_PiB inv(K_BB) trans(L)
    // => F = L [I + inv(K_BB) K_BPi inv(S_PiPi) K_PiB] inv(K_BB) trans(L)

    double wtime = MPI::Wtime();

    void *ctx;
    MatShellGetContext(mat, &ctx);
    FetiData* data = static_cast<FetiData*>(ctx);

    MatMultTranspose(*(data->mat_lagrange), x, data->tmp_vec_b0);

    double wtime01 = MPI::Wtime();
    data->subSolver->solveGlobal(data->tmp_vec_b0, data->tmp_vec_b0);

    FetiTimings::fetiSolve01 += (MPI::Wtime() - wtime01);

    MatMult(*(data->mat_lagrange), data->tmp_vec_b0, data->tmp_vec_lagrange);

    MatMult(data->subSolver->getMatCoarseInterior(), 
	    data->tmp_vec_b0, data->tmp_vec_primal0);

    wtime01 = MPI::Wtime();
    KSPSolve(*(data->ksp_schur_primal), data->tmp_vec_primal0, data->tmp_vec_primal0);
    FetiTimings::fetiSolve02 += (MPI::Wtime() - wtime01);

    MatMult(data->subSolver->getMatInteriorCoarse(), 
	    data->tmp_vec_primal0, data->tmp_vec_b0);

    wtime01 = MPI::Wtime();
    data->subSolver->solveGlobal(data->tmp_vec_b0, data->tmp_vec_b0);
    FetiTimings::fetiSolve01 += (MPI::Wtime() - wtime01);

    MatMult(*(data->mat_lagrange), data->tmp_vec_b0, y);

    VecAXPBY(y, 1.0, 1.0, data->tmp_vec_lagrange);

    FetiTimings::fetiSolve += (MPI::Wtime() - wtime);

    return 0;
  }


  int petscMultMatFetiInterface(Mat mat, Vec x, Vec y)
  {
    FUNCNAME("petscMultMatFetiInterface()");

    double wtime = MPI::Wtime();

    Vec x_interface, x_lagrange, y_interface, y_lagrange;    
    VecNestGetSubVec(x, 0, &x_interface);
    VecNestGetSubVec(x, 1, &x_lagrange);
    VecNestGetSubVec(y, 0, &y_interface);
    VecNestGetSubVec(y, 1, &y_lagrange);

    void *ctx;
    MatShellGetContext(mat, &ctx);
    FetiData* data = static_cast<FetiData*>(ctx);


    // === Calculation of v_{B} ===

    // tmp_vec_b0 = J^{T} \lambda
    MatMultTranspose(*(data->mat_lagrange), x_lagrange, data->tmp_vec_b0);
    // tmp_vec_b1 = A_{B\Gamma} u_{\Gamma}
    MatMult(data->subSolver->getMatInteriorCoarse(1), x_interface, data->tmp_vec_b1);
    // tmp_vec_b0 = A_{B\Gamma} u_{\Gamma} + J^{T} \lambda
    VecAXPY(data->tmp_vec_b0, 1.0, data->tmp_vec_b1);


    // == Calculation of v_{\Pi}
    
    // tmp_vec_primal0 = A_{\Pi\Gamma} u_{\Gamma}
    MatMult(data->subSolver->getMatCoarse(0, 1), x_interface, data->tmp_vec_primal0);
    

    // === Calculate action of FETI-DP operator ===

    double wtime01 = MPI::Wtime();
    // tmp_vec_b0 = A_{BB}^{-1} v_{B}
    data->subSolver->solveGlobal(data->tmp_vec_b0, data->tmp_vec_b0);
    FetiTimings::fetiSolve01 += (MPI::Wtime() - wtime01);

    // tmp_vec_primal1 = A_{\Pi B} A_{BB}^{-1} v_{B}
    MatMult(data->subSolver->getMatCoarseInterior(), 
	    data->tmp_vec_b0, data->tmp_vec_primal1);

    // tmp_vec_primal0 = v_{\Pi} - A_{\Pi B} A_{BB}^{-1} v_{B}
    VecAXPY(data->tmp_vec_primal0, -1.0, data->tmp_vec_primal1);

    wtime01 = MPI::Wtime();
    // tmp_vec_primal0 = S_{\Pi\Pi}^{-1} (v_{\Pi} - A_{\Pi B} A_{BB}^{-1} v_{B})
    KSPSolve(*(data->ksp_schur_primal), data->tmp_vec_primal0, data->tmp_vec_primal0);
    FetiTimings::fetiSolve02 += (MPI::Wtime() - wtime01);

    // tmp_vec_b1 = A_{B\Pi} S_{\Pi\Pi}^{-1} (v_{\Pi} - A_{\Pi B} A_{BB}^{-1} v_{B})
    MatMult(data->subSolver->getMatInteriorCoarse(), 
	    data->tmp_vec_primal0, data->tmp_vec_b1);

    wtime01 = MPI::Wtime();
    // tmp_vec_b1 = A_{BB}^{-1} A_{B\Pi} S_{\Pi\Pi}^{-1} (v_{\Pi} - A_{\Pi B} A_{BB}^{-1} v_{B})
    data->subSolver->solveGlobal(data->tmp_vec_b1, data->tmp_vec_b1);
    FetiTimings::fetiSolve01 += (MPI::Wtime() - wtime01);

    // tmp_vec_b0 = A_{BB}^{-1} v_{B} - A_{BB}^{-1} A_{B\Pi} S_{\Pi\Pi}^{-1} (v_{\Pi} - A_{\Pi B} A_{BB}^{-1} v_{B})
    VecAXPY(data->tmp_vec_b0, -1.0, data->tmp_vec_b1);


    // === Calculate projection to interface and constraint variable ===

    // y_interface = A_{\Gamma B} tmp_vec_b0
    MatMult(data->subSolver->getMatCoarseInterior(1), data->tmp_vec_b0, y_interface);

    // tmp_vec_primal0 = S_{\Pi\Pi}^{-1} (v_{\Pi} - A_{\Pi B} A_{BB}^{-1} v_{B})
    // tmp_vec_interface = A_{\Gamma \Pi} tmp_vec_primal0
    MatMult(data->subSolver->getMatCoarse(1, 0), data->tmp_vec_primal0, data->tmp_vec_interface);
    // y_interface = A_{\Gamma B} tmp_vec_b0 + A_{\Gamma \Pi} tmp_vec_primal0
    VecAXPY(y_interface, 1.0, data->tmp_vec_interface);

    // y_lagrange = J tmp_vec_b0
    MatMult(*(data->mat_lagrange), data->tmp_vec_b0, y_lagrange);

    FetiTimings::fetiSolve += (MPI::Wtime() - wtime);

    return 0;
  }


  PetscErrorCode petscApplyFetiDirichletPrecon(PC pc, Vec x, Vec y)
  {
    double wtime = MPI::Wtime();

    // Get data for the preconditioner
    void *ctx;
    PCShellGetContext(pc, &ctx);
    FetiDirichletPreconData* data = static_cast<FetiDirichletPreconData*>(ctx);

    // Multiply with scaled Lagrange constraint matrix.
    MatMultTranspose(*(data->mat_lagrange_scaled), x, data->tmp_vec_b);


    // === Restriction of the B nodes to the boundary nodes. ===

    int nLocalB;
    int nLocalDuals;
    VecGetLocalSize(data->tmp_vec_b, &nLocalB);
    VecGetLocalSize(data->tmp_vec_duals0, &nLocalDuals);

    PetscScalar *local_b, *local_duals;
    VecGetArray(data->tmp_vec_b, &local_b);
    VecGetArray(data->tmp_vec_duals0, &local_duals);

    for (map<int, int>::iterator it = data->localToDualMap.begin();
	 it != data->localToDualMap.end(); ++it)
      local_duals[it->second] = local_b[it->first];

    VecRestoreArray(data->tmp_vec_b, &local_b);
    VecRestoreArray(data->tmp_vec_duals0, &local_duals);


    // === K_DD - K_DI inv(K_II) K_ID ===

    MatMult(*(data->mat_duals_duals), data->tmp_vec_duals0, data->tmp_vec_duals1);

    MatMult(*(data->mat_interior_duals), data->tmp_vec_duals0, data->tmp_vec_interior);
    KSPSolve(*(data->ksp_interior), data->tmp_vec_interior, data->tmp_vec_interior);
    MatMult(*(data->mat_duals_interior), data->tmp_vec_interior, data->tmp_vec_duals0);

    VecAXPBY(data->tmp_vec_duals0, 1.0, -1.0, data->tmp_vec_duals1);


    // === Prolongation from local dual nodes to B nodes.

    VecGetArray(data->tmp_vec_b, &local_b);
    VecGetArray(data->tmp_vec_duals0, &local_duals);

    for (map<int, int>::iterator it = data->localToDualMap.begin();
	 it != data->localToDualMap.end(); ++it)
      local_b[it->first] = local_duals[it->second];

    VecRestoreArray(data->tmp_vec_b, &local_b);
    VecRestoreArray(data->tmp_vec_duals0, &local_duals);


    // Multiply with scaled Lagrange constraint matrix.
    MatMult(*(data->mat_lagrange_scaled), data->tmp_vec_b, y);

    FetiTimings::fetiPreconditioner += (MPI::Wtime() - wtime);

    return 0;
  }


  PetscErrorCode petscApplyFetiLumpedPrecon(PC pc, Vec x, Vec y)
  {
    // Get data for the preconditioner
    void *ctx;
    PCShellGetContext(pc, &ctx);
    FetiLumpedPreconData* data = static_cast<FetiLumpedPreconData*>(ctx);

    // Multiply with scaled Lagrange constraint matrix.
    MatMultTranspose(*(data->mat_lagrange_scaled), x, data->tmp_vec_b);


    // === Restriction of the B nodes to the boundary nodes. ===

    int nLocalB;
    int nLocalDuals;
    VecGetLocalSize(data->tmp_vec_b, &nLocalB);
    VecGetLocalSize(data->tmp_vec_duals0, &nLocalDuals);

    PetscScalar *local_b, *local_duals;
    VecGetArray(data->tmp_vec_b, &local_b);
    VecGetArray(data->tmp_vec_duals0, &local_duals);

    for (map<int, int>::iterator it = data->localToDualMap.begin();
	 it != data->localToDualMap.end(); ++it)
      local_duals[it->second] = local_b[it->first];

    VecRestoreArray(data->tmp_vec_b, &local_b);
    VecRestoreArray(data->tmp_vec_duals0, &local_duals);


    // === K_DD ===

    MatMult(*(data->mat_duals_duals), data->tmp_vec_duals0, data->tmp_vec_duals1);


    // === Prolongation from local dual nodes to B nodes.

    VecGetArray(data->tmp_vec_b, &local_b);
    VecGetArray(data->tmp_vec_duals1, &local_duals);

    for (map<int, int>::iterator it = data->localToDualMap.begin();
	 it != data->localToDualMap.end(); ++it)
      local_b[it->first] = local_duals[it->second];

    VecRestoreArray(data->tmp_vec_b, &local_b);
    VecRestoreArray(data->tmp_vec_duals0, &local_duals);


    // Multiply with scaled Lagrange constraint matrix.
    MatMult(*(data->mat_lagrange_scaled), data->tmp_vec_b, y);

    return 0;
  }

}