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

#include "allocate.h"
#include "parameter.h"
#include "solver.h"
#include <alloca.h>
#include <mpi.h>
#include <stddef.h>
#include <stdlib.h>

#define PI        3.14159265358979323846
#define P(i, j)   p[(i) * (solver->ljmax + 2) + (j)]
#define RHS(i, j) rhs[(i) * (solver->ljmax + 2) + (j)]

int distribute_dim(int dim_rank, int dim_comm_size, int dim_size)
{
    int base_size = dim_size / dim_comm_size;
    return dim_rank < (dim_size % dim_comm_size) ? base_size + 1 : base_size;
}

int get_dim_start(int dim_rank, int dim_comm_size, int dim_size)
{
    int dim_start = 0;
    for (int i = 0; i < dim_rank; i++)
        dim_start += distribute_dim(i, dim_comm_size, dim_size);
    return dim_start;
}

void write_matrix_to_file(
    const char* _Nonnull filename, double* matrix, size_t rows, size_t cols)
{
    FILE* fp;
    fp = fopen(filename, "w");

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

    for (int i = 0; i < rows; i++) {
        for (int j = 0; j < cols; j++) {
            fprintf(fp, "%f ", matrix[i * (cols) + j]);
        }
        fprintf(fp, "\n");
    }

    fclose(fp);
}

void initSolver(Solver* solver, Parameter* params, int problem)
{
    int x_low, x_high, y_low, y_high;
    int rank, size;
    MPI_Datatype contiguos_boundary, vector_boundary;

    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    solver->imax    = params->imax;
    solver->jmax    = params->jmax;
    solver->dx      = params->xlength / params->imax;
    solver->dy      = params->ylength / params->jmax;
    solver->eps     = params->eps;
    solver->omega   = params->omg;
    solver->rho     = params->rho;
    solver->itermax = params->itermax;
    solver->dims[0] = 0;
    solver->dims[1] = 0;
    solver->rank    = 0;
    solver->size    = 0;
    solver->rank    = rank;
    solver->size    = size;

    MPI_Dims_create(size, 2, solver->dims);
    MPI_Cart_create(MPI_COMM_WORLD,
        2,
        solver->dims,
        (int[2]) { 0, 0 },
        0,
        &solver->cart_comm);
    MPI_Cart_coords(solver->cart_comm, rank, 2, solver->coords);
    solver->boundary = NB;
    MPI_Cart_shift(solver->cart_comm, 0, 1, &x_low, &x_high);
    MPI_Cart_shift(solver->cart_comm, 1, 1, &y_low, &y_high);

    if (x_low == MPI_PROC_NULL) solver->boundary += BOTTOM;
    if (y_low == MPI_PROC_NULL) solver->boundary += LEFT;
    if (x_high == MPI_PROC_NULL) solver->boundary += TOP;
    if (y_high == MPI_PROC_NULL) solver->boundary += RIGHT;

    solver->limax = distribute_dim(solver->coords[0], solver->dims[0], solver->imax);
    solver->ljmax = distribute_dim(solver->coords[1], solver->dims[1], solver->jmax);

    MPI_Type_contiguous(solver->ljmax, MPI_DOUBLE, &contiguos_boundary);
    MPI_Type_vector(solver->limax, 1, solver->ljmax + 2, MPI_DOUBLE, &vector_boundary);
    MPI_Type_commit(&contiguos_boundary);
    MPI_Type_commit(&vector_boundary);
    solver->types[0] = contiguos_boundary;
    solver->types[1] = contiguos_boundary;
    solver->types[2] = vector_boundary;
    solver->types[3] = vector_boundary;

    solver->snd_displacements[0] = ((solver->ljmax + 2) + 1) * sizeof(double);
    solver->snd_displacements[1] = ((solver->ljmax + 2) * (solver->limax) + 1) *
                                   sizeof(double);
    solver->snd_displacements[2] = ((solver->ljmax + 2) + 1) * sizeof(double);
    solver->snd_displacements[3] = ((solver->ljmax + 2) + solver->ljmax) * sizeof(double);

    solver->rcv_displacements[0] = (1) * sizeof(double);
    solver->rcv_displacements[1] = ((solver->ljmax + 2) * (solver->limax + 1) + 1) *
                                   sizeof(double);
    solver->rcv_displacements[2] = ((solver->ljmax + 2)) * sizeof(double);
    solver->rcv_displacements[3] = ((solver->ljmax + 2) + solver->ljmax + 1) *
                                   sizeof(double);

    size_t bytesize = (solver->limax + 2) * (solver->ljmax + 2) * sizeof(double);
    solver->p       = allocate(64, bytesize);
    solver->rhs     = allocate(64, bytesize);
    double dx       = solver->dx;
    double dy       = solver->dy;
    double* p       = solver->p;
    double* rhs     = solver->rhs;
    solver->i_start = get_dim_start(solver->coords[0], solver->dims[0], solver->imax);
    solver->j_start = get_dim_start(solver->coords[1], solver->dims[1], solver->jmax);

#ifdef DEBUG
    printf("Prc %d; {%d,%d}, {%d,%d}, {%d,%d}\n",
        rank,
        solver->coords[0],
        solver->coords[1],
        solver->limax,
        solver->ljmax,
        solver->i_start,
        solver->j_start);
    fflush(stdout);
#endif /* ifdef DEBUG */

    for (int i = 0; i < solver->limax + 2; i++) {
        for (int j = 0; j < solver->ljmax + 2; j++) {
            P(i, j) = sin(2.0 * PI * (solver->i_start + i) * dx * 2.0) +
                      sin(2.0 * PI * (solver->j_start + j) * dy * 2.0);
        }
    }

    if (problem == 2) {
        for (int i = 0; i < solver->limax + 2; i++) {
            for (int j = 0; j < solver->ljmax + 2; j++) {
                RHS(i, j) = sin(2.0 * PI * (solver->i_start + i) * dx);
            }
        }
    } else {
        for (int i = 0; i < solver->limax + 2; i++) {
            for (int j = 0; j < solver->ljmax + 2; j++) {
                RHS(i, j) = 0.0;
            }
        }
    }
}

