//  run with MPI:  ff-mpirun -np 4 script.edp
// NBPROC 4

load "hpddm"                        // HPDDM plugin
macro dimension()2// EOM            // 2D or 3D
include "macro_ddm.idp"             // additional DDM functions

macro def(i)i// EOM                         // scalar field definition
macro init(i)i// EOM                        // scalar field initialization
macro grad(u)[dx(u), dy(u)]// EOM           // two-dimensional gradient
func Pk = P2;                               // finite element space

string deflation = getARGV("-deflation", "geneo");              // coarse space construction
int overlap = getARGV("-overlap", 1);                           // geometric overlap between subdomains
int fakeInterface = getARGV("-interface", 10);                  // interface between subdomains
int s = getARGV("-split", 1);                                   // refinement factor
real dt = getARGV("-dt", 0.1);                                 // time step
real mu = 0.01;
int iMax = getARGV("-iMax", 10);                                // number of iterations

mpiComm comm;
int p = getARGV("-hpddm_level_2_p", 1);
bool excluded = splitComm(mpiCommWorld, p, comm, topology = getARGV("-hpddm_level_2_topology", 0), exclude = (usedARGV("-hpddm_level_2_exclude") != -1));

if(verbosity > 0 && mpirank == 0) {
    cout << " --- " << mpirank << "/" << mpisize;
    cout << " - heat-2d.edp - input parameters: refinement factor = " << s << " - overlap = " << overlap << endl;
}

mesh ThBorder, Th = square(1, 1);
fespace Wh(Th, Pk);           // local finite element space
int[int][int] intersection;   // local-to-neighbors renumbering
real[int] D;                  // partition of unity
{
    int[int] l = [1, 1, 1, 1];
    Th = square(getARGV("-global", 40), getARGV("-global", 40), label = l);    // global mesh
    buildOverlap(Th, ThBorder, fakeInterface, s, overlap, intersection, D, Pk, comm, excluded)
}

real[int] rhs;                              // local right-hand side
matrix<real> Mat;                           // local operator
matrix<real> M;                             // local mass matrix
{                                           // local weak form
    mesh ThAugmented = Th + ThBorder;
    fespace Ph(ThAugmented, P0);
    Ph f = x < 0.25 ? 1.0e+3 : 1.0e-3;
    varf vPb(u, v) = int2d(ThAugmented)(u * v + dt * (grad(u)' * grad(v))) + int2d(ThAugmented)(dt * f *  v) + on(1, u = 0.0);
    fespace WhAugmented(ThAugmented, Pk);
    Mat = vPb(WhAugmented, WhAugmented, tgv = -1);
    real[int] rhsFull = vPb(0, WhAugmented, tgv = -1);
    matrix R = interpolate(Wh, WhAugmented);
    varf vPbM(u, v) = int2d(ThAugmented)(u * v);
    M = vPbM(WhAugmented, WhAugmented);
    renumbering(M, R, rhsFull, rhs);
    renumbering(Mat, R, rhsFull, rhs);
}
ThBorder = square(1, 1);

schwarz A(Mat, intersection, D, scaled = false);

matrix<real> Opt;                           // local operator with optimized boundary conditions
pair ret;
{
    int solver = getOption("schwarz_method");
    if(solver == 1 || solver == 2 || solver == 4) { // optimized Schwarz methods
        fespace Ph(Th, P0);
        real kZero = getARGV("-kZero", 10.0);
        Ph transmission = kZero;
        varf vOptimized(u, v) = int2d(Th)(u * v + dt * (grad(u)' * grad(v))) + int1d(Th, fakeInterface)(transmission * (u * v)) + on(1, u = 0.0);
        Opt = vOptimized(Wh, Wh, tgv = -1);
    }
    if(mpisize > 1 && isSetOption("schwarz_coarse_correction")) { // two-level Schwarz methods
        if(excluded)
            AttachCoarseOperator(mpiCommWorld, A/*, A = noPen, B = overlapRestriction, threshold = 2. * h[].max / diam*/);
        else {
            varf vPbNoPen(u, v) = int2d(Th)(u * v + dt * (grad(u)' * grad(v))) + on(1, u = 0.0);
            matrix<real> noPen = vPbNoPen(Wh, Wh, sym = 1);
            if(deflation == "geneo") // standard GenEO, no need for RHS -> deduced from LHS (Neumann matrix)
                AttachCoarseOperator(mpiCommWorld, A, A = noPen/*, threshold = 2. * h[].max / diam,*/, ret = ret);
            else if(deflation == "dtn") {
                varf vMass(def(u), def(v)) = int1d(Th, fakeInterface)(u * v);
                matrix<real> massMatrix = vMass(Wh, Wh, sym = 1);
                AttachCoarseOperator(mpiCommWorld, A, A = noPen, B = massMatrix, pattern = Opt/*, threshold = k,*/, ret = ret);
            }
            else if(deflation == "geneo-2") // GenEO-2 for optimized Schwarz methods, need for RHS (LHS is still Neumann matrix)
                AttachCoarseOperator(mpiCommWorld, A, A = noPen, B = Opt, pattern = Opt/*, threshold = 2. * h[].max / diam,*/, ret = ret);
        }
    }
}

set(A, sparams = "-hpddm_reuse_preconditioner=1");
Wh<real> def(u) = init(0.0);    // local solution
for(int i = 0; i < iMax; ++i) {
    real[int] newRhs(rhs.n);
    dmv(A, M, u[], newRhs);
    newRhs += rhs;

    if(Opt.n > 0)       // optimized Schwarz methods
        DDM(A, newRhs, u[], excluded = excluded, ret = ret, O = Opt);
    else
        u[] = A^-1 * newRhs;

    plotMPI(Th, u, Pk, def, real, cmm = "Global solution");
}
