/*
 * 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 <float.h>
#include <math.h>
#include <stdio.h>
#include <string.h>

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

static double distance(
    double i, double j, double k, double iCenter, double jCenter, double kCenter)
{
    return sqrt(pow(iCenter - i, 2) + pow(jCenter - j, 2) + pow(kCenter - k, 2) * 1.0);
}

static double sumOffset(double* sizes, int init, int offset, int coord)
{
    double sum = 0;

    for (int i = init - offset; coord > 0; i -= offset, --coord) {
        sum += sizes[i];
    }

    return sum;
}

static void printConfig(Discretization* d)
{
    if (commIsMaster(&d->comm)) {
        printf("Parameters for #%s#\n", d->problem);
        printf("BC Left:%d Right:%d Bottom:%d Top:%d Front:%d Back:%d\n",
            d->bcLeft,
            d->bcRight,
            d->bcBottom,
            d->bcTop,
            d->bcFront,
            d->bcBack);
        printf("\tReynolds number: %.2f\n", d->re);
        printf("\tGx Gy: %.2f %.2f %.2f\n", d->gx, d->gy, d->gz);
        printf("Geometry data:\n");
        printf("\tDomain box size (x, y, z): %.2f, %.2f, %.2f\n",
            d->grid.xlength,
            d->grid.ylength,
            d->grid.zlength);
        printf("\tCells (x, y, z): %d, %d, %d\n",
            d->grid.imax,
            d->grid.jmax,
            d->grid.kmax);
        printf("\tCell size (dx, dy, dz): %f, %f, %f\n",
            d->grid.dx,
            d->grid.dy,
            d->grid.dz);
        printf("Timestep parameters:\n");
        printf("\tDefault stepsize: %.2f, Final time %.2f\n", d->dt, d->te);
        printf("\tdt bound: %.6f\n", d->dtBound);
        printf("\tTau factor: %.2f\n", d->tau);
        printf("Iterative parameters:\n");
        printf("\tMax iterations: %d\n", d->itermax);
        printf("\tepsilon (stopping tolerance) : %f\n", d->eps);
        printf("\tgamma factor: %f\n", d->gamma);
        printf("\tomega (SOR relaxation): %f\n", d->omega);
    }
    commPrintConfig(&d->comm);
}

void printGrid(Discretization* d, double* s)
{
    int imaxLocal = d->comm.imaxLocal;
    int jmaxLocal = d->comm.jmaxLocal;
    int kmaxLocal = d->comm.kmaxLocal;

    for (int k = 0; k < kmaxLocal + 2; k++) {
        printf("K : %02d:\n", k);
        for (int j = 0; j < jmaxLocal + 2; j++) {
            printf("J : %02d: ", j);
            for (int i = 0; i <imaxLocal + 2; i++) {
                switch ((int)S(i, j, k)) {
                case FRONTFACE:
                    printf("FF  ");
                    break;
                case BACKFACE:
                    printf("BF  ");
                    break;
                case LEFTFACE:
                    printf("LF  ");
                    break;
                case RIGHTFACE:
                    printf("RF  ");
                    break;
                case TOPFACE:
                    printf("TF  ");
                    break;
                case BOTTOMFACE:
                    printf("BMF  ");
                    break;

                case FRONTTOPLEFTCORNER:
                    printf("FTLC  ");
                    break;
                case FRONTTOPRIGHTCORNER:
                    printf("FTRC  ");
                    break;
                case FRONTBOTTOMLEFTCORNER:
                    printf("FBLC  ");
                    break;
                case FRONTBOTTOMRIGHTCORNER:
                    printf("FBRC  ");
                    break;
                case BACKTOPLEFTCORNER:
                    printf("BTLC  ");
                    break;
                case BACKTOPRIGHTCORNER:
                    printf("BTRC  ");
                    break;
                case BACKBOTTOMLEFTCORNER:
                    printf("BBLC  ");
                    break;
                case BACKBOTTOMRIGHTCORNER:
                    printf("BBRC  ");
                    break;

                case FRONTTOPLINE:
                    printf("FTL  ");
                    break;
                case FRONTBOTTOMLINE:
                    printf("FBL  ");
                    break;
                case FRONTLEFTLINE:
                    printf("FLL  ");
                    break;
                case FRONTRIGHTLINE:
                    printf("FRL  ");
                    break;
                case MIDTOPLEFTLINE:
                    printf("MTLL  ");
                    break;
                case MIDTOPRIGHTLINE:
                    printf("MTRL  ");
                    break;
                case MIDBOTTOMLEFTLINE:
                    printf("MBTL  ");
                    break;
                case MIDBOTTOMRIGHTLINE:
                    printf("MBRL  ");
                    break;
                case BACKTOPLINE:
                    printf("BTL  ");
                    break;
                case BACKBOTTOMLINE:
                    printf("BBL  ");
                    break;
                case BACKLEFTLINE:
                    printf("BLL  ");
                    break;
                case BACKRIGHTLINE:
                    printf("BRL  ");
                    break;

                case OBSTACLE:
                    printf("O  ");
                    break;
                case FLUID:
                    printf("N  ");
                    break;
                }
            }
            printf("\n");
        }
        printf("\n\n");
    }
    fflush(stdout);
}

void initDiscretization(Discretization* d, Parameter* params)
{
    d->problem  = params->name;
    d->bcLeft   = params->bcLeft;
    d->bcRight  = params->bcRight;
    d->bcBottom = params->bcBottom;
    d->bcTop    = params->bcTop;
    d->bcFront  = params->bcFront;
    d->bcBack   = params->bcBack;

    d->grid.imax    = params->imax;
    d->grid.jmax    = params->jmax;
    d->grid.kmax    = params->kmax;
    d->grid.xlength = params->xlength;
    d->grid.ylength = params->ylength;
    d->grid.zlength = params->zlength;
    d->grid.dx      = params->xlength / params->imax;
    d->grid.dy      = params->ylength / params->jmax;
    d->grid.dz      = params->zlength / params->kmax;

    d->eps     = params->eps;
    d->omega   = params->omg;
    d->itermax = params->itermax;
    d->re      = params->re;
    d->gx      = params->gx;
    d->gy      = params->gy;
    d->gz      = params->gz;
    d->dt      = params->dt;
    d->te      = params->te;
    d->tau     = params->tau;
    d->gamma   = params->gamma;

    /* allocate arrays */
    int imaxLocal = d->comm.imaxLocal;
    int jmaxLocal = d->comm.jmaxLocal;
    int kmaxLocal = d->comm.kmaxLocal;
    size_t size   = (imaxLocal + 2) * (jmaxLocal + 2) * (kmaxLocal + 2);

    d->u   = allocate(64, size * sizeof(double));
    d->v   = allocate(64, size * sizeof(double));
    d->w   = allocate(64, size * sizeof(double));
    d->p   = allocate(64, size * sizeof(double));
    d->rhs = allocate(64, size * sizeof(double));
    d->f   = allocate(64, size * sizeof(double));
    d->g   = allocate(64, size * sizeof(double));
    d->h   = allocate(64, size * sizeof(double));
    d->s   = allocate(64, size * sizeof(double));

    for (int i = 0; i < size; i++) {
        d->u[i]   = params->u_init;
        d->v[i]   = params->v_init;
        d->w[i]   = params->w_init;
        d->p[i]   = params->p_init;
        d->rhs[i] = 0.0;
        d->f[i]   = 0.0;
        d->g[i]   = 0.0;
        d->h[i]   = 0.0;
        d->s[i]   = FLUID;
    }

    double dx = d->grid.dx;
    double dy = d->grid.dy;
    double dz = d->grid.dz;

    double invSqrSum = 1.0 / (dx * dx) + 1.0 / (dy * dy) + 1.0 / (dz * dz);
    d->dtBound       = 0.5 * d->re * 1.0 / invSqrSum;

    double xCenter = 0, yCenter = 0, zCenter = 0, radius = 0;
    double x1 = 0, x2 = 0, y1 = 0, y2 = 0, z1 = 0, z2 = 0;
    int iOffset = 0, jOffset = 0, kOffset = 0;

    d->xLocal = d->comm.imaxLocal * d->grid.dx;
    d->yLocal = d->comm.jmaxLocal * d->grid.dy;
    d->zLocal = d->comm.kmaxLocal * d->grid.dz;

