#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <omp.h>

#define N      1001
#define ITERS  1000000
#define OMEGA  1.9
#define TOL    0.0000000001

int main(void) {
    double (*A)[N] = malloc(sizeof(double[N][N]));
    double maxdiff;

    for (int j = 0; j < N; j++) {
        for (int i = 0; i < N; i++) {
            double v = (j == 0) ? 1.0 : 0.0;
            A[j][i] = v;
        }
    }

    #pragma omp target data map(tofrom: A[0:N][0:N])
    {
        for (int iter = 0; iter < ITERS; iter++) {
            maxdiff = 0.;

            // update A in 2 passes, like red and black on checkerboard
            // second pass uses updated values of A (Gauss-Seidel)
            for (int k = 0; k < 2; k++) {

                #pragma omp target teams distribute parallel for \
                            collapse(2) reduction(max:maxdiff)
                for (int j = 1; j < N - 1; j++) {
                    for (int i = 1; i < N - 1; i++) {

                        // on each row, half the threads do no work
                        if (((i + j) & 1) != k) continue;

                        double old = A[j][i];
                        double gs = 0.25 * ( A[j][i+1] + A[j][i-1]
                                           + A[j-1][i] + A[j+1][i] );
                        // update by successive over-relaxation (SOR)
                        // Gauss-Seidel is recovered if OMEGA = 1
                        A[j][i] = old + OMEGA * (gs - old);

                        maxdiff = fmax(maxdiff, fabs(A[j][i] - old));
                    }
                }

            }

            if (iter % 100 == 0) printf("maxdiff = %.12f\n", maxdiff);
            if (maxdiff <= TOL) {
                printf("maxdiff = %.12f, iter = %d\n", maxdiff, iter);
                break;
            }
        }
    }

    printf("A[N/2][N/2] = %.6f\n", A[N/2][N/2]);
    free(A);
    return 0;
}