/*
 * Copyright (C)  NHR@FAU, University Erlangen-Nuremberg.
 * All rights reserved. This file is part of nusif-solver.
 * Use of this source code is governed by a MIT style
 * license that can be found in the LICENSE file.
 */
#include <math.h>
#include <stdio.h>

#include "allocate.h"
#include "comm.h"
#include "discretization.h"
#include "parameter.h"
#include "solver.h"
#include "util.h"

void printSolver(Solver* s, double* grid, char* gridname)
{
    printf("Grid name : %s", gridname);
    int imaxLocal = s->comm->imaxLocal;
    int jmaxLocal = s->comm->jmaxLocal;

    for (int i = 0; i < s->comm->size; i++) {
        if (i == s->comm->rank) {
            sleep(1 * s->comm->rank);
            printf("### RANK %d LVL "
                   "###################################################### #\n ",
                s->comm->rank);
            for (int j = 0; j < jmaxLocal + 2; j++) {
                printf("%02d: ", j);
                for (int i = 0; i < imaxLocal + 2; i++) {
                    printf("%2.2f  ", grid[j * (imaxLocal + 2) + i]);
                }
                printf("\n");
            }
            fflush(stdout);
        }
    }
}


void initSolver(Solver* s, Discretization* d, Parameter* p)
{
    s->grid    = &d->grid;
    s->eps     = p->eps;
    s->omega   = p->omg;
    s->itermax = p->itermax;
    s->comm    = &d->comm;
}

double solve(Solver* s, double* p, double* rhs)
{
    int imax      = s->grid->imax;
    int jmax      = s->grid->jmax;
    int imaxLocal = s->comm->imaxLocal;
    int jmaxLocal = s->comm->jmaxLocal;
    double eps    = s->eps;
    int itermax   = s->itermax;
    double dx2    = s->grid->dx * s->grid->dx;
    double dy2    = s->grid->dy * s->grid->dy;
    double idx2   = 1.0 / dx2;
    double idy2   = 1.0 / dy2;
    double factor = s->omega * 0.5 * (dx2 * dy2) / (dx2 + dy2);
    double epssq  = eps * eps;
    int pass, jsw, isw;
    int it     = 0;
    double res = 1.0;

    while ((res >= epssq) && (it < itermax)) {
        jsw = 1;
        for (pass = 0; pass < 2; pass++) {
            isw = jsw;
            commExchange(s->comm, p);

            for (int j = 1; j < jmaxLocal + 1; j++) {
                for (int i = isw; i < imaxLocal + 1; i += 2) {

                    double r = RHS(i, j) -
                               ((P(i + 1, j) - 2.0 * P(i, j) + P(i - 1, j)) * idx2 +
                                   (P(i, j + 1) - 2.0 * P(i, j) + P(i, j - 1)) * idy2);

                    P(i, j) -= (factor * r);
                    res += (r * r);
                }
                isw = 3 - isw;
            }
            jsw = 3 - jsw;
        }

        if (commIsBoundary(s->comm, B)) { // set bottom bc
            for (int i = 1; i < imaxLocal + 1; i++) {
                P(i, 0) = P(i, 1);
            }
        }

        if (commIsBoundary(s->comm, T)) { // set top bc
            for (int i = 1; i < imaxLocal + 1; i++) {
                P(i, jmaxLocal + 1) = P(i, jmaxLocal);
            }
        }

        if (commIsBoundary(s->comm, L)) { // set left bc
            for (int j = 1; j < jmaxLocal + 1; j++) {
                P(0, j) = P(1, j);
            }
        }

        if (commIsBoundary(s->comm, R)) { // set right bc
            for (int j = 1; j < jmaxLocal + 1; j++) {
                P(imaxLocal + 1, j) = P(imaxLocal, j);
            }
        }

        commReduction(&res, SUM);
        res = res / (double)(imax * jmax);
#ifdef DEBUG
        if (commIsMaster(s->comm)) {
            printf("%d Residuum: %e\n", it, res);
        }
#endif
        it++;
    }

#ifdef VERBOSE
    if (commIsMaster(s->comm)) {
        printf("Solver took %d iterations to reach %f\n", it, sqrt(res));
    }
#endif

return res;
}