#ifdef _MPI
    double xLocal[d->comm.size];
    double yLocal[d->comm.size];
    double zLocal[d->comm.size];

    MPI_Allgather(&d->xLocal, 1, MPI_DOUBLE, xLocal, 1, MPI_DOUBLE, d->comm.comm);
    MPI_Allgather(&d->yLocal, 1, MPI_DOUBLE, yLocal, 1, MPI_DOUBLE, d->comm.comm);
    MPI_Allgather(&d->zLocal, 1, MPI_DOUBLE, zLocal, 1, MPI_DOUBLE, d->comm.comm);

    d->xOffset    = sumOffset(xLocal,
        d->comm.rank,
        (d->comm.dims[1] * d->comm.dims[2]),
        d->comm.coords[0]);
    d->yOffset    = sumOffset(yLocal, d->comm.rank, d->comm.dims[2], d->comm.coords[1]);
    d->zOffset    = sumOffset(zLocal, d->comm.rank, 1, d->comm.coords[2]);
    d->xOffsetEnd = d->xOffset + d->xLocal;
    d->yOffsetEnd = d->yOffset + d->yLocal;
    d->zOffsetEnd = d->zOffset + d->zLocal;

#else
    d->xOffset    = 0;
    d->yOffset    = 0;
    d->zOffset    = 0;
    d->xOffsetEnd = d->xOffset + d->xLocal;
    d->yOffsetEnd = d->yOffset + d->yLocal;
    d->zOffsetEnd = d->zOffset + d->zLocal;
#endif

    iOffset = d->xOffset / dx;
    jOffset = d->yOffset / dy;
    kOffset = d->zOffset / dz;

    double* s = d->s;

    switch (params->shape) {
    case NOSHAPE:
        break;
    case RECT:
        x1 = params->xCenter - params->xRectLength / 2;
        x2 = params->xCenter + params->xRectLength / 2;
        y1 = params->yCenter - params->yRectLength / 2;
        y2 = params->yCenter + params->yRectLength / 2;
        z1 = params->zCenter - params->zRectLength / 2;
        z2 = params->zCenter + params->zRectLength / 2;

        for (int k = 1; k < kmaxLocal + 1; k++) {
            for (int j = 1; j < jmaxLocal + 1; j++) {
                for (int i = 1; i < imaxLocal + 1; i++) {
                    if ((x1 <= ((i + iOffset) * dx)) && (((i + iOffset) * dx) <= x2) &&
                        (y1 <= ((j + jOffset) * dy)) && (((j + jOffset) * dy) <= y2) &&
                        ((z1 <= ((k + kOffset) * dz)) && (((k + kOffset) * dz) <= z2))) {
                        S(i, j, k) = OBSTACLE;
                    }
                }
            }
        }

        break;
    case CIRCLE:

        xCenter = params->xCenter;
        yCenter = params->yCenter;
        zCenter = params->zCenter;
        radius  = params->circleRadius;

        for (int k = 1; k < kmaxLocal + 1; k++) {
            for (int j = 1; j < jmaxLocal + 1; j++) {
                for (int i = 1; i < imaxLocal + 1; i++) {
                    if (distance(((i + iOffset) * dx),
                            ((j + jOffset) * dy),
                            ((k + kOffset) * dz),
                            xCenter,
                            yCenter,
                            zCenter) <= radius) {
                        S(i, j, k) = OBSTACLE;
                    }
                }
            }
        }
        break;
    default:
        break;
    }

#ifdef _MPI
    commExchange(&d->comm, s);
