/*
 * Copyright (C) 2022 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 <mpi.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

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

#define P(i, j)   p[(j) * (imaxLocal + 2) + (i)]
#define F(i, j)   f[(j) * (imaxLocal + 2) + (i)]
#define G(i, j)   g[(j) * (imaxLocal + 2) + (i)]
#define U(i, j)   u[(j) * (imaxLocal + 2) + (i)]
#define V(i, j)   v[(j) * (imaxLocal + 2) + (i)]
#define RHS(i, j) rhs[(j) * (imaxLocal + 2) + (i)]

#define NDIMS 2
#define IDIM  0
#define JDIM  1

static int sizeOfRank(int rank, int size, int N)
{
    return N / size + ((N % size > rank) ? 1 : 0);
}

void print(Solver* solver, double* grid)
{
    int imaxLocal = solver->imaxLocal;

    for (int i = 0; i < solver->size; i++) {
        if (i == solver->rank) {
            printf(
                "### RANK %d #######################################################\n",
                solver->rank);
            for (int j = 0; j < solver->jmaxLocal + 2; j++) {
                printf("%02d: ", j);
                for (int i = 0; i < solver->imaxLocal + 2; i++) {
                    printf("%12.8f  ", grid[j * (imaxLocal + 2) + i]);
                }
                printf("\n");
            }
            fflush(stdout);
        }
        MPI_Barrier(MPI_COMM_WORLD);
    }
}

static void exchange(Solver* solver, double* grid)
{
    double* buf[8];
    MPI_Request requests[8];
    for (int i = 0; i < 8; i++)
        requests[i] = MPI_REQUEST_NULL;

    buf[0] = grid + 1;                                                     // recv bottom
    buf[1] = grid + (solver->imaxLocal + 2) + 1;                           // send bottom
    buf[2] = grid + (solver->jmaxLocal + 1) * (solver->imaxLocal + 2) + 1; // recv top
    buf[3] = grid + (solver->jmaxLocal) * (solver->imaxLocal + 2) + 1;     // send top
    buf[4] = grid + (solver->imaxLocal + 2);                               // recv left
    buf[5] = grid + (solver->imaxLocal + 2) + 1;                           // send left
    buf[6] = grid + (solver->imaxLocal + 2) + (solver->imaxLocal + 1);     // recv right
    buf[7] = grid + (solver->imaxLocal + 2) + (solver->imaxLocal);         // send right

    for (int i = 0; i < 2; i++) {
        int tag = 0;
        if (solver->jNeighbours[i] != MPI_PROC_NULL) {
            tag = solver->jNeighbours[i];
        }
        /* exchange ghost cells with bottom/top neighbor */
        MPI_Irecv(buf[i * 2],
            1,
            solver->jBufferType,
            solver->jNeighbours[i],
            tag,
            solver->comm,
            &requests[i * 2]);
        MPI_Isend(buf[(i * 2) + 1],
            1,
            solver->jBufferType,
            solver->jNeighbours[i],
            solver->rank,
            solver->comm,
            &requests[i * 2 + 1]);

        tag = 0;
        if (solver->iNeighbours[i] != MPI_PROC_NULL) {
            tag = solver->iNeighbours[i];
        }
        /* exchange ghost cells with left/right neighbor */
        MPI_Irecv(buf[i * 2 + 4],
            1,
            solver->iBufferType,
            solver->iNeighbours[i],
            tag,
            solver->comm,
            &requests[i * 2 + 4]);
        MPI_Isend(buf[i * 2 + 5],
            1,
            solver->iBufferType,
            solver->iNeighbours[i],
            solver->rank,
            solver->comm,
            &requests[(i * 2) + 5]);
    }

    MPI_Waitall(8, requests, MPI_STATUSES_IGNORE);
}

