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

load "PETSc-complex"                // PETSc plugin
include "macro_ddm.idp"             // additional DDM functions

macro grad(u)[dx(u), dy(u)]// EOM   // two-dimensional gradient
func Pk = P1;                       // finite element space

mesh Th = square(getARGV("-global", 40), getARGV("-global", 40)); // global mesh
Mat<complex> A;
int s = getARGV("-split", 1);
macro ThRefinementFactor()s//
MatCreate(Th, A, Pk);

fespace Wh(Th, Pk);                 // local finite element space

func real wedge(real a, real b) {
    if(y < 0.4 + 0.1 * 0.75 * x)
        return 2;
    else if(y < 0.8 - 0.2 * 0.75 * x)
        return 1.5;
    else
        return 3;
}
real omega = 2 * pi * 5;
real xi = 0, yi = 0;
func f = 80 * 100 * s * exp(-20 * 100 * s * ((x-0.5-xi)^2 + (y-0.25-yi)^2));
complex[int] rhs(Wh.ndof);                  // local right-hand side
matrix<complex> Loc;                        // local operator
{                                           // local weak form
    fespace Ph(Th, P0);
    Ph val = wedge(x, y);
    Ph k = omega / val;
    varf vPb(u, v) = int2d(Th)(-k^2 * u * v + (grad(u)' * grad(v))) + int1d(Th, -111111)(1i * k * u * v) + int1d(Th, 2)(1i * k * u * v) + int2d(Th)(f * v) + on(1, u = 0.0);
    Loc = vPb(Wh, Wh, tgv = -1);
    rhs = vPb(0, Wh, tgv = -1);
}

A = Loc;
if(usedARGV("-optimized_Schwarz") != -1)
    set(A, sparams = "-ksp_view -sub_pc_type lu", O = Loc);
else
    set(A, sparams = "-ksp_view -pc_type lu");
Wh<complex> u;                      // local solution

u[] = A^-1 * rhs;

macro def(u)u//
plotMPI(Th, u, Pk, def, complex, cmm = "Global solution");

int M = 4;
complex[int, int] RHS(rhs.n, M);
RHS(:, 0) = rhs;
for(int i = 1; i < M; ++i) {
    xi += 0.1;
    yi += 0.1;
    varf vPb(u, v) = int2d(Th)(f * v) + on(1, u = 0.0);
    RHS(:, i) = vPb(0, Wh, tgv = -1);
}
set(A, sparams = "-ksp_type hpddm -ksp_converged_reason -ksp_hpddm_type bgmres");
complex[int, int] B = A^-1 * RHS;
for(int i = 0; i < M; ++i) {
    u[] = B(:, i);
    macro params()cmm = "Global solution #" + i, wait = 1, fill = 1, value = 1, dim = 3//
    plotMPI(Th, u, Pk, def, complex, params);
}
complex[int] C(B.n * B.m);
C = A^-1 * RHS.asarray;
for(int i = 0; i < M; ++i) {
    u[] = B(:, i);
    u[] -= C(i * Wh.ndof:(i + 1) * Wh.ndof - 1);
    macro params()cmm = "Error #" + i, wait = 1, fill = 1, value = 1, dim = 3//
    plotMPI(Th, u, Pk, def, complex, params);
}