#endif

    for (int k = 1; k < kmaxLocal + 1; k++) {
        for (int j = 1; j < jmaxLocal + 1; j++) {
            for (int i = 1; i < imaxLocal + 1; i++) {

                /* Assigning enum values to Corners */
                if (S(i - 1, j + 1, k - 1) == FLUID && S(i - 1, j, k) == FLUID &&
                    S(i, j + 1, k) == FLUID && S(i, j, k - 1) == FLUID &&
                    S(i + 1, j - 1, k + 1) == OBSTACLE && S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = FRONTTOPLEFTCORNER; //
                }
                if (S(i + 1, j + 1, k - 1) == FLUID && S(i + 1, j, k) == FLUID &&
                    S(i, j + 1, k) == FLUID && S(i, j, k - 1) == FLUID &&
                    S(i - 1, j - 1, k + 1) == OBSTACLE && S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = FRONTTOPRIGHTCORNER; //
                }
                if (S(i - 1, j - 1, k - 1) == FLUID && S(i - 1, j, k) == FLUID &&
                    S(i, j - 1, k) == FLUID && S(i, j, k - 1) == FLUID &&
                    S(i + 1, j + 1, k + 1) == OBSTACLE && S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = FRONTBOTTOMLEFTCORNER; //
                }
                if (S(i + 1, j - 1, k - 1) == FLUID && S(i + 1, j, k) == FLUID &&
                    S(i, j - 1, k) == FLUID && S(i, j, k - 1) == FLUID &&
                    S(i - 1, j + 1, k + 1) == OBSTACLE && S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = FRONTBOTTOMRIGHTCORNER; //
                }
                if (S(i - 1, j + 1, k + 1) == FLUID && S(i - 1, j, k) == FLUID &&
                    S(i, j + 1, k) == FLUID && S(i, j, k + 1) == FLUID &&
                    S(i + 1, j - 1, k - 1) == OBSTACLE && S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = BACKTOPLEFTCORNER; //
                }
                if (S(i + 1, j + 1, k + 1) == FLUID && S(i + 1, j, k) == FLUID &&
                    S(i, j + 1, k) == FLUID && S(i, j, k + 1) == FLUID &&
                    S(i - 1, j - 1, k - 1) == OBSTACLE && S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = BACKTOPRIGHTCORNER;
                }
                if (S(i - 1, j - 1, k + 1) == FLUID && S(i - 1, j, k) == FLUID &&
                    S(i, j - 1, k) == FLUID && S(i, j, k + 1) == FLUID &&
                    S(i + 1, j + 1, k - 1) == OBSTACLE && S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = BACKBOTTOMLEFTCORNER;
                }
                if (S(i + 1, j - 1, k + 1) == FLUID && S(i + 1, j, k) == FLUID &&
                    S(i, j - 1, k) == FLUID && S(i, j, k + 1) == FLUID &&
                    S(i - 1, j + 1, k - 1) == OBSTACLE && S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = BACKBOTTOMRIGHTCORNER;
                }
                /* Assigning enum values to Lines */
                if (S(i - 1, j, k - 1) == FLUID && S(i, j, k - 1) == FLUID &&
                    S(i - 1, j, k) == FLUID && S(i + 1, j, k + 1) == OBSTACLE &&
                    S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = FRONTLEFTLINE;
                }
                if (S(i + 1, j, k - 1) == FLUID && S(i + 1, j, k) == FLUID &&
                    S(i, j, k - 1) == FLUID && S(i - 1, j, k + 1) == OBSTACLE &&
                    S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = FRONTRIGHTLINE;
                }
                if (S(i, j + 1, k - 1) == FLUID && S(i, j + 1, k) == FLUID &&
                    S(i, j, k - 1) == FLUID && S(i, j - 1, k + 1) == OBSTACLE &&
                    S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = FRONTTOPLINE;
                }
                if (S(i, j - 1, k - 1) == FLUID && S(i, j, k - 1) == FLUID &&
                    S(i, j - 1, k) == FLUID && S(i, j + 1, k + 1) == OBSTACLE &&
                    S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = FRONTBOTTOMLINE;
                }
                if (S(i - 1, j + 1, k) == FLUID && S(i, j + 1, k) == FLUID &&
                    S(i - 1, j, k) == FLUID && S(i + 1, j - 1, k) == OBSTACLE &&
                    S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = MIDTOPLEFTLINE;
                }
                if (S(i + 1, j + 1, k) == FLUID && S(i + 1, j, k) == FLUID &&
                    S(i, j + 1, k) == FLUID && S(i - 1, j - 1, k) == OBSTACLE &&
                    S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = MIDTOPRIGHTLINE;
                }
                if (S(i - 1, j - 1, k) == FLUID && S(i - 1, j, k) == FLUID &&
                    S(i, j - 1, k) == FLUID && S(i + 1, j + 1, k) == OBSTACLE &&
                    S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = MIDBOTTOMLEFTLINE;
                }
                if (S(i + 1, j - 1, k) == FLUID && S(i + 1, j, k) == FLUID &&
                    S(i, j - 1, k) == FLUID && S(i - 1, j + 1, k) == OBSTACLE &&
                    S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = MIDBOTTOMRIGHTLINE;
                }
                if (S(i - 1, j, k + 1) == FLUID && S(i - 1, j, k) == FLUID &&
                    S(i, j, k + 1) == FLUID && S(i + 1, j, k - 1) == OBSTACLE &&
                    S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = BACKLEFTLINE;
                }
                if (S(i + 1, j, k + 1) == FLUID && S(i + 1, j, k) == FLUID &&
                    S(i, j, k + 1) == FLUID && S(i - 1, j, k - 1) == OBSTACLE &&
                    S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = BACKRIGHTLINE;
                }
                if (S(i, j + 1, k + 1) == FLUID && S(i, j + 1, k) == FLUID &&
                    S(i, j, k + 1) == FLUID && S(i, j - 1, k - 1) == OBSTACLE &&
                    S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = BACKTOPLINE;
                }
                if (S(i, j - 1, k + 1) == FLUID && S(i, j - 1, k) == FLUID &&
                    S(i, j, k + 1) == FLUID && S(i, j + 1, k - 1) == OBSTACLE &&
                    S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = BACKBOTTOMLINE;
                }
                /* Assigning enum values to Faces */
                if (S(i, j, k - 1) == FLUID && S(i, j, k + 1) == OBSTACLE &&
                    S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = FRONTFACE; //
                }
                if (S(i, j, k + 1) == FLUID && S(i, j, k - 1) == OBSTACLE &&
                    S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = BACKFACE; //
                }
                if (S(i, j - 1, k) == FLUID && S(i, j + 1, k) == OBSTACLE &&
                    S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = BOTTOMFACE; //
                }
                if (S(i, j + 1, k) == FLUID && S(i, j - 1, k) == OBSTACLE &&
                    S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = TOPFACE; //
                }
                if (S(i - 1, j, k) == FLUID && S(i + 1, j, k) == OBSTACLE &&
                    S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = LEFTFACE; //
                }
                if (S(i + 1, j, k) == FLUID && S(i - 1, j, k) == OBSTACLE &&
                    S(i, j, k) == OBSTACLE) {
                    S(i, j, k) = RIGHTFACE; //
                }
            }
        }
    }

#ifdef VERBOSE
    printConfig(d);
#endif /* VERBOSE */
}