static void shift(Solver* solver)
{
    MPI_Request requests[4] = { MPI_REQUEST_NULL,
        MPI_REQUEST_NULL,
        MPI_REQUEST_NULL,
        MPI_REQUEST_NULL };
    double* f               = solver->f;
    double* g               = solver->g;

    /* shift G */
    double* buf = g + 1;
    /* receive ghost cells from bottom neighbor */
    MPI_Irecv(buf,
        1,
        solver->jBufferType,
        solver->jNeighbours[0],
        0,
        solver->comm,
        &requests[0]);

    buf = g + (solver->jmaxLocal) * (solver->imaxLocal + 2) + 1;
    /* send ghost cells to top neighbor */
    MPI_Isend(buf,
        1,
        solver->jBufferType,
        solver->jNeighbours[1],
        0,
        solver->comm,
        &requests[1]);

    /* shift F */
    buf = f + (solver->imaxLocal + 2);
    /* receive ghost cells from left neighbor */
    MPI_Irecv(buf,
        1,
        solver->iBufferType,
        solver->iNeighbours[0],
        1,
        solver->comm,
        &requests[2]);

    buf = f + (solver->imaxLocal + 2) + (solver->imaxLocal);
    /* send ghost cells to right neighbor */
    MPI_Isend(buf,
        1,
        solver->iBufferType,
        solver->iNeighbours[1],
        1,
        solver->comm,
        &requests[3]);

    MPI_Waitall(4, requests, MPI_STATUSES_IGNORE);
}

void debugExchange(Solver* solver)
{
    int imaxLocal = solver->imaxLocal;
    int jmaxLocal = solver->jmaxLocal;

    for (int j = 0; j < jmaxLocal + 2; j++) {
        for (int i = 0; i < solver->imaxLocal + 2; i++) {
            solver->p[j * (imaxLocal + 2) + i] = solver->rank + 0.01 * i + 0.0001 * j;
        }
    }
    collectResult(solver);
    /* print(solver, solver->p); */
}

void debugBC(Solver* solver)
{
    int imaxLocal = solver->imaxLocal;
    int jmaxLocal = solver->jmaxLocal;
    double* v     = solver->v;

    // Northern boundary
    if (solver->coords[JDIM] == (solver->dims[JDIM] - 1)) { // set top bc
        for (int i = 1; i < imaxLocal + 1; i++) {
            V(i, jmaxLocal + 1) = 10.0 + solver->rank;
        }
    }

    // Eastern boundary
    if (solver->coords[IDIM] == (solver->dims[IDIM] - 1)) { // set right bc
        for (int j = 1; j < jmaxLocal + 1; j++) {
            V(imaxLocal + 1, j) = 20.0 + solver->rank;
        }
    }

    // Southern boundary
    if (solver->coords[JDIM] == 0) { // set bottom bc
        for (int i = 1; i < imaxLocal + 1; i++) {
            V(i, 0) = 30.0 + solver->rank;
        }
    }

    // Western boundary
    if (solver->coords[IDIM] == 0) { // set left bc
        for (int j = 1; j < jmaxLocal + 1; j++) {
            V(0, j) = 40.0 + solver->rank;
        }
    }
    print(solver, solver->v);
}

static void assembleResult(Solver* solver,
    double* src,
    double* dst,
    int imaxLocal[],
    int jmaxLocal[],
    int offset[])
{
    MPI_Request* requests;
    int numRequests = 1;

    if (solver->rank == 0) {
        numRequests = solver->size + 1;
    } else {
        numRequests = 1;
    }

    requests = (MPI_Request*)malloc(numRequests * sizeof(MPI_Request));

    /* all ranks send their bulk array */
    MPI_Datatype bulkType;
    const int ndims     = 2;
    int oldSizes[ndims] = { solver->jmaxLocal + 2, solver->imaxLocal + 2 };
    int newSizes[ndims] = { solver->jmaxLocal, solver->imaxLocal };
    int starts[ndims]   = { 1, 1 };
    MPI_Type_create_subarray(2,
        oldSizes,
        newSizes,
        starts,
        MPI_ORDER_C,
        MPI_DOUBLE,
        &bulkType);
    MPI_Type_commit(&bulkType);

    MPI_Isend(src, 1, bulkType, 0, 0, solver->comm, &requests[0]);

    /* rank 0 assembles the subdomains */
    if (solver->rank == 0) {
        for (int i = 0; i < solver->size; i++) {
            MPI_Datatype domainType;
            MPI_Type_vector(jmaxLocal[i],
                imaxLocal[i],
                solver->imax,
                MPI_DOUBLE,
                &domainType);
            MPI_Type_commit(&domainType);

            MPI_Irecv(dst + offset[i],
                1,
                domainType,
                i,
                0,
                solver->comm,
                &requests[i + 1]);
        }
    }

    MPI_Waitall(numRequests, requests, MPI_STATUSES_IGNORE);
}

