/*
 * 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 <stdio.h>
#include <stdlib.h>

#include "allocate.h"
#include "solver.h"
#include "util.h"

#define FINEST_LEVEL   0
#define COARSEST_LEVEL (s->levels - 1)

#define E(i, j)   e[(j) * (imaxLocal + 2) + (i)]
#define R(i, j)   r[(j) * (imaxLocal + 2) + (i)]
#define OLD(i, j) old[(j) * (imaxLocal + 2) + (i)]

void printSolver(Solver* s, double* grid, char* gridname)
{
    if (0 == s->comm->rank) 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);
        }
    }
}

static void restrictMG(Solver* s, int level, Comm* comm)
{
    int imaxLocal = comm->imaxLocal;
    int jmaxLocal = comm->jmaxLocal;

    double* r   = s->r[level + 1];
    double* old = s->r[level];

#ifdef _MPI
    commExchange(comm, old);
#endif

    for (int j = 1; j < (jmaxLocal / 2) + 1; j++) {
        for (int i = 1; i < (imaxLocal / 2) + 1; i++) {
            R(i, j) = (OLD(2 * i - 1, 2 * j - 1) + OLD(2 * i, 2 * j - 1) * 2 +
                          OLD(2 * i + 1, 2 * j - 1) + OLD(2 * i - 1, 2 * j) * 2 +
                          OLD(2 * i, 2 * j) * 4 + OLD(2 * i + 1, 2 * j) * 2 +
                          OLD(2 * i - 1, 2 * j + 1) + OLD(2 * i, 2 * j + 1) * 2 +
                          OLD(2 * i + 1, 2 * j + 1)) /
                      16.0;
        }
    }
}

static void prolongate(Solver* s, int level, Comm* comm)
{
    int imaxLocal = comm->imaxLocal;
    int jmaxLocal = comm->jmaxLocal;

    double* old = s->r[level + 1];
    double* e   = s->r[level];

    for (int j = 2; j < jmaxLocal + 1; j += 2) {
        for (int i = 2; i < imaxLocal + 1; i += 2) {
            E(i, j) = OLD(i / 2, j / 2);
        }
    }
}

static void correct(Solver* s, double* p, int level, Comm* comm)
{
    double* e     = s->e[level];
    int imaxLocal = comm->imaxLocal;
    int jmaxLocal = comm->jmaxLocal;

    for (int j = 1; j < jmaxLocal + 1; ++j) {
        for (int i = 1; i < imaxLocal + 1; ++i) {
            P(i, j) += E(i, j);
        }
    }
}

static void setBoundaryCondition(Solver* s, double* p, int imaxLocal, int jmaxLocal)
{
#ifdef _MPI
    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);
        }
    }
#else
    for (int i = 1; i < imaxLocal + 1; i++) {
        P(i, 0)             = P(i, 1);
        P(i, jmaxLocal + 1) = P(i, jmaxLocal);
    }

    for (int j = 1; j < jmaxLocal + 1; j++) {
        P(0, j)             = P(1, j);
        P(imaxLocal + 1, j) = P(imaxLocal, j);
    }
#endif
}

static void smooth(Solver* s, double* p, double* rhs, int level, Comm* comm)
{
    int imaxLocal = comm->imaxLocal;
    int jmaxLocal = comm->jmaxLocal;

    int imax = s->grid->imax;
    int jmax = s->grid->jmax;

    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* r     = s->r[level];

    double res = 1.0;
    int pass, jsw, isw;

    jsw = 1;

    for (pass = 0; pass < 2; pass++) {
        isw = jsw;

#ifdef _MPI
        commExchange(comm, p);
#endif

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

                P(i, j) -= factor *
                           (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));
            }
            isw = 3 - isw;
        }
        jsw = 3 - jsw;
    }
}

static double calculateResidual(Solver* s, double* p, double* rhs, int level, Comm* comm)
{
    int imax      = s->grid->imax;
    int jmax      = s->grid->jmax;
    int imaxLocal = comm->imaxLocal;
    int jmaxLocal = comm->jmaxLocal;

    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* r     = s->r[level];
    double res    = 1.0;
    int pass, jsw, isw;

    jsw = 1;

    for (pass = 0; pass < 2; pass++) {
        isw = jsw;

#ifdef _MPI
        commExchange(comm, p);
#endif

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

                R(i, j) = 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);

                res += (R(i, j) * R(i, j));
            }
            isw = 3 - isw;
        }
        jsw = 3 - jsw;
    }

#ifdef _MPI
    commReduction(&res, SUM);
#endif

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

static double multiGrid(Solver* s, double* p, double* rhs, int level, Comm* comm)
{
    double res = 0.0;

    // coarsest level
    if (level == COARSEST_LEVEL) {
        for (int i = 0; i < 5; i++) {
            smooth(s, p, rhs, level, comm);
        }
        return res;
    }

    // pre-smoothing
    for (int i = 0; i < s->presmooth; i++) {
        smooth(s, p, rhs, level, comm);

        if (level == FINEST_LEVEL)
            setBoundaryCondition(s, p, comm->imaxLocal, comm->jmaxLocal);
    }

    // calculate residuals
    res = calculateResidual(s, p, rhs, level, comm);

    // restrict
    restrictMG(s, level, comm);

    // Create a new comm object withupdated imaxLocal and jmaxLocal
    // along with their updated bufferTypes, sdispls, rdispls
    Comm newcomm;
    commUpdateDatatypes(s->comm, &newcomm, comm->imaxLocal / 2, comm->jmaxLocal / 2);

    // MGSolver on residual and error.
    multiGrid(s, s->e[level + 1], s->r[level + 1], level + 1, &newcomm);

    commFreeCommunicator(&newcomm);

    // prolongate
    prolongate(s, level, comm);

    // correct p on finer level using residual
    correct(s, p, level, comm);

    if (level == FINEST_LEVEL)
        setBoundaryCondition(s, p, comm->imaxLocal, comm->jmaxLocal);

    // post-smoothing
    for (int i = 0; i < s->postsmooth; i++) {
        smooth(s, p, rhs, level, comm);
        if (level == FINEST_LEVEL)
            setBoundaryCondition(s, p, comm->imaxLocal, comm->jmaxLocal);
    }

    return res;
}

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

    int imax   = s->grid->imax;
    int jmax   = s->grid->jmax;
    int levels = s->levels;
    printf("Using Multigrid solver with %d levels\n", levels);

    s->r = malloc(levels * sizeof(double*));
    s->e = malloc(levels * sizeof(double*));

    size_t size = (imax + 2) * (jmax + 2) * sizeof(double);

    for (int j = 0; j < levels; j++) {
        s->r[j] = allocate(64, size);
        s->e[j] = allocate(64, size);

        for (int i = 0; i < (imax + 2) * (jmax + 2); i++) {
            s->r[j][i] = 0.0;
            s->e[j][i] = 0.0;
        }
    }
}

double solve(Solver* s, double* p, double* rhs)
{
    double res = multiGrid(s, p, rhs, 0, s->comm);

#ifdef VERBOSE
    if (commIsMaster(s->comm)) {
        printf("Residuum: %.6f\n", res);
    }
#endif
return res;
}