void setBoundaryConditions(Discretization* d)
{
    int imaxLocal = d->comm.imaxLocal;
    int jmaxLocal = d->comm.jmaxLocal;
    int kmaxLocal = d->comm.kmaxLocal;

    double* u = d->u;
    double* v = d->v;
    double* w = d->w;

    if (commIsBoundary(&d->comm, TOP)) {
        switch (d->bcTop) {
        case NOSLIP:
            for (int k = 1; k < kmaxLocal + 1; k++) {
                for (int i = 1; i < imaxLocal + 1; i++) {
                    U(i, jmaxLocal + 1, k) = -U(i, jmaxLocal, k);
                    V(i, jmaxLocal, k)     = 0.0;
                    W(i, jmaxLocal + 1, k) = -W(i, jmaxLocal, k);
                }
            }
            break;
        case SLIP:
            for (int k = 1; k < kmaxLocal + 1; k++) {
                for (int i = 1; i < imaxLocal + 1; i++) {
                    U(i, jmaxLocal + 1, k) = U(i, jmaxLocal, k);
                    V(i, jmaxLocal, k)     = 0.0;
                    W(i, jmaxLocal + 1, k) = W(i, jmaxLocal, k);
                }
            }
            break;
        case OUTFLOW:
            for (int k = 1; k < kmaxLocal + 1; k++) {
                for (int i = 1; i < imaxLocal + 1; i++) {
                    U(i, jmaxLocal + 1, k) = U(i, jmaxLocal, k);
                    V(i, jmaxLocal, k)     = V(i, jmaxLocal - 1, k);
                    W(i, jmaxLocal + 1, k) = W(i, jmaxLocal, k);
                }
            }
            break;
        case PERIODIC:
            break;
        }
    }

    if (commIsBoundary(&d->comm, BOTTOM)) {
        switch (d->bcBottom) {
        case NOSLIP:
            for (int k = 1; k < kmaxLocal + 1; k++) {
                for (int i = 1; i < imaxLocal + 1; i++) {
                    U(i, 0, k) = -U(i, 1, k);
                    V(i, 0, k) = 0.0;
                    W(i, 0, k) = -W(i, 1, k);
                }
            }
            break;
        case SLIP:
            for (int k = 1; k < kmaxLocal + 1; k++) {
                for (int i = 1; i < imaxLocal + 1; i++) {
                    U(i, 0, k) = U(i, 1, k);
                    V(i, 0, k) = 0.0;
                    W(i, 0, k) = W(i, 1, k);
                }
            }
            break;
        case OUTFLOW:
            for (int k = 1; k < kmaxLocal + 1; k++) {
                for (int i = 1; i < imaxLocal + 1; i++) {
                    U(i, 0, k) = U(i, 1, k);
                    V(i, 0, k) = V(i, 1, k);
                    W(i, 0, k) = W(i, 1, k);
                }
            }
            break;
        case PERIODIC:
            break;
        }
    }

    if (commIsBoundary(&d->comm, LEFT)) {
        switch (d->bcLeft) {
        case NOSLIP:
            for (int k = 1; k < kmaxLocal + 1; k++) {
                for (int j = 1; j < jmaxLocal + 1; j++) {
                    U(0, j, k) = 0.0;
                    V(0, j, k) = -V(1, j, k);
                    W(0, j, k) = -W(1, j, k);
                }
            }
            break;
        case SLIP:
            for (int k = 1; k < kmaxLocal + 1; k++) {
                for (int j = 1; j < jmaxLocal + 1; j++) {
                    U(0, j, k) = 0.0;
                    V(0, j, k) = V(1, j, k);
                    W(0, j, k) = W(1, j, k);
                }
            }
            break;
        case OUTFLOW:
            for (int k = 1; k < kmaxLocal + 1; k++) {
                for (int j = 1; j < jmaxLocal + 1; j++) {
                    U(0, j, k) = U(1, j, k);
                    V(0, j, k) = V(1, j, k);
                    W(0, j, k) = W(1, j, k);
                }
            }
            break;
        case PERIODIC:
            break;
        }
    }

    if (commIsBoundary(&d->comm, RIGHT)) {
        switch (d->bcRight) {
        case NOSLIP:
            for (int k = 1; k < kmaxLocal + 1; k++) {
                for (int j = 1; j < jmaxLocal + 1; j++) {
                    U(imaxLocal, j, k)     = 0.0;
                    V(imaxLocal + 1, j, k) = -V(imaxLocal, j, k);
                    W(imaxLocal + 1, j, k) = -W(imaxLocal, j, k);
                }
            }
            break;
        case SLIP:
            for (int k = 1; k < kmaxLocal + 1; k++) {
                for (int j = 1; j < jmaxLocal + 1; j++) {
                    U(imaxLocal, j, k)     = 0.0;
                    V(imaxLocal + 1, j, k) = V(imaxLocal, j, k);
                    W(imaxLocal + 1, j, k) = W(imaxLocal, j, k);
                }
            }
            break;
        case OUTFLOW:
            for (int k = 1; k < kmaxLocal + 1; k++) {
                for (int j = 1; j < jmaxLocal + 1; j++) {
                    U(imaxLocal, j, k)     = U(imaxLocal - 1, j, k);
                    V(imaxLocal + 1, j, k) = V(imaxLocal, j, k);
                    W(imaxLocal + 1, j, k) = W(imaxLocal, j, k);
                }
            }
            break;
        case PERIODIC:
            break;
        }
    }

    if (commIsBoundary(&d->comm, FRONT)) {
        switch (d->bcFront) {
        case NOSLIP:
            for (int j = 1; j < jmaxLocal + 1; j++) {
                for (int i = 1; i < imaxLocal + 1; i++) {
                    U(i, j, 0) = -U(i, j, 1);
                    V(i, j, 0) = -V(i, j, 1);
                    W(i, j, 0) = 0.0;
                }
            }
            break;
        case SLIP:
            for (int j = 1; j < jmaxLocal + 1; j++) {
                for (int i = 1; i < imaxLocal + 1; i++) {
                    U(i, j, 0) = U(i, j, 1);
                    V(i, j, 0) = V(i, j, 1);
                    W(i, j, 0) = 0.0;
                }
            }
            break;
        case OUTFLOW:
            for (int j = 1; j < jmaxLocal + 1; j++) {
                for (int i = 1; i < imaxLocal + 1; i++) {
                    U(i, j, 0) = U(i, j, 1);
                    V(i, j, 0) = V(i, j, 1);
                    W(i, j, 0) = W(i, j, 1);
                }
            }
            break;
        case PERIODIC:
            break;
        }
    }

    if (commIsBoundary(&d->comm, BACK)) {
        switch (d->bcBack) {
        case NOSLIP:
            for (int j = 1; j < jmaxLocal + 1; j++) {
                for (int i = 1; i < imaxLocal + 1; i++) {
                    U(i, j, kmaxLocal + 1) = -U(i, j, kmaxLocal);
                    V(i, j, kmaxLocal + 1) = -V(i, j, kmaxLocal);
                    W(i, j, kmaxLocal)     = 0.0;
                }
            }
            break;
        case SLIP:
            for (int j = 1; j < jmaxLocal + 1; j++) {
                for (int i = 1; i < imaxLocal + 1; i++) {
                    U(i, j, kmaxLocal + 1) = U(i, j, kmaxLocal);
                    V(i, j, kmaxLocal + 1) = V(i, j, kmaxLocal);
                    W(i, j, kmaxLocal)     = 0.0;
                }
            }
            break;
        case OUTFLOW:
            for (int j = 1; j < jmaxLocal + 1; j++) {
                for (int i = 1; i < imaxLocal + 1; i++) {
                    U(i, j, kmaxLocal + 1) = U(i, j, kmaxLocal);
                    V(i, j, kmaxLocal + 1) = V(i, j, kmaxLocal);
                    W(i, j, kmaxLocal)     = W(i, j, kmaxLocal - 1);
                }
            }
            break;
        case PERIODIC:
            break;
        }
    }
}