static int sum(int* sizes, int position)
{
    int sum = 0;

    for (int i = 0; i < position; i++) {
        sum += sizes[i];
    }

    return sum;
}

void collectResult(Solver* solver)
{
    double* Pall = NULL;
    double* Uall = NULL;
    double* Vall = NULL;
    int offset[solver->size];
    int imaxLocal[solver->size];
    int jmaxLocal[solver->size];

    MPI_Gather(&solver->imaxLocal, 1, MPI_INT, imaxLocal, 1, MPI_INT, 0, MPI_COMM_WORLD);
    MPI_Gather(&solver->jmaxLocal, 1, MPI_INT, jmaxLocal, 1, MPI_INT, 0, MPI_COMM_WORLD);

    if (solver->rank == 0) {
        Pall = allocate(64, (solver->imax) * (solver->jmax) * sizeof(double));
        Uall = allocate(64, (solver->imax) * (solver->jmax) * sizeof(double));
        Vall = allocate(64, (solver->imax) * (solver->jmax) * sizeof(double));

        for (int i = 0; i < solver->size; i++) {
            int coords[2];
            MPI_Cart_coords(solver->comm, i, 2, coords);
            int ioffset = sum(imaxLocal, coords[0]);
            int joffset = sum(jmaxLocal, coords[1]);
            offset[i]   = (joffset * solver->imax) + ioffset;
            printf("Rank: %d, Coords(i,j): %d %d, Size(i,j): %d %d, Offset(i,j): %d %d\n",
                i,
                coords[0],
                coords[1],
                imaxLocal[i],
                jmaxLocal[i],
                ioffset,
                joffset);
        }
    }

    /* collect P */
    assembleResult(solver, solver->p, Pall, imaxLocal, jmaxLocal, offset);

    /* collect U */
    assembleResult(solver, solver->u, Uall, imaxLocal, jmaxLocal, offset);

    /* collect V */
    assembleResult(solver, solver->v, Vall, imaxLocal, jmaxLocal, offset);

    /* write to disk */
    if (solver->rank == 0) writeResult(solver, Pall, Uall, Vall);
}

static void printConfig(Solver* solver)
{
    if (solver->rank == 0) {
        printf("Parameters for #%s#\n", solver->problem);
        printf("Boundary conditions N:%d E:%d S:%d W:%d\n",
            solver->bcN,
            solver->bcE,
            solver->bcS,
            solver->bcW);
        printf("\tReynolds number: %.2f\n", solver->re);
        printf("\tGx Gy: %.2f %.2f\n", solver->gx, solver->gy);
        printf("Geometry data:\n");
        printf("\tDomain box size (x, y): %.2f, %.2f\n",
            solver->xlength,
            solver->ylength);
        printf("\tCells (x, y): %d, %d\n", solver->imax, solver->jmax);
        printf("Timestep parameters:\n");
        printf("\tDefault stepsize: %.2f, Final time %.2f\n", solver->dt, solver->te);
        printf("\tdt bound: %.6f\n", solver->dtBound);
        printf("\tTau factor: %.2f\n", solver->tau);
        printf("Iterative solver parameters:\n");
        printf("\tMax iterations: %d\n", solver->itermax);
        printf("\tepsilon (stopping tolerance) : %f\n", solver->eps);
        printf("\tgamma factor: %f\n", solver->gamma);
        printf("\tomega (SOR relaxation): %f\n", solver->omega);
        printf("Communication parameters:\n");
    }
    for (int i = 0; i < solver->size; i++) {
        if (i == solver->rank) {
            printf("\tRank %d of %d\n", solver->rank, solver->size);
            printf("\tNeighbours (b, t, l, r): %d, %d, %d, %d\n",
                solver->jNeighbours[0],
                solver->jNeighbours[1],
                solver->iNeighbours[0],
                solver->iNeighbours[1]);
            printf("\tCoordinates %d,%d\n", solver->coords[0], solver->coords[1]);
            printf("\tLocal domain size: %dx%d\n", solver->imaxLocal, solver->jmaxLocal);
            fflush(stdout);
        }
    }
}