void solveRB(Solver* solver)
{
    int imax      = solver->imax;
    int jmax      = solver->jmax;
    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;
    int pass, jsw, isw;

    while ((res >= epssq) && (it < itermax)) {
        res = 0.0;
        jsw = solver->j_start % 2 ? 1 : 2;

        MPI_Neighbor_alltoallw(solver->p,
            (int[4]) { 1, 1, 1, 1 },
            solver->snd_displacements,
            solver->types,
            solver->p,
            (int[4]) { 1, 1, 1, 1 },
            solver->rcv_displacements,
            solver->types,
            solver->cart_comm);

        for (pass = 0; pass < 2; pass++) { // why ??
            isw = jsw;
            for (int i = 1; i < solver->limax + 1; i++) {
                for (int j = isw; j < solver->ljmax + 1; j += 2) {
                    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);
                }
                isw = 3 - isw;
            }
            jsw = 3 - jsw;
        }

        MPI_Allreduce(MPI_IN_PLACE, &res, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);

        if (solver->boundary & BOTTOM)
            for (int j = 1; j < solver->ljmax + 1; j++)
                P(0, j) = P(1, j);
        if (solver->boundary & TOP)
            for (int j = 1; j < solver->ljmax + 1; j++)
                P(solver->limax + 1, j) = P(solver->limax, j);
        if (solver->boundary & LEFT)
            for (int i = 1; i < solver->limax + 1; i++)
                P(i, 0) = P(i, 1);
        if (solver->boundary & RIGHT)
            for (int i = 1; i < solver->limax + 1; i++)
                P(i, solver->ljmax + 1) = P(i, solver->ljmax);

        res = res / (double)(imax * jmax);

#ifdef DEBUG
        printf("%d Residuum: %e\n", it, res);
#endif
        it++;
    }

    if (solver->rank == 0)
        printf("Solver took %d iterations to reach %f\n", it, sqrt(res));
}

void writeResult(Solver* solver, char* filename)
{
    int imax            = solver->imax;
    int jmax            = solver->jmax;
    double* p           = solver->p;
    int* rcv_count      = NULL;
    int* displacements  = NULL;
    double* p_total     = NULL;
    double* send_buffer = NULL;
    int local_size      = solver->limax * solver->ljmax;

    if (solver->rank == 0) {

        rcv_count     = allocate(64, solver->size * sizeof(int));
        displacements = allocate(64, solver->size * sizeof(int));
        p_total       = allocate(64, sizeof(double) * solver->imax * solver->jmax);
        double* p_total_ordered = allocate(64,
            sizeof(double) * solver->imax * solver->jmax);

        MPI_Gather(&local_size, 1, MPI_INT, rcv_count, 1, MPI_INT, 0, MPI_COMM_WORLD);

        displacements[0] = 0;

        for (int i = 0; i < solver->limax; i++)
            for (int j = 0; j < solver->ljmax; j++)
                p_total_ordered[i * (solver->ljmax) + j] =
                    solver->p[(i + 1) * (solver->ljmax + 2) + (j + 1)];

        for (int i = 1; i < solver->size; i++)
            displacements[i] = displacements[i - 1] + rcv_count[i - 1];

        MPI_Gatherv(p_total_ordered,
            local_size,
            MPI_DOUBLE,
            p_total,
            rcv_count,
            displacements,
            MPI_DOUBLE,
            0,
            MPI_COMM_WORLD);

        for (int process = 0; process < solver->size; process++) {

            int coords[2] = { 0 };

            MPI_Cart_coords(solver->cart_comm, process, 2, coords);

            int i_max_local   = distribute_dim(coords[0], solver->dims[0], solver->imax);
            int j_max_local   = distribute_dim(coords[1], solver->dims[1], solver->jmax);
            int i_start_local = get_dim_start(coords[0], solver->dims[0], solver->imax);
            int j_start_local = get_dim_start(coords[1], solver->dims[1], solver->jmax);

            for (int i = 0; i < i_max_local; i++)
                for (int j = 0; j < j_max_local; j++)
                    p_total_ordered[(i + i_start_local) * (solver->jmax) +
                                    (j + j_start_local)] =
                        p_total[displacements[process] + i * (j_max_local) + j];
        }

        write_matrix_to_file(filename, p_total_ordered, solver->imax, solver->jmax);

    } else {

        send_buffer = allocate(64, local_size * sizeof(double));
        for (int i = 0; i < solver->limax; i++)
            for (int j = 0; j < solver->ljmax; j++)
                send_buffer[i * (solver->ljmax) + j] =
                    solver->p[(i + 1) * (solver->ljmax + 2) + (j + 1)];

        MPI_Gather(&local_size, 1, MPI_INT, rcv_count, 1, MPI_INT, 0, MPI_COMM_WORLD);

        MPI_Gatherv(send_buffer,
            local_size,
            MPI_DOUBLE,
            p_total,
            rcv_count,
            displacements,
            MPI_DOUBLE,
            0,
            MPI_COMM_WORLD);
    }
}