void computeRHS(Discretization* d)
{
    int imaxLocal = d->comm.imaxLocal;
    int jmaxLocal = d->comm.jmaxLocal;
    int kmaxLocal = d->comm.kmaxLocal;

    double idx = 1.0 / d->grid.dx;
    double idy = 1.0 / d->grid.dy;
    double idz = 1.0 / d->grid.dz;
    double idt = 1.0 / d->dt;

    double* rhs = d->rhs;
    double* f   = d->f;
    double* g   = d->g;
    double* h   = d->h;

    commShift(&d->comm, f, g, h);

    for (int k = 1; k < kmaxLocal + 1; k++) {
        for (int j = 1; j < jmaxLocal + 1; j++) {
            for (int i = 1; i < imaxLocal + 1; i++) {
                RHS(i, j, k) = ((F(i, j, k) - F(i - 1, j, k)) * idx +
                                   (G(i, j, k) - G(i, j - 1, k)) * idy +
                                   (H(i, j, k) - H(i, j, k - 1)) * idz) *
                               idt;
            }
        }
    }
}

void setSpecialBoundaryCondition(Discretization* d)
{
    int imaxLocal = d->comm.imaxLocal;
    int jmaxLocal = d->comm.jmaxLocal;
    int kmaxLocal = d->comm.kmaxLocal;

    double* u = d->u;

    if (strcmp(d->problem, "dcavity") == 0) {
        if (commIsBoundary(&d->comm, TOP)) {
            for (int k = 1; k < kmaxLocal; k++) {
                for (int i = 1; i < imaxLocal; i++) {
                    U(i, jmaxLocal + 1, k) = 2.0 - U(i, jmaxLocal, k);
                }
            }
        }
    } else if (strcmp(d->problem, "canal") == 0) {
        if (commIsBoundary(&d->comm, LEFT)) {
            for (int k = 1; k < kmaxLocal + 1; k++) {
                for (int j = 1; j < jmaxLocal + 1; j++) {
                    U(0, j, k) = 2.0;
                }
            }
        }
    } else if (strcmp(d->problem, "backstep") == 0) {
        if (commIsBoundary(&d->comm, LEFT)) {
            double* s = d->s;
            for (int k = 1; k < kmaxLocal + 1; k++) {
                for (int j = 1; j < jmaxLocal + 1; j++) {
                    if (S(1, j, k) == FLUID) U(0, j, k) = 1.0;
                    else {
                        U(0, j, k) = 0.0;
                        U(1, j, k) = 0.0;
                    }
                }
            }
        }
    } else if (strcmp(d->problem, "karman") == 0) {
        if (commIsBoundary(&d->comm, LEFT)) {
            for (int k = 1; k < kmaxLocal + 1; k++) {
                for (int j = 1; j < jmaxLocal + 1; j++) {
                    U(1, j, k) = 1.0;
                }
            }
        }
    }
}

static double maxElement(Discretization* d, double* m)
{
    int size = (d->comm.imaxLocal + 2) * (d->comm.jmaxLocal + 2) *
               (d->comm.kmaxLocal + 2);
    double maxval = DBL_MIN;

    for (int i = 0; i < size; i++) {
        maxval = MAX(maxval, fabs(m[i]));
    }
    commReduction(&maxval, MAX);
    return maxval;
}

void normalizePressure(Discretization* d)
{
    int imaxLocal = d->comm.imaxLocal;
    int jmaxLocal = d->comm.jmaxLocal;
    int kmaxLocal = d->comm.kmaxLocal;

    double* p   = d->p;
    double avgP = 0.0;

    for (int k = 1; k < kmaxLocal + 1; k++) {
        for (int j = 1; j < jmaxLocal + 1; j++) {
            for (int i = 1; i < imaxLocal + 1; i++) {
                avgP += P(i, j, k);
            }
        }
    }
    commReduction(&avgP, SUM);
    avgP /= (d->grid.imax * d->grid.jmax * d->grid.kmax);

    for (int k = 1; k < kmaxLocal + 1; k++) {
        for (int j = 1; j < jmaxLocal + 1; j++) {
            for (int i = 1; i < imaxLocal + 1; i++) {
                P(i, j, k) = P(i, j, k) - avgP;
            }
        }
    }
}

void computeTimestep(Discretization* d)
{
    double dt = d->dtBound;
    double dx = d->grid.dx;
    double dy = d->grid.dy;
    double dz = d->grid.dz;

    double umax = maxElement(d, d->u);
    double vmax = maxElement(d, d->v);
    double wmax = maxElement(d, d->w);

    if (umax > 0) {
        dt = (dt > dx / umax) ? dx / umax : dt;
    }
    if (vmax > 0) {
        dt = (dt > dy / vmax) ? dy / vmax : dt;
    }
    if (wmax > 0) {
        dt = (dt > dz / wmax) ? dz / wmax : dt;
    }

    d->dt = dt * d->tau;
}