void initSolver(Solver* solver, Parameter* params)
{
    solver->problem = params->name;
    solver->bcN     = params->bcN;
    solver->bcS     = params->bcS;
    solver->bcW     = params->bcW;
    solver->bcE     = params->bcE;
    solver->imax    = params->imax;
    solver->jmax    = params->jmax;
    solver->xlength = params->xlength;
    solver->ylength = params->ylength;
    solver->dx      = params->xlength / params->imax;
    solver->dy      = params->ylength / params->jmax;
    solver->eps     = params->eps;
    solver->omega   = params->omg;
    solver->itermax = params->itermax;
    solver->re      = params->re;
    solver->gx      = params->gx;
    solver->gy      = params->gy;
    solver->dt      = params->dt;
    solver->te      = params->te;
    solver->tau     = params->tau;
    solver->gamma   = params->gamma;

    /* setup communication */
    MPI_Comm_rank(MPI_COMM_WORLD, &(solver->rank));
    MPI_Comm_size(MPI_COMM_WORLD, &(solver->size));
    int dims[NDIMS]    = { 0, 0 };
    int periods[NDIMS] = { 0, 0 };
    MPI_Dims_create(solver->size, NDIMS, dims);
    MPI_Cart_create(MPI_COMM_WORLD, NDIMS, dims, periods, 0, &solver->comm);
    MPI_Cart_shift(solver->comm,
        IDIM,
        1,
        &solver->iNeighbours[0],
        &solver->iNeighbours[1]);
    MPI_Cart_shift(solver->comm,
        JDIM,
        1,
        &solver->jNeighbours[0],
        &solver->jNeighbours[1]);
    MPI_Cart_get(solver->comm, NDIMS, solver->dims, periods, solver->coords);

    solver->imaxLocal = sizeOfRank(solver->rank, dims[IDIM], solver->imax);
    solver->jmaxLocal = sizeOfRank(solver->rank, dims[JDIM], solver->jmax);

    MPI_Type_contiguous(solver->imaxLocal, MPI_DOUBLE, &solver->jBufferType);
    MPI_Type_commit(&solver->jBufferType);

    MPI_Type_vector(solver->jmaxLocal,
        1,
        solver->imaxLocal + 2,
        MPI_DOUBLE,
        &solver->iBufferType);
    MPI_Type_commit(&solver->iBufferType);

    /* allocate arrays */
    int imaxLocal   = solver->imaxLocal;
    int jmaxLocal   = solver->jmaxLocal;
    size_t bytesize = (imaxLocal + 2) * (jmaxLocal + 2) * sizeof(double);
    solver->u       = allocate(64, bytesize);
    solver->v       = allocate(64, bytesize);
    solver->p       = allocate(64, bytesize);
    solver->rhs     = allocate(64, bytesize);
    solver->f       = allocate(64, bytesize);
    solver->g       = allocate(64, bytesize);

    for (int i = 0; i < (imaxLocal + 2) * (jmaxLocal + 2); i++) {
        solver->u[i]   = params->u_init;
        solver->v[i]   = params->v_init;
        solver->p[i]   = params->p_init;
        solver->rhs[i] = 0.0;
        solver->f[i]   = 0.0;
        solver->g[i]   = 0.0;
    }

    double dx          = solver->dx;
    double dy          = solver->dy;
    double inv_sqr_sum = 1.0 / (dx * dx) + 1.0 / (dy * dy);
    solver->dtBound    = 0.5 * solver->re * 1.0 / inv_sqr_sum;
#ifdef VERBOSE
    printConfig(solver);
#endif
}

