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

int n = 4;
real[int] b(n), u(n);
func real[int] matId(real[int]& u) { return u; };
func real[int] A(real[int]& u) {
    for(int i = 0; i < u.n; ++i)
        u[i] = (10 * mpirank + i + 1) * u[i];
    return u;
}
func real error(real[int]& u , real[int]& b) {
    real s = 0;
    for(int i = 0; i < u.n; ++i)
        s += abs((10 * mpirank + i + 1) * u[i] - b[i]);
    return s;
}

b = 1;
u = 0;
IterativeMethod(A, b, u, sparams = "-hpddm_tol 1e-6 -hpddm_foo_krylov_method gcrodr -hpddm_foo_recycle 4 -hpddm_foo_variant flexible" + (mpirank == 0 ? " -hpddm_foo_verbosity 4" : ""), prefix = "foo_");
assert(error(u, b) < 1e-5);
u = 0;
IterativeMethod(A, b, u, precon = matId, prefix = "foo_");
assert(error(u, b) < 1e-5);
func real[int] B(real[int]& u) {
    for(int i = 0; i < u.n; ++i)
        u[i] = (2 * mpirank + i + 1) * u[i];
    return u;
}
mpiWtime();
func real[int] time(real[int]& u) {
    for(int i = 0;i < u.n; ++i)
        u[i] *= (1.0 + 1.0 / mpiWtime());
    return u;
}
u = 0;
IterativeMethod(B, b, u, precon = time, prefix = "foo_");
func real errorB(real[int]& u , real[int]& b) {
    real s = 0;
    for(int i = 0; i < u.n; ++i)
        s += abs((2 * mpirank + i + 1) * u[i] - b[i]);
    return s;
}
assert(errorB(u, b) < 1e-5);
