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

#include "allocate.h"
#include "comm.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)]

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

void initSolver(Solver* s, Parameter* params)
{
    s->problem  = params->name;
    s->bcLeft   = params->bcLeft;
    s->bcRight  = params->bcRight;
    s->bcBottom = params->bcBottom;
    s->bcTop    = params->bcTop;
    s->imax     = params->imax;
    s->jmax     = params->jmax;
    s->xlength  = params->xlength;
    s->ylength  = params->ylength;
    s->dx       = params->xlength / params->imax;
    s->dy       = params->ylength / params->jmax;
    s->eps      = params->eps;
    s->omega    = params->omg;
    s->itermax  = params->itermax;
    s->re       = params->re;
    s->gx       = params->gx;
    s->gy       = params->gy;
    s->dt       = params->dt;
    s->te       = params->te;
    s->tau      = params->tau;
    s->gamma    = params->gamma;

    commInit(&s->comm, s->jmax, s->imax);

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

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

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

    double dx = s->dx;
    double dy = s->dy;

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

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

    commShift(&s->comm, f, g);

    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* s)
{
    int imax      = s->imax;
    int jmax      = s->jmax;
    int imaxLocal = s->comm.imaxLocal;
    int jmaxLocal = s->comm.jmaxLocal;
    double eps    = s->eps;
    int itermax   = s->itermax;
    double dx2    = s->dx * s->dx;
    double dy2    = s->dy * s->dy;
    double idx2   = 1.0 / dx2;
    double idy2   = 1.0 / dy2;
    double factor = s->omega * 0.5 * (dx2 * dy2) / (dx2 + dy2);
    double* p     = s->p;
    double* rhs   = s->rhs;
    double epssq  = eps * eps;
    int it        = 0;
    double res    = 1.0;
    commExchange(&s->comm, p);

    while ((res >= epssq) && (it < itermax)) {
        res = 0.0;

        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 (commIsBoundary(&s->comm, BOTTOM)) { // set bottom bc
            for (int i = 1; i < imaxLocal + 1; i++) {
                P(i, 0) = P(i, 1);
            }
        }

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

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

        if (commIsBoundary(&s->comm, RIGHT)) { // 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
    if (res < eps) {
        return 0;
    } else {
        return 1;
    }
}

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

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

    commReduction(&maxval, MAX);
    return maxval;
}

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

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

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

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

    while ((res >= epssq) && (it < itermax)) {
        res = 0.0;

        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 (commIsBoundary(&s->comm, BOTTOM)) { // set bottom bc
            for (int i = 1; i < imaxLocal + 1; i++) {
                P(i, 0) = P(i, 1);
            }
        }

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

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

        if (commIsBoundary(&s->comm, RIGHT)) { // 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++;
    }
}
int solveRBA(Solver* s)
{
    int imax      = s->imax;
    int jmax      = s->jmax;
    int imaxLocal = s->comm.imaxLocal;
    int jmaxLocal = s->comm.jmaxLocal;
    double eps    = s->eps;
    int itermax   = s->itermax;
    double dx2    = s->dx * s->dx;
    double dy2    = s->dy * s->dy;
    double idx2   = 1.0 / dx2;
    double idy2   = 1.0 / dy2;
    double factor = s->omega * 0.5 * (dx2 * dy2) / (dx2 + dy2);
    double* p     = s->p;
    double* rhs   = s->rhs;
    double epssq  = eps * eps;
    int it        = 0;
    double res    = 1.0;
    commExchange(&s->comm, p);

    while ((res >= epssq) && (it < itermax)) {
        res = 0.0;

        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 (commIsBoundary(&s->comm, BOTTOM)) { // set bottom bc
            for (int i = 1; i < imaxLocal + 1; i++) {
                P(i, 0) = P(i, 1);
            }
        }

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

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

        if (commIsBoundary(&s->comm, RIGHT)) { // 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
    if (res < eps) {
        return 0;
    } else {
        return 1;
    }
}

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

    if (commIsBoundary(&s->comm, TOP)) {
        switch (s->bcTop) {
        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;
        }
    }

    if (commIsBoundary(&s->comm, BOTTOM)) {
        switch (s->bcBottom) {
        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;
        }
    }

    if (commIsBoundary(&s->comm, RIGHT)) {
        switch (s->bcRight) {
        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;
        }
    }

    if (commIsBoundary(&s->comm, LEFT)) {
        switch (s->bcLeft) {
        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* s)
{
    int imaxLocal = s->comm.imaxLocal;
    int jmaxLocal = s->comm.jmaxLocal;
    double* u     = s->u;

    if (strcmp(s->problem, "dcavity") == 0) {
        if (commIsBoundary(&s->comm, TOP)) {
            for (int i = 1; i < imaxLocal + 1; i++) {
                U(i, jmaxLocal + 1) = 2.0 - U(i, jmaxLocal);
            }
        }
    } else if (strcmp(s->problem, "canal") == 0) {
        if (commIsBoundary(&s->comm, LEFT)) {
            double ylength = s->ylength;
            double dy      = s->dy;
            int rest       = s->jmax % s->comm.size;
            int yc    = s->comm.rank * (s->jmax / s->comm.size) + MIN(rest, s->comm.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* s)
{
    double* u = s->u;
    double* v = s->v;
    double* f = s->f;
    double* g = s->g;

    int imaxLocal = s->comm.imaxLocal;
    int jmaxLocal = s->comm.jmaxLocal;

    double gx        = s->gx;
    double gy        = s->gy;
    double gamma     = s->gamma;
    double dt        = s->dt;
    double inverseRe = 1.0 / s->re;
    double inverseDx = 1.0 / s->dx;
    double inverseDy = 1.0 / s->dy;
    double du2dx, dv2dy, duvdx, duvdy;
    double du2dx2, du2dy2, dv2dx2, dv2dy2;

    commExchange(&s->comm, u);
    commExchange(&s->comm, 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 (commIsBoundary(&s->comm, LEFT)) {
        for (int j = 1; j < jmaxLocal + 1; j++) {
            F(0, j) = U(0, j);
        }
    }

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

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

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

void adaptUV(Solver* s)
{
    int imaxLocal = s->comm.imaxLocal;
    int jmaxLocal = s->comm.jmaxLocal;

    double* p = s->p;
    double* u = s->u;
    double* v = s->v;
    double* f = s->f;
    double* g = s->g;

    double factorX = s->dt / s->dx;
    double factorY = s->dt / s->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* s, double* u, double* v, double* p)
{
    int imax  = s->imax;
    int jmax  = s->jmax;
    double dx = s->dx;
    double dy = s->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);
}