void computeRHS(Solver* solver)
{
    int imaxLocal = solver->imaxLocal;
    int jmaxLocal = solver->jmaxLocal;
    double idx    = 1.0 / solver->dx;
    double idy    = 1.0 / solver->dy;
    double idt    = 1.0 / solver->dt;
    double* rhs   = solver->rhs;
    double* f     = solver->f;
    double* g     = solver->g;

    shift(solver);

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

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

    while ((res >= epssq) && (it < itermax)) {
        res = 0.0;
        exchange(solver, p);

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

                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);
            }
        }

        if (solver->coords[JDIM] == 0) { // set bottom bc
            for (int i = 1; i < imaxLocal + 1; i++) {
                P(i, 0) = P(i, 1);
            }
        }

        if (solver->coords[JDIM] == (solver->dims[JDIM] - 1)) { // set top bc
            for (int i = 1; i < imaxLocal + 1; i++) {
                P(i, jmaxLocal + 1) = P(i, jmaxLocal);
            }
        }

        if (solver->coords[IDIM] == 0) { // set left bc
            for (int j = 1; j < jmaxLocal + 1; j++) {
                P(0, j) = P(1, j);
            }
        }

        if (solver->coords[IDIM] == (solver->dims[IDIM] - 1)) { // set right bc
            for (int j = 1; j < jmaxLocal + 1; j++) {
                P(imaxLocal + 1, j) = P(imaxLocal, j);
            }
        }

        MPI_Allreduce(MPI_IN_PLACE, &res, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
        res = res / (double)(imax * jmax);
#ifdef DEBUG
        if (solver->rank == 0) {
            printf("%d Residuum: %e\n", it, res);
        }
#endif
        it++;
    }

#ifdef VERBOSE
    if (solver->rank == 0) {
        printf("Solver took %d iterations to reach %f\n", it, sqrt(res));
    }
#endif
    if (res < eps) {
        return 0;
    } else {
        return 1;
    }
}

static double maxElement(Solver* solver, double* m)
{
    int size      = (solver->imaxLocal + 2) * (solver->jmaxLocal + 2);
    double maxval = DBL_MIN;

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

    MPI_Allreduce(MPI_IN_PLACE, &maxval, 1, MPI_DOUBLE, MPI_MAX, MPI_COMM_WORLD);
    return maxval;
}