void setObjectBoundaryCondition(Discretization* d)
{
    int imaxLocal = d->comm.imaxLocal;
    int jmaxLocal = d->comm.jmaxLocal;
    int kmaxLocal = d->comm.kmaxLocal;
    double* u     = d->u;
    double* v     = d->v;
    double* w     = d->w;
    double* s     = d->s;
    for (int k = 1; k < kmaxLocal + 1; k++) {
        for (int j = 1; j < jmaxLocal + 1; j++) {
            for (int i = 1; i < imaxLocal + 1; i++) {
                switch ((int)S(i, j, k)) {
                case TOPFACE:
                    U(i, j, k) = -U(i, j + 1, k);
                    V(i, j, k) = 0.0;
                    W(i, j, k) = -W(i, j + 1, k);
                    ;
                    break;
                case BOTTOMFACE:
                    U(i, j, k) = -U(i, j - 1, k);
                    V(i, j, k) = 0.0;
                    W(i, j, k) = -W(i, j - 1, k);
                    break;
                case LEFTFACE:
                    U(i, j, k) = 0.0;
                    V(i, j, k) = -V(i - 1, j, k);
                    W(i, j, k) = -W(i - 1, j, k);
                    break;
                case RIGHTFACE:
                    U(i, j, k) = 0.0;
                    V(i, j, k) = -V(i + 1, j, k);
                    W(i, j, k) = -W(i + 1, j, k);
                    break;
                case FRONTFACE:
                    U(i, j, k) = -U(i, j, k - 1);
                    V(i, j, k) = -V(i, j, k - 1);
                    W(i, j, k) = 0.0;
                    break;
                case BACKFACE:
                    U(i, j, k) = -U(i, j, k + 1);
                    V(i, j, k) = -V(i, j, k + 1);
                    W(i, j, k) = 0.0;
                    break;
                case FRONTLEFTLINE:
                    U(i, j, k) = 0.0;
                    V(i, j, k) = -V(i, j, k - 1);
                    W(i, j, k) = -W(i - 1, j, k);
                    break;
                case FRONTRIGHTLINE:
                    U(i, j, k) = 0.0;
                    V(i, j, k) = -V(i, j, k - 1);
                    W(i, j, k) = -W(i + 1, j, k);
                    break;
                case FRONTTOPLINE:
                    U(i, j, k) = -U(i, j, k - 1);
                    V(i, j, k) = 0.0;
                    W(i, j, k) = -W(i, j + 1, k);
                    break;
                case FRONTBOTTOMLINE:
                    U(i, j, k) = -U(i, j, k - 1);
                    V(i, j, k) = 0.0;
                    W(i, j, k) = -W(i, j - 1, k);
                    break;
                case MIDTOPLEFTLINE:
                    U(i, j, k) = -U(i, j + 1, k);
                    V(i, j, k) = -V(i - 1, j, k);
                    W(i, j, k) = 0.0;
                    break;
                case MIDTOPRIGHTLINE:
                    U(i, j, k)     = 0.0;
                    V(i, j, k)     = 0.0;
                    U(i - 1, j, k) = -U(i - 1, j + 1, k);
                    V(i, j - 1, k) = -V(i + 1, j - 1, k);
                    W(i, j, k)     = 0.0;
                    break;
                case MIDBOTTOMLEFTLINE:
                    U(i, j, k) = -U(i, j - 1, k);
                    V(i, j, k) = -V(i - 1, j, k);
                    W(i, j, k) = 0.0;
                    break;
                case MIDBOTTOMRIGHTLINE:
                    U(i, j, k) = -U(i, j - 1, k);
                    V(i, j, k) = -V(i + 1, j, k);
                    W(i, j, k) = 0.0;
                    break;
                case BACKLEFTLINE:
                    U(i, j, k) = 0.0;
                    V(i, j, k) = -V(i, j, k + 1);
                    W(i, j, k) = -W(i - 1, j, k);
                    break;
                case BACKRIGHTLINE:
                    U(i, j, k) = 0.0;
                    V(i, j, k) = -V(i, j, k + 1);
                    W(i, j, k) = -W(i + 1, j, k);
                    break;
                case BACKTOPLINE:
                    U(i, j, k) = -U(i, j, k + 1);
                    V(i, j, k) = 0.0;
                    W(i, j, k) = -W(i, j + 1, k);
                    break;
                case BACKBOTTOMLINE:
                    U(i, j, k) = -U(i, j, k + 1);
                    V(i, j, k) = 0.0;
                    W(i, j, k) = -W(i, j - 1, k);
                    break;
                case FRONTTOPLEFTCORNER:
                    U(i, j, k) = -U(i, j, k - 1);
                    V(i, j, k) = -V(i - 1, j, k);
                    W(i, j, k) = -W(i, j + 1, k);
                    break;
                case FRONTTOPRIGHTCORNER:
                    U(i, j, k) = -U(i, j, k - 1);
                    V(i, j, k) = -V(i + 1, j, k);
                    W(i, j, k) = -W(i, j + 1, k);
                    break;
                case FRONTBOTTOMLEFTCORNER:
                    U(i, j, k) = -U(i, j, k - 1);
                    V(i, j, k) = -V(i - 1, j, k);
                    W(i, j, k) = -W(i, j - 1, k);
                    break;
                case FRONTBOTTOMRIGHTCORNER:
                    U(i, j, k) = -U(i, j, k - 1);
                    V(i, j, k) = -V(i + 1, j, k);
                    W(i, j, k) = -W(i, j - 1, k);
                    break;
                case BACKTOPLEFTCORNER:
                    U(i, j, k) = -U(i, j, k + 1);
                    V(i, j, k) = -V(i - 1, j, k);
                    W(i, j, k) = -W(i, j + 1, k);
                    break;
                case BACKTOPRIGHTCORNER:
                    U(i, j, k) = -U(i, j, k + 1);
                    V(i, j, k) = -V(i + 1, j, k);
                    W(i, j, k) = -W(i, j + 1, k);
                    break;
                case BACKBOTTOMLEFTCORNER:
                    U(i, j, k) = -U(i, j, k + 1);
                    V(i, j, k) = -V(i - 1, j, k);
                    W(i, j, k) = -W(i, j - 1, k);
                    break;
                case BACKBOTTOMRIGHTCORNER:
                    U(i, j, k) = -U(i, j, k + 1);
                    V(i, j, k) = -V(i + 1, j, k);
                    W(i, j, k) = -W(i, j - 1, k);
                    break;
                }
            }
        }
    }
}

