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

load "mshmet"
load "PETSc"
macro with()PETSc//
load "parmmg"
load "mmg"
macro dimension()3// EOM            // 2D or 3D
include "macro_ddm.idp"             // additional DDM functions

int n = 12;
int[int] L=[1,1,1,1], LU=[0,2], LD=[1,1], LR=[0,1];
mesh3 Th3 = buildlayers(square(n, n, region=0, label=L), n, zbound=[0,1], labelmid=LD, labelup=LU, labeldown=LR);
Th3 = trunc(Th3, x < 0.5 || y < 0.5 || z < 0.5, label=3);

macro grad(u)[dx(u), dy(u), dz(u)]// EOM    // three-dimensional gradient
macro MmgParameters(ThGather, met, rt, verbose)ThGather, metric = met, hgrad = getARGV("-hgrad", 2.3), requiredTriangle = rt, verbose = verbose// EOM
macro ParMmgCommunicatorsAndMetric(Th, met, ThParMmg, metParMmg, communicators) {
int[int] n2o;
ParMmgCreateCommunicators(Th, ThParMmg, n2o, communicators);
int[int] rest(ThParMmg.nv * (met.n / Th.nv));
if(met.n == 6 * Th.nv) {
    fespace VhMet(Th, [P1, P1, P1, P1, P1, P1]);
    fespace VhParMmg(ThParMmg, [P1, P1, P1, P1, P1, P1]);
    rest = restrict(VhParMmg, VhMet, n2o);
}
else if(met.n == Th.nv) {
    fespace VhMet(Th, P1);
    fespace VhParMmg(ThParMmg, P1);
    rest = restrict(VhParMmg, VhMet, n2o);
}
else
    assert(0);
metParMmg.resize(rest.n);
metParMmg = met(rest);
}//

varf vPoisson(u, v) = int3d(Th3)(grad(u)' * grad(v)) + int3d(Th3)(v) + on(1,2,3, u=0);

Mat<PetscScalar> A;
DmeshCreate(Th3);
MatCreate(Th3, A, P1);
fespace Vh(Th3, P1);
Vh<PetscScalar> u;
real err = 8.0e-3;

int iMax = getARGV("-iMax", 1);
int noTransfer = (getARGV("-transfer", 1) == 0);
int niter = getARGV("-niter", 3);
int verbose = getARGV("-verbose", 1);
bool isotropic = (usedARGV("-isotropic") != -1);
int elementsPerProcess = max(getARGV("-elements_per_process", 1), 8000);
int maxP = max(min(getARGV("-max_P", mpisize), mpisize), 1);
int[int] rt(2);
rt = [1,2];
for(int i = 0; i < iMax; ++i) {
    A = vPoisson(Vh, Vh, tgv = -2);
    set(A, sparams = "-ksp_monitor -pc_type gamg");
    PetscScalar[int] rhs = vPoisson(0, Vh, tgv = -1);
    u[] = A^-1 * rhs;
    mesh3 ThParMmg;
    DmeshInitialize(ThParMmg);
    int P = 1;
    if(maxP != 1) {
        fespace Ph(Th3, P0);
        Ph part;
        PartitionCreate(Th3, part[], P0);
        part = abs(part - 1.0) < 1e-1;
        int nt, ntLocal = part[].l1;
        mpiAllReduce(ntLocal, nt, mpiCommWorld, mpiSUM);
        P = min(maxP, min(mpisize, max(1, nt / elementsPerProcess)));
    }
    if(P != mpisize) {
        int div = mpisize / P;
        mpiComm commThGather(mpiCommWorld, (mpirank % div == 0 && mpirank / div < P) ? 0 : mpiUndefined, mpirank / div);
        mpiComm comm(mpiCommWorld, min(mpirank / div, P - 1), mpirank - div * min(mpirank / div, P - 1));
        macro ThGatherComm()commThGather//
        mesh3 ThGather;
        DmeshGather(Th3, comm, ThGather);
        fespace VhGather(ThGather, P1);
        VhGather<PetscScalar> uGather;
        VecGather(Th3, comm, ThGather, P1, u, uGather);
        macro ThGatherParMmgComm()commThGather//
        mesh3 ThGatherParMmg;
        DmeshInitialize(ThGatherParMmg);
        if((mpirank % div == 0 && mpirank / div < P) != 0) {
            real[int] met = mshmet(ThGather, abs(uGather), aniso = !isotropic, hmin = 1.0e-3, hmax = 1.0e-1, err = err);
            if(mpiSize(commThGather) > 1) {
                real[int] metParMmg;
                int[int][int] communicators;
                ParMmgCommunicatorsAndMetric(ThGather, met, ThGatherParMmg, metParMmg, communicators);
                ThGatherParMmg = parmmg3d(MmgParameters(ThGatherParMmg, metParMmg, rt, verbose), nodeCommunicators = communicators, niter = niter, comm = commThGather);
            }
            else
                ThGatherParMmg = mmg3d(MmgParameters(ThGather, met, rt, verbose));
        }
        DmeshScatter(ThGatherParMmg, comm, ThParMmg);
    }
    else {
        real[int] met = mshmet(Th3, abs(u), aniso = !isotropic, hmin = 1.0e-3, hmax = 1.0e-1, err = err);
        real[int] metParMmg;
        int[int][int] communicators;
        ParMmgCommunicatorsAndMetric(Th3, met, ThParMmg, metParMmg, communicators);
        ThParMmg = parmmg3d(MmgParameters(ThParMmg, metParMmg, rt, verbose), nodeCommunicators = communicators, niter = niter);
        DmeshReconstruct(ThParMmg);
    }
    int[int] fforder = [1];
    savevtk("laplace-adapt-dist-3d.vtu", Th3, abs(u), bin = 1, order = fforder, append = i ? true : false);
    if(!noTransfer) {
        fespace VhAdapt(ThParMmg, P1);
        VhAdapt<PetscScalar> uAdapt;
        VecInterpolate(Th3, P1, u, ThParMmg, P1, uAdapt);
        savevtk("laplace-adapt-dist-3d.vtu", ThParMmg, abs(uAdapt), bin = 1, order = fforder, append = true);
    }
    DmeshCopy(ThParMmg, Th3);
    Mat<PetscScalar> Adapt;
    MatCreate(Th3, Adapt, P1);
    A = Adapt;
    u = 0.0;
    err *= 0.5;
}