void computeTimestep(Solver* solver)
{
    double dt   = solver->dtBound;
    double dx   = solver->dx;
    double dy   = solver->dy;
    double umax = maxElement(solver, solver->u);
    double vmax = maxElement(solver, solver->v);

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

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

void setBoundaryConditions(Solver* solver)
{
    int imaxLocal = solver->imaxLocal;
    int jmaxLocal = solver->jmaxLocal;
    double* u     = solver->u;
    double* v     = solver->v;

    // Northern boundary
    if (solver->coords[JDIM] == (solver->dims[JDIM] - 1)) { // set top bc
        switch (solver->bcN) {
        case NOSLIP:
            for (int i = 1; i < imaxLocal + 1; i++) {
                V(i, jmaxLocal)     = 0.0;
                U(i, jmaxLocal + 1) = -U(i, jmaxLocal);
            }
            break;
        case SLIP:
            for (int i = 1; i < imaxLocal + 1; i++) {
                V(i, jmaxLocal)     = 0.0;
                U(i, jmaxLocal + 1) = U(i, jmaxLocal);
            }
            break;
        case OUTFLOW:
            for (int i = 1; i < imaxLocal + 1; i++) {
                U(i, jmaxLocal + 1) = U(i, jmaxLocal);
                V(i, jmaxLocal)     = V(i, jmaxLocal - 1);
            }
            break;
        case PERIODIC:
            break;
        }
    }

    // Southern boundary
    if (solver->coords[JDIM] == 0) { // set bottom bc
        switch (solver->bcS) {
        case NOSLIP:
            for (int i = 1; i < imaxLocal + 1; i++) {
                V(i, 0) = 0.0;
                U(i, 0) = -U(i, 1);
            }
            break;
        case SLIP:
            for (int i = 1; i < imaxLocal + 1; i++) {
                V(i, 0) = 0.0;
                U(i, 0) = U(i, 1);
            }
            break;
        case OUTFLOW:
            for (int i = 1; i < imaxLocal + 1; i++) {
                U(i, 0) = U(i, 1);
                V(i, 0) = V(i, 1);
            }
            break;
        case PERIODIC:
            break;
        }
    }

    // Eastern boundary
    if (solver->coords[IDIM] == (solver->dims[IDIM] - 1)) { // set right bc
        switch (solver->bcE) {
        case NOSLIP:
            for (int j = 1; j < jmaxLocal + 1; j++) {
                U(imaxLocal, j)     = 0.0;
                V(imaxLocal + 1, j) = -V(imaxLocal, j);
            }
            break;
        case SLIP:
            for (int j = 1; j < jmaxLocal + 1; j++) {
                U(imaxLocal, j)     = 0.0;
                V(imaxLocal + 1, j) = V(imaxLocal, j);
            }
            break;
        case OUTFLOW:
            for (int j = 1; j < jmaxLocal + 1; j++) {
                U(imaxLocal, j)     = U(imaxLocal - 1, j);
                V(imaxLocal + 1, j) = V(imaxLocal, j);
            }
            break;
        case PERIODIC:
            break;
        }
    }

    // Western boundary
    if (solver->coords[IDIM] == 0) { // set left bc
        switch (solver->bcW) {
        case NOSLIP:
            for (int j = 1; j < jmaxLocal + 1; j++) {
                U(0, j) = 0.0;
                V(0, j) = -V(1, j);
            }
            break;
        case SLIP:
            for (int j = 1; j < jmaxLocal + 1; j++) {
                U(0, j) = 0.0;
                V(0, j) = V(1, j);
            }
            break;
        case OUTFLOW:
            for (int j = 1; j < jmaxLocal + 1; j++) {
                U(0, j) = U(1, j);
                V(0, j) = V(1, j);
            }
            break;
        case PERIODIC:
            break;
        }
    }
}

void setSpecialBoundaryCondition(Solver* solver)
{
    int imaxLocal = solver->imaxLocal;
    int jmaxLocal = solver->jmaxLocal;
    double* u     = solver->u;

    if (strcmp(solver->problem, "dcavity") == 0) {
        if (solver->coords[JDIM] == (solver->dims[JDIM] - 1)) { // set top bc
            for (int i = 1; i < imaxLocal + 1; i++) {
                U(i, jmaxLocal + 1) = 2.0 - U(i, jmaxLocal);
            }
        }
    } else if (strcmp(solver->problem, "canal") == 0) {
        if (solver->coords[IDIM] == 0) { // set left bc
            double ylength = solver->ylength;
            double dy      = solver->dy;
            int rest       = solver->jmax % solver->size;
            int yc         = solver->rank * (solver->jmax / solver->size) +
                     MIN(rest, solver->rank);
            double ys = dy * (yc + 0.5);
            double y;

            /* printf("RANK %d yc: %d ys: %f\n", solver->rank, yc, ys); */

            for (int j = 1; j < jmaxLocal + 1; j++) {
                y       = ys + dy * (j - 0.5);
                U(0, j) = y * (ylength - y) * 4.0 / (ylength * ylength);
            }
        }
    }
    /* print(solver, solver->u); */
}

void computeFG(Solver* solver)
{
    double* u        = solver->u;
    double* v        = solver->v;
    double* f        = solver->f;
    double* g        = solver->g;
    int imaxLocal    = solver->imaxLocal;
    int jmaxLocal    = solver->jmaxLocal;
    double gx        = solver->gx;
    double gy        = solver->gy;
    double gamma     = solver->gamma;
    double dt        = solver->dt;
    double inverseRe = 1.0 / solver->re;
    double inverseDx = 1.0 / solver->dx;
    double inverseDy = 1.0 / solver->dy;
    double du2dx, dv2dy, duvdx, duvdy;
    double du2dx2, du2dy2, dv2dx2, dv2dy2;

    exchange(solver, u);
    exchange(solver, v);

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

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

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

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

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

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

    /* ----------------------------- boundary of F --------------------------- */
    if (solver->coords[IDIM] == 0) { // set left bc
        for (int j = 1; j < jmaxLocal + 1; j++) {
            F(0, j) = U(0, j);
        }
    }

    if (solver->coords[IDIM] == (solver->dims[IDIM] - 1)) { // set right bc
        for (int j = 1; j < jmaxLocal + 1; j++) {
            F(imaxLocal, j) = U(imaxLocal, j);
        }
    }

    /* ----------------------------- boundary of G --------------------------- */
    if (solver->coords[JDIM] == 0) { // set bottom bc
        for (int i = 1; i < imaxLocal + 1; i++) {
            G(i, 0) = V(i, 0);
        }
    }

    if (solver->coords[JDIM] == (solver->dims[JDIM] - 1)) { // set top bc
        for (int i = 1; i < imaxLocal + 1; i++) {
            G(i, jmaxLocal) = V(i, jmaxLocal);
        }
    }
}

void adaptUV(Solver* solver)
{
    int imaxLocal  = solver->imaxLocal;
    int jmaxLocal  = solver->jmaxLocal;
    double* p      = solver->p;
    double* u      = solver->u;
    double* v      = solver->v;
    double* f      = solver->f;
    double* g      = solver->g;
    double factorX = solver->dt / solver->dx;
    double factorY = solver->dt / solver->dy;

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

void writeResult(Solver* solver, double* p, double* u, double* v)
{
    int imax  = solver->imax;
    int jmax  = solver->jmax;
    double dx = solver->dx;
    double dy = solver->dy;
    double x = 0.0, y = 0.0;

    FILE* fp;
    fp = fopen("pressure.dat", "w");

    if (fp == NULL) {
        printf("Error!\n");
        exit(EXIT_FAILURE);
    }

    for (int j = 1; j < jmax; j++) {
        y = (double)(j - 0.5) * dy;
        for (int i = 1; i < imax; i++) {
            x = (double)(i - 0.5) * dx;
            fprintf(fp, "%.2f %.2f %f\n", x, y, p[j * (imax) + i]);
        }
        fprintf(fp, "\n");
    }

    fclose(fp);

    fp = fopen("velocity.dat", "w");

    if (fp == NULL) {
        printf("Error!\n");
        exit(EXIT_FAILURE);
    }

    for (int j = 1; j < jmax; j++) {
        y = dy * (j - 0.5);
        for (int i = 1; i < imax; i++) {
            x            = dx * (i - 0.5);
            double vel_u = (u[j * (imax) + i] + u[j * (imax) + (i - 1)]) / 2.0;
            double vel_v = (v[j * (imax) + i] + v[(j - 1) * (imax) + i]) / 2.0;
            double len   = sqrt((vel_u * vel_u) + (vel_v * vel_v));
            fprintf(fp, "%.2f %.2f %f %f %f\n", x, y, vel_u, vel_v, len);
        }
    }

    fclose(fp);
}