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

#include "mpi.h"
#include "vtkWriter.h"
#define G(v, i, j, k) v[(k)*imax * jmax + (j)*imax + (i)]

static float floatSwap(float f)
{
    union {
        float f;
        char b[4];
    } dat1, dat2;

    dat1.f    = f;
    dat2.b[0] = dat1.b[3];
    dat2.b[1] = dat1.b[2];
    dat2.b[2] = dat1.b[1];
    dat2.b[3] = dat1.b[0];
    return dat2.f;
}

static void writeHeader(VtkOptions* o)
{
    const size_t MAX_HEADER = 200;

    char* header = (char*)malloc(MAX_HEADER);
    char* cursor = header;

    cursor += sprintf(cursor, "# vtk DataFile Version 3.0\n");
    cursor += sprintf(cursor, "PAMPI cfd solver output\n");
    cursor += sprintf(cursor, "BINARY\n");

    cursor += sprintf(cursor, "DATASET STRUCTURED_POINTS\n");
    cursor += sprintf(cursor,
        "DIMENSIONS %d %d %d\n",
        o->grid.imax,
        o->grid.jmax,
        o->grid.kmax);
    cursor += sprintf(cursor,
        "ORIGIN %f %f %f\n",
        o->grid.dx * 0.5,
        o->grid.dy * 0.5,
        o->grid.dz * 0.5);
    cursor += sprintf(cursor, "SPACING %f %f %f\n", o->grid.dx, o->grid.dy, o->grid.dz);
    cursor += sprintf(cursor,
        "POINT_DATA %d\n",
        o->grid.imax * o->grid.jmax * o->grid.kmax);
}

void vtkOpen(VtkOptions* o, char* problem)
{
    char filename[50];

    snprintf(filename, 50, "%s-p%d.vtk", problem, o->comm.size);
    MPI_File_open(o->comm.comm,
        filename,
        MPI_MODE_WRONLY | MPI_MODE_CREATE,
        MPI_INFO_NULL,
        &o->fh);

    if (commIsMaster(&o->comm)) printf("Writing VTK output for %s\n", problem);
    writeHeader(o);
}

static void writeScalar(VtkOptions* o, double* s)
{
    int imax = o->grid.imax;
    int jmax = o->grid.jmax;
    int kmax = o->grid.kmax;

    for (int k = 0; k < kmax; k++) {
        for (int j = 0; j < jmax; j++) {
            for (int i = 0; i < imax; i++) {
                if (o->fmt == ASCII) {
                    fprintf(o->fh, "%f\n", G(s, i, j, k));
                } else if (o->fmt == BINARY) {
                    fwrite((float[1]) { floatSwap(G(s, i, j, k)) },
                        sizeof(float),
                        1,
                        o->fh);
                }
            }
        }
    }
    if (o->fmt == BINARY) fprintf(o->fh, "\n");
}

static bool isInitialized(FILE* ptr)
{
    if (ptr == NULL) {
        printf("vtkWriter not initialize! Call vtkOpen first!\n");
        return false;
    }
    return true;
}

void vtkScalar(VtkOptions* o, char* name, double* s)
{
    if (commIsMaster(&o->comm)) printf("Register scalar %s\n", name);

    if (o->mode == UNIX) {
        if (commIsMaster(&o->comm)) {
            if (!isInitialized(o->fh)) return;
            fprintf(o->fh, "SCALARS %s float 1\n", name);
            fprintf(o->fh, "LOOKUP_TABLE default\n");
            writeScalar(o, s);
        }
    } else if (o->mode == MPI) {
    }
}

static void writeVector(VtkOptions* o, VtkVector vec)
{
    int imax = o->grid.imax;
    int jmax = o->grid.jmax;
    int kmax = o->grid.kmax;

    for (int k = 0; k < kmax; k++) {
        for (int j = 0; j < jmax; j++) {
            for (int i = 0; i < imax; i++) {
                if (o->fmt == ASCII) {
                    fprintf(o->fh,
                        "%f %f %f\n",
                        G(vec.u, i, j, k),
                        G(vec.v, i, j, k),
                        G(vec.w, i, j, k));
                } else if (o->fmt == BINARY) {
                    fwrite((float[3]) { floatSwap(G(vec.u, i, j, k)),
                               floatSwap(G(vec.v, i, j, k)),
                               floatSwap(G(vec.w, i, j, k)) },
                        sizeof(float),
                        3,
                        o->fh);
                }
            }
        }
    }
    if (o->fmt == BINARY) fprintf(o->fh, "\n");
}

void vtkVector(VtkOptions* o, char* name, VtkVector vec)
{
    if (commIsMaster(&o->comm)) printf("Register vector %s\n", name);

    if (o->mode == UNIX) {
        if (commIsMaster(&o->comm)) {
            if (!isInitialized(o->fh)) return;
            fprintf(o->fh, "VECTORS %s float\n", name);
            writeVector(o, vec);
        }
    } else if (o->mode == MPI) {
    }
}

void vtkClose(VtkOptions* o) { MPI_File_close(o->fh); }