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

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

macro def(i)[i, i#B]// EOM          // vector field definition
macro init(i)[i, i]// EOM           // vector field initialization
real Sqrt = sqrt(2.0);
macro epsilon(u)[dx(u), dy(u#B), (dy(u) + dx(u#B)) / Sqrt]// EOM
macro div(u)(dx(u) + dy(u#B))// EOM
func Pk = [vectorialfe, vectorialfe];                 // 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 << " - elasticity-2d.edp - input parameters: refinement factor = " << s << " - overlap = " << overlap << endl;
}

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

real f = -90000.0;
real strain = 100.0;
real Young = 1.0e8;
real poisson = 0.45;
real tmp = 1.0 + poisson;
real mu = Young  / (2.0 * tmp);
real lambda = Young * poisson / (tmp * (1.0 - 2.0 * poisson));
matrix<real> Mat;                           // local operator
varf vPb(def(u), def(v)) = int2d(Th)(lambda * div(u) * div(v) + 2.0 * mu * (epsilon(u)' * epsilon(v))) + int2d(Th)(f * vB) + on(1, u = 0.0, uB = 0.0);
Mat = vPb(Wh, Wh, sym = 1);
real[int] rhs = vPb(0, Wh);

schwarz A(Mat, intersection, D);

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 = 2 * kZero * mu * (2 * mu + lambda) / (lambda + 3 * mu);
        varf vOptimized(def(u), def(v)) = int2d(Th)(lambda * div(u) * div(v) + 2.0 * mu * (epsilon(u)' * epsilon(v))) + int1d(Th, fakeInterface)(transmission * (def(u)' * def(v))) + on(1, u = 0.0, uB = 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(def(u), def(v)) = int2d(Th)(lambda * div(u) * div(v) + 2.0 * mu * (epsilon(u)' * epsilon(v))) + on(1, u = 0.0, uB = 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);
        }
    }
}

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;

real[int] err(u[].n);
err = A * u[];      // global matrix-vector product
err -= rhs;

plotMPI(Th, def(u), Pk, def, real, cmm = "Global solution");
u[] = err;
macro def1(u)u// EOM
plotMPI(Th, u, vectorialfe, def1, real, cmm = "Global residual");
real alpha = 1.0e+1;
mesh ThMoved = movemesh(Th, [x + alpha * u, y + alpha * uB]);
u[] = mpirank;
plotMPI(ThMoved, u, vectorialfe, def1, real, cmm = "Global moved solution");
