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

#include "allocate.h"
#include "parameter.h"
#include "progress.h"
#include "solver.h"
#include "test.h"
#include "timing.h"
#include "vtkWriter.h"

enum VARIANT { SOR = 1, RB, RBA };

int main(int argc, char** argv)
{
    int rank;
    int variant = RB;

    double timeStart, timeStop;
    Parameter params;
    Solver solver;

    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    initParameter(&params);

    if (argc < 2) {
        printf("Usage: %s <configFile>\n", argv[0]);
        exit(EXIT_SUCCESS);
    }

    readParameter(&params, argv[1]);
    if (argc == 3) 
    {
        variant = atoi(argv[2]);
    }
    if (commIsMaster(&solver.comm)) {
        printParameter(&params);
    }
    initSolver(&solver, &params);
#ifndef VERBOSE
    initProgress(solver.te);
#endif

    double tau = solver.tau;
    double te  = solver.te;
    double t   = 0.0;
    int nt     = 0;


    void (*solver_generic[])(solver) = {solve, solveRB, solveRBA};
    
    timeStart = getTimeStamp();

    while (t <= te) {
        if (tau > 0.0) computeTimestep(&solver);
        setBoundaryConditions(&solver);
        setSpecialBoundaryCondition(&solver);
        computeFG(&solver);
        computeRHS(&solver);
        // if (nt % 100 == 0) normalizePressure(&solver);
        (*solver_generic[variant - 1])(&solver);
        adaptUV(&solver);
        t += solver.dt;
        nt++;

#ifdef VERBOSE
        if (rank == 0) {
            printf("TIME %f , TIMESTEP %f\n", t, solver.dt);
        }
#else
        printProgress(t);
#endif
    }



    timeStop = getTimeStamp();
    stopProgress();
    if (commIsMaster(&solver.comm)) {
        printf("Solution took %.2fs\n", timeStop - timeStart);
    }

    // testInit(&solver);
    // commExchange(&solver.comm, solver.p);
    // testPrintHalo(&solver, solver.p);

    double *pg, *ug, *vg, *wg;

    if (commIsMaster(&solver.comm)) {
        size_t bytesize = solver.grid.imax * solver.grid.jmax * solver.grid.kmax *
                          sizeof(double);

        pg = allocate(64, bytesize);
        ug = allocate(64, bytesize);
        vg = allocate(64, bytesize);
        wg = allocate(64, bytesize);
    }

    commCollectResult(&solver.comm,
        ug,
        vg,
        wg,
        pg,
        solver.u,
        solver.v,
        solver.w,
        solver.p,
        solver.grid.kmax,
        solver.grid.jmax,
        solver.grid.imax);

    VtkOptions opts = { .grid = solver.grid, .comm = solver.comm };
    vtkOpen(&opts, solver.problem);
    vtkScalar(&opts, "pressure", pg);
    vtkVector(&opts, "velocity", (VtkVector) { ug, vg, wg });
    vtkClose(&opts);

    commFree(&solver.comm);
    MPI_Finalize();
    return EXIT_SUCCESS;
}
