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

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

macro def(i)[i, i#B, i#C, i#D]// EOM// vector field definition
macro init(i)[i, i, i, i]// EOM     // vector field initialization
macro grad(u)[dx(u), dy(u), dz(u)]//// two-dimensional gradient
real Sqrt = sqrt(2.);
macro div(u)(dx(u) + dy(u#B) + dz(u#C))// EOM
func Pk = [P2, P2, P2, P1];             // 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

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 << " - stokes-3d.edp - input parameters: refinement factor = " << s << " - overlap = " << overlap << endl;
}

mesh3 ThBorder, Th = buildlayers(square(1, 1), 1);
fespace Wh(Th, Pk);           // local finite element space
int[int][int] intersection;   // local-to-neighbors renumbering
real[int] D;                  // partition of unity
{
    mesh ThGlobal2d = square(getARGV("-global", 12), getARGV("-global", 12), [x, y]);    // global mesh
    ThGlobal2d = trunc(ThGlobal2d, (x <= 0.5) || (y <= 0.5), label = 5);
    ThGlobal2d = trunc(ThGlobal2d, (y >= 0.25) || (x >= 0.25), label = 5);
    mesh Th2d = movemesh(ThGlobal2d, [-x, y]);
    ThGlobal2d = ThGlobal2d + Th2d;
    Th = buildlayers(ThGlobal2d, getARGV("-global", 12) / 2, zbound = [0, 0.4]);
    buildOverlap(Th, ThBorder, fakeInterface, s, overlap, intersection, D, Pk, comm, excluded)
}

real[int] rhs;                              // local right-hand side
matrix<real> Mat;                           // local operator
{                                           // local weak form
    mesh3 ThAugmented = Th + ThBorder;
    varf vPb([u, uB, uC, p], [v, vB, vC, q]) = int3d(ThAugmented)(grad(u)' * grad(v) + grad(uB)' * grad(vB) + grad(uC)' * grad(vC) - div(u) * q - div(v) * p + 1e-10 * p * q) + on(0, 1, 3, 5, u = 0, uB = 0, uC = 0) + on(2, u = 1000*y*(0.5-y)*z*(0.4-z), uB = 0, uC = 0);
    fespace WhAugmented(ThAugmented, Pk);
    Mat = vPb(WhAugmented, WhAugmented, tgv = -1);
    real[int] rhsFull = vPb(0, WhAugmented, tgv = -1);
    matrix R = interpolate(Wh, WhAugmented);
    renumbering(Mat, R, rhsFull, rhs, eps = -1);
}
ThBorder = buildlayers(square(1, 1), 1);

schwarz A(Mat, intersection, D, scaled = false);
set(A, sparams = "-hpddm_mkl_pardiso_iparm_13 1 --hpddm_schwarz_method oras -hpddm_schwarz_coarse_correction deflated -hpddm_geneo_nu 1");

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, uB, uC, p], [v, vB, vC, q]) = int3d(Th)(grad(u)' * grad(v) + grad(uB)' * grad(vB) + grad(uC)' * grad(vC) - div(u) * q - div(v) * p + 1e-10 * p * q) + int2d(Th, fakeInterface)(transmission * ([u, uB, uC]' * [v, vB, vC])) + on(0, 1, 3, 5, u = 0, uB = 0, uC = 0) + on(2, u = y*(0.5-y), uB = 0, uC = 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, uB, uC, p], [v, vB, vC, q]) = int3d(Th)(grad(u)' * grad(v) + grad(uB)' * grad(vB) + grad(uC)' * grad(vC) - div(u) * q - div(v) * p + 1e-10 * p * q) + on(0, 1, 3, 5, u = 0, uB = 0, uC = 0) + on(2, u = y*(0.5-y), uB = 0, uC = 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 == "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);
        }
    }
}

Wh<real> def(u);    // local solution

if(Opt.n > 0)       // optimized Schwarz methods
    DDM(A, rhs, u[], excluded = excluded, ret = ret, O = Opt);
else
    u[] = A^-1 * rhs;
int[int] fforder(2);
fforder = [1, 1];
savevtk("stokes-io-3d.vtu", Th, [u, uB, uC], uD, order = fforder, dataname = "velocity pressure");