void computeFG(Discretization* d)
{
    int imaxLocal = d->comm.imaxLocal;
    int jmaxLocal = d->comm.jmaxLocal;
    int kmaxLocal = d->comm.kmaxLocal;

    double* u = d->u;
    double* v = d->v;
    double* w = d->w;
    double* f = d->f;
    double* g = d->g;
    double* h = d->h;
    double* s = d->s;

    double gx = d->gx;
    double gy = d->gy;
    double gz = d->gz;
    double dt = d->dt;

    double gamma     = d->gamma;
    double inverseRe = 1.0 / d->re;
    double inverseDx = 1.0 / d->grid.dx;
    double inverseDy = 1.0 / d->grid.dy;
    double inverseDz = 1.0 / d->grid.dz;
    double du2dx, dv2dy, dw2dz;
    double duvdx, duwdx, duvdy, dvwdy, duwdz, dvwdz;
    double du2dx2, du2dy2, du2dz2;
    double dv2dx2, dv2dy2, dv2dz2;
    double dw2dx2, dw2dy2, dw2dz2;

    commExchange(&d->comm, u);
    commExchange(&d->comm, v);
    commExchange(&d->comm, w);

    for (int k = 1; k < kmaxLocal + 1; k++) {
        for (int j = 1; j < jmaxLocal + 1; j++) {
            for (int i = 1; i < imaxLocal + 1; i++) {
                if (S(i, j, k) == FLUID) {
                    du2dx = inverseDx * 0.25 *
                                ((U(i, j, k) + U(i + 1, j, k)) *
                                        (U(i, j, k) + U(i + 1, j, k)) -
                                    (U(i, j, k) + U(i - 1, j, k)) *
                                        (U(i, j, k) + U(i - 1, j, k))) +
                            gamma * inverseDx * 0.25 *
                                (fabs(U(i, j, k) + U(i + 1, j, k)) *
                                        (U(i, j, k) - U(i + 1, j, k)) +
                                    fabs(U(i, j, k) + U(i - 1, j, k)) *
                                        (U(i, j, k) - U(i - 1, j, k)));

                    duvdy = inverseDy * 0.25 *
                                ((V(i, j, k) + V(i + 1, j, k)) *
                                        (U(i, j, k) + U(i, j + 1, k)) -
                                    (V(i, j - 1, k) + V(i + 1, j - 1, k)) *
                                        (U(i, j, k) + U(i, j - 1, k))) +
                            gamma * inverseDy * 0.25 *
                                (fabs(V(i, j, k) + V(i + 1, j, k)) *
                                        (U(i, j, k) - U(i, j + 1, k)) +
                                    fabs(V(i, j - 1, k) + V(i + 1, j - 1, k)) *
                                        (U(i, j, k) - U(i, j - 1, k)));

                    duwdz = inverseDz * 0.25 *
                                ((W(i, j, k) + W(i + 1, j, k)) *
                                        (U(i, j, k) + U(i, j, k + 1)) -
                                    (W(i, j, k - 1) + W(i + 1, j, k - 1)) *
                                        (U(i, j, k) + U(i, j, k - 1))) +
                            gamma * inverseDz * 0.25 *
                                (fabs(W(i, j, k) + W(i + 1, j, k)) *
                                        (U(i, j, k) - U(i, j, k + 1)) +
                                    fabs(W(i, j, k - 1) + W(i + 1, j, k - 1)) *
                                        (U(i, j, k) - U(i, j, k - 1)));

                    du2dx2 = inverseDx * inverseDx *
                             (U(i + 1, j, k) - 2.0 * U(i, j, k) + U(i - 1, j, k));
                    du2dy2 = inverseDy * inverseDy *
                             (U(i, j + 1, k) - 2.0 * U(i, j, k) + U(i, j - 1, k));
                    du2dz2 = inverseDz * inverseDz *
                             (U(i, j, k + 1) - 2.0 * U(i, j, k) + U(i, j, k - 1));
                    F(i, j, k) = U(i, j, k) +
                                 dt * (inverseRe * (du2dx2 + du2dy2 + du2dz2) - du2dx -
                                          duvdy - duwdz + gx);

                    duvdx = inverseDx * 0.25 *
                                ((U(i, j, k) + U(i, j + 1, k)) *
                                        (V(i, j, k) + V(i + 1, j, k)) -
                                    (U(i - 1, j, k) + U(i - 1, j + 1, k)) *
                                        (V(i, j, k) + V(i - 1, j, k))) +
                            gamma * inverseDx * 0.25 *
                                (fabs(U(i, j, k) + U(i, j + 1, k)) *
                                        (V(i, j, k) - V(i + 1, j, k)) +
                                    fabs(U(i - 1, j, k) + U(i - 1, j + 1, k)) *
                                        (V(i, j, k) - V(i - 1, j, k)));

                    dv2dy = inverseDy * 0.25 *
                                ((V(i, j, k) + V(i, j + 1, k)) *
                                        (V(i, j, k) + V(i, j + 1, k)) -
                                    (V(i, j, k) + V(i, j - 1, k)) *
                                        (V(i, j, k) + V(i, j - 1, k))) +
                            gamma * inverseDy * 0.25 *
                                (fabs(V(i, j, k) + V(i, j + 1, k)) *
                                        (V(i, j, k) - V(i, j + 1, k)) +
                                    fabs(V(i, j, k) + V(i, j - 1, k)) *
                                        (V(i, j, k) - V(i, j - 1, k)));

                    dvwdz = inverseDz * 0.25 *
                                ((W(i, j, k) + W(i, j + 1, k)) *
                                        (V(i, j, k) + V(i, j, k + 1)) -
                                    (W(i, j, k - 1) + W(i, j + 1, k - 1)) *
                                        (V(i, j, k) + V(i, j, k + 1))) +
                            gamma * inverseDz * 0.25 *
                                (fabs(W(i, j, k) + W(i, j + 1, k)) *
                                        (V(i, j, k) - V(i, j, k + 1)) +
                                    fabs(W(i, j, k - 1) + W(i, j + 1, k - 1)) *
                                        (V(i, j, k) - V(i, j, k + 1)));

                    dv2dx2 = inverseDx * inverseDx *
                             (V(i + 1, j, k) - 2.0 * V(i, j, k) + V(i - 1, j, k));
                    dv2dy2 = inverseDy * inverseDy *
                             (V(i, j + 1, k) - 2.0 * V(i, j, k) + V(i, j - 1, k));
                    dv2dz2 = inverseDz * inverseDz *
                             (V(i, j, k + 1) - 2.0 * V(i, j, k) + V(i, j, k - 1));
                    G(i, j, k) = V(i, j, k) +
                                 dt * (inverseRe * (dv2dx2 + dv2dy2 + dv2dz2) - duvdx -
                                          dv2dy - dvwdz + gy);

                    duwdx = inverseDx * 0.25 *
                                ((U(i, j, k) + U(i, j, k + 1)) *
                                        (W(i, j, k) + W(i + 1, j, k)) -
                                    (U(i - 1, j, k) + U(i - 1, j, k + 1)) *
                                        (W(i, j, k) + W(i - 1, j, k))) +
                            gamma * inverseDx * 0.25 *
                                (fabs(U(i, j, k) + U(i, j, k + 1)) *
                                        (W(i, j, k) - W(i + 1, j, k)) +
                                    fabs(U(i - 1, j, k) + U(i - 1, j, k + 1)) *
                                        (W(i, j, k) - W(i - 1, j, k)));

                    dvwdy = inverseDy * 0.25 *
                                ((V(i, j, k) + V(i, j, k + 1)) *
                                        (W(i, j, k) + W(i, j + 1, k)) -
                                    (V(i, j - 1, k + 1) + V(i, j - 1, k)) *
                                        (W(i, j, k) + W(i, j - 1, k))) +
                            gamma * inverseDy * 0.25 *
                                (fabs(V(i, j, k) + V(i, j, k + 1)) *
                                        (W(i, j, k) - W(i, j + 1, k)) +
                                    fabs(V(i, j - 1, k + 1) + V(i, j - 1, k)) *
                                        (W(i, j, k) - W(i, j - 1, k)));

                    dw2dz = inverseDz * 0.25 *
                                ((W(i, j, k) + W(i, j, k + 1)) *
                                        (W(i, j, k) + W(i, j, k + 1)) -
                                    (W(i, j, k) + W(i, j, k - 1)) *
                                        (W(i, j, k) + W(i, j, k - 1))) +
                            gamma * inverseDz * 0.25 *
                                (fabs(W(i, j, k) + W(i, j, k + 1)) *
                                        (W(i, j, k) - W(i, j, k + 1)) +
                                    fabs(W(i, j, k) + W(i, j, k - 1)) *
                                        (W(i, j, k) - W(i, j, k - 1)));

                    dw2dx2 = inverseDx * inverseDx *
                             (W(i + 1, j, k) - 2.0 * W(i, j, k) + W(i - 1, j, k));
                    dw2dy2 = inverseDy * inverseDy *
                             (W(i, j + 1, k) - 2.0 * W(i, j, k) + W(i, j - 1, k));
                    dw2dz2 = inverseDz * inverseDz *
                             (W(i, j, k + 1) - 2.0 * W(i, j, k) + W(i, j, k - 1));
                    H(i, j, k) = W(i, j, k) +
                                 dt * (inverseRe * (dw2dx2 + dw2dy2 + dw2dz2) - duwdx -
                                          dvwdy - dw2dz + gz);
                } else {
                    switch ((int)S(i, j, k)) {
                    case TOPFACE:
                        G(i, j, k) = V(i, j, k);
                        break;
                    case BOTTOMFACE:
                        G(i, j, k) = V(i, j, k);
                        break;
                    case LEFTFACE:
                        F(i, j, k) = U(i, j, k);
                        break;
                    case RIGHTFACE:
                        F(i, j, k) = U(i, j, k);
                        break;
                    case FRONTFACE:
                        H(i, j, k) = W(i, j, k);
                        break;
                    case BACKFACE:
                        H(i, j, k) = W(i, j, k);
                        break;
                    case FRONTLEFTLINE:
                        F(i, j, k) = 0.0;
                        H(i, j, k) = 0.0;
                        break;
                    case FRONTRIGHTLINE:
                        F(i, j, k) = 0.0;
                        H(i, j, k) = 0.0;
                        break;
                    case FRONTTOPLINE:
                        G(i, j, k) = 0.0;
                        H(i, j, k) = 0.0;
                        break;
                    case FRONTBOTTOMLINE:
                        G(i, j, k) = 0.0;
                        H(i, j, k) = 0.0;
                        break;
                    case MIDTOPLEFTLINE:
                        F(i, j, k) = 0.0;
                        G(i, j, k) = 0.0;
                        break;
                    case MIDTOPRIGHTLINE:
                        F(i, j, k) = U(i, j, k);
                        G(i, j, k) = V(i, j, k);
                        break;
                    case MIDBOTTOMLEFTLINE:
                        F(i, j, k) = 0.0;
                        G(i, j, k) = 0.0;
                        break;
                    case MIDBOTTOMRIGHTLINE:
                        F(i, j, k) = 0.0;
                        G(i, j, k) = 0.0;
                        break;
                    case BACKLEFTLINE:
                        F(i, j, k) = 0.0;
                        H(i, j, k) = 0.0;
                        break;
                    case BACKRIGHTLINE:
                        F(i, j, k) = 0.0;
                        H(i, j, k) = 0.0;
                        break;
                    case BACKTOPLINE:
                        G(i, j, k) = 0.0;
                        H(i, j, k) = 0.0;
                        break;
                    case BACKBOTTOMLINE:
                        G(i, j, k) = 0.0;
                        H(i, j, k) = 0.0;
                        break;
                    case FRONTTOPLEFTCORNER:
                        F(i, j, k) = 0.0;
                        G(i, j, k) = 0.0;
                        H(i, j, k) = 0.0;
                        break;
                    case FRONTTOPRIGHTCORNER:
                        F(i, j, k) = 0.0;
                        G(i, j, k) = 0.0;
                        H(i, j, k) = 0.0;
                        break;
                    case FRONTBOTTOMLEFTCORNER:
                        F(i, j, k) = 0.0;
                        G(i, j, k) = 0.0;
                        H(i, j, k) = 0.0;
                        break;
                    case FRONTBOTTOMRIGHTCORNER:
                        F(i, j, k) = 0.0;
                        G(i, j, k) = 0.0;
                        H(i, j, k) = 0.0;
                        break;
                    case BACKTOPLEFTCORNER:
                        F(i, j, k) = 0.0;
                        G(i, j, k) = 0.0;
                        H(i, j, k) = 0.0;
                        break;
                    case BACKTOPRIGHTCORNER:
                        F(i, j, k) = 0.0;
                        G(i, j, k) = 0.0;
                        H(i, j, k) = 0.0;
                        break;
                    case BACKBOTTOMLEFTCORNER:
                        F(i, j, k) = 0.0;
                        G(i, j, k) = 0.0;
                        H(i, j, k) = 0.0;
                        break;
                    case BACKBOTTOMRIGHTCORNER:
                        F(i, j, k) = 0.0;
                        G(i, j, k) = 0.0;
                        H(i, j, k) = 0.0;
                        break;
                    }
                }
            }
        }
    }

    /* ----------------------------- boundary of F ---------------------------
     */
    if (commIsBoundary(&d->comm, LEFT)) {
        for (int k = 1; k < kmaxLocal + 1; k++) {
            for (int j = 1; j < jmaxLocal + 1; j++) {
                F(0, j, k) = U(0, j, k);
            }
        }
    }

    if (commIsBoundary(&d->comm, RIGHT)) {
        for (int k = 1; k < kmaxLocal + 1; k++) {
            for (int j = 1; j < jmaxLocal + 1; j++) {
                F(imaxLocal, j, k) = U(imaxLocal, j, k);
            }
        }
    }

    /* ----------------------------- boundary of G ---------------------------
     */
    if (commIsBoundary(&d->comm, BOTTOM)) {
        for (int k = 1; k < kmaxLocal + 1; k++) {
            for (int i = 1; i < imaxLocal + 1; i++) {
                G(i, 0, k) = V(i, 0, k);
            }
        }
    }

    if (commIsBoundary(&d->comm, TOP)) {
        for (int k = 1; k < kmaxLocal + 1; k++) {
            for (int i = 1; i < imaxLocal + 1; i++) {
                G(i, jmaxLocal, k) = V(i, jmaxLocal, k);
            }
        }
    }

    /* ----------------------------- boundary of H ---------------------------
     */
    if (commIsBoundary(&d->comm, FRONT)) {
        for (int j = 1; j < jmaxLocal + 1; j++) {
            for (int i = 1; i < imaxLocal + 1; i++) {
                H(i, j, 0) = W(i, j, 0);
            }
        }
    }

    if (commIsBoundary(&d->comm, BACK)) {
        for (int j = 1; j < jmaxLocal + 1; j++) {
            for (int i = 1; i < imaxLocal + 1; i++) {
                H(i, j, kmaxLocal) = W(i, j, kmaxLocal);
            }
        }
    }
}

void adaptUV(Discretization* d)
{
    int imaxLocal = d->comm.imaxLocal;
    int jmaxLocal = d->comm.jmaxLocal;
    int kmaxLocal = d->comm.kmaxLocal;

    double* p = d->p;
    double* u = d->u;
    double* v = d->v;
    double* w = d->w;
    double* f = d->f;
    double* g = d->g;
    double* h = d->h;

    double factorX = d->dt / d->grid.dx;
    double factorY = d->dt / d->grid.dy;
    double factorZ = d->dt / d->grid.dz;

    for (int k = 1; k < kmaxLocal + 1; k++) {
        for (int j = 1; j < jmaxLocal + 1; j++) {
            for (int i = 1; i < imaxLocal + 1; i++) {
                U(i, j, k) = F(i, j, k) - (P(i + 1, j, k) - P(i, j, k)) * factorX;
                V(i, j, k) = G(i, j, k) - (P(i, j + 1, k) - P(i, j, k)) * factorY;
                W(i, j, k) = H(i, j, k) - (P(i, j, k + 1) - P(i, j, k)) * factorZ;
            }
        }
    }
}