/*
 * 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 "vtkWriter.h"

#define U(i, j, k) u[(k) * (imax + 2) * (jmax + 2) + (j) * (imax + 2) + (i)]
#define V(i, j, k) v[(k) * (imax + 2) * (jmax + 2) + (j) * (imax + 2) + (i)]
#define W(i, j, k) w[(k) * (imax + 2) * (jmax + 2) + (j) * (imax + 2) + (i)]
#define S(i, j, k) seg[(k) * (imax + 2) * (jmax + 2) + (j) * (imax + 2) + (i)]

static int ts     = 0;
unsigned int seed = 32767;
void printParticles(ParticleTracer* particletracer)
{
    for (int i = 0; i < particletracer->totalParticles; ++i) {
        printf("Particle position X : %.2f, Y : %.2f, flag : %d\n",
            particletracer->particlePool[i].x,
            particletracer->particlePool[i].y,
            particletracer->particlePool[i].flag);
    }
}
void injectParticles(ParticleTracer* particletracer, int* seg)
{

    int imax = particletracer->imax;
    int jmax = particletracer->jmax;
    int kmax = particletracer->kmax;
    for (int i = 0; i < particletracer->numberOfParticles; ++i) {

        // particletracer->particlePool[particletracer->pointer].x =
        // particletracer->linSpaceLine[i].x;
        // particletracer->particlePool[particletracer->pointer].y =
        // particletracer->linSpaceLine[i].y;
        // particletracer->particlePool[particletracer->pointer].z =
        // particletracer->linSpaceLine[i].z;

        particletracer->particlePool[particletracer->pointer].x = particletracer->x1;
        particletracer->particlePool[particletracer->pointer].y = (double)rand() /
                                                                  RAND_MAX *
                                                                  particletracer->ylength;
        particletracer->particlePool[particletracer->pointer].z = (double)rand() /
                                                                  RAND_MAX *
                                                                  particletracer->zlength;

        int i = particletracer->particlePool[particletracer->pointer].x /
                particletracer->dx;
        int j = particletracer->particlePool[particletracer->pointer].y /
                particletracer->dy;
        int k = particletracer->particlePool[particletracer->pointer].z /
                particletracer->dz;

        if (S(i, j, k) == NONE) {
            particletracer->particlePool[particletracer->pointer].flag = true;
            ++(particletracer->pointer);
            ++(particletracer->totalParticles);
        }
    }
}

void advanceParticles(ParticleTracer* particletracer,
    double* restrict u,
    double* restrict v,
    double* restrict w,
    int* restrict seg,
    double time)
{
    int imax = particletracer->imax;
    int jmax = particletracer->jmax;
    int kmax = particletracer->kmax;

    double dx = particletracer->dx;
    double dy = particletracer->dy;
    double dz = particletracer->dz;

    double xlength = particletracer->xlength;
    double ylength = particletracer->ylength;
    double zlength = particletracer->zlength;

    for (int i = 0; i < particletracer->totalParticles; ++i) {
        if (particletracer->particlePool[i].flag == true) {
            double x = particletracer->particlePool[i].x;
            double y = particletracer->particlePool[i].y;
            double z = particletracer->particlePool[i].z;

            int iCoord = (int)(x / dx) + 1;
            int jCoord = (int)((y + 0.5 * dy) / dy) + 1;
            int kCoord = (int)((z + 0.5 * dz) / dz) + 1;

            double x1 = (double)(iCoord - 1) * dx;
            double y1 = ((double)(jCoord - 1) - 0.5) * dy;
            double z1 = ((double)(kCoord - 1) - 0.5) * dz;

            double x2 = (double)iCoord * dx;
            double y2 = ((double)jCoord - 0.5) * dy;
            double z2 = ((double)kCoord - 0.5) * dz;

            double u_n =
                (1.0 / (dx * dy * dz)) *
                ((x2 - x) * (y2 - y) * (z2 - z) * U(iCoord - 1, jCoord - 1, kCoord - 1) +
                    (x - x1) * (y2 - y) * (z2 - z) * U(iCoord, jCoord - 1, kCoord - 1) +
                    (x2 - x) * (y - y1) * (z2 - z) * U(iCoord - 1, jCoord, kCoord - 1) +
                    (x - x1) * (y - y1) * (z2 - z) * U(iCoord, jCoord, kCoord - 1) +
                    (x2 - x) * (y2 - y) * (z - z1) * U(iCoord - 1, jCoord - 1, kCoord) +
                    (x - x1) * (y2 - y) * (z - z1) * U(iCoord, jCoord - 1, kCoord) +
                    (x2 - x) * (y - y1) * (z - z1) * U(iCoord - 1, jCoord, kCoord) +
                    (x - x1) * (y - y1) * (z - z1) * U(iCoord, jCoord, kCoord));

            double new_x                      = x + particletracer->dt * u_n;
            particletracer->particlePool[i].x = new_x;

            iCoord = (int)((x + 0.5 * dx) / dx) + 1;
            jCoord = (int)(y / dy) + 1;
            kCoord = (int)((z + 0.5 * dz) / dz) + 1;

            x1 = ((double)(iCoord - 1) - 0.5) * dx;
            y1 = (double)(jCoord - 1) * dy;
            z1 = ((double)(kCoord - 1) - 0.5) * dz;

            x2 = ((double)iCoord - 0.5) * dx;
            y2 = (double)jCoord * dy;
            z2 = ((double)kCoord - 0.5) * dz;

            double v_n =
                (1.0 / (dx * dy * dz)) *
                ((x2 - x) * (y2 - y) * (z2 - z) * V(iCoord - 1, jCoord - 1, kCoord - 1) +
                    (x - x1) * (y2 - y) * (z2 - z) * V(iCoord, jCoord - 1, kCoord - 1) +
                    (x2 - x) * (y - y1) * (z2 - z) * V(iCoord - 1, jCoord, kCoord - 1) +
                    (x - x1) * (y - y1) * (z2 - z) * V(iCoord, jCoord, kCoord - 1) +
                    (x2 - x) * (y2 - y) * (z - z1) * V(iCoord - 1, jCoord - 1, kCoord) +
                    (x - x1) * (y2 - y) * (z - z1) * V(iCoord, jCoord - 1, kCoord) +
                    (x2 - x) * (y - y1) * (z - z1) * V(iCoord - 1, jCoord, kCoord) +
                    (x - x1) * (y - y1) * (z - z1) * V(iCoord, jCoord, kCoord));

            double new_y                      = y + particletracer->dt * v_n;
            particletracer->particlePool[i].y = new_y;

            iCoord = (int)((x + 0.5 * dx) / dx) + 1;
            jCoord = (int)((y + 0.5 * dy) / dy) + 1;
            kCoord = (int)(z / dz) + 1;

            x1 = ((double)(iCoord - 1) - 0.5) * dx;
            y1 = ((double)(jCoord - 1) - 0.5) * dy;
            z1 = (double)(kCoord - 1) * dz;

            x2 = ((double)iCoord - 0.5) * dx;
            y2 = ((double)jCoord - 0.5) * dy;
            z2 = (double)kCoord * dz;

            double w_n =
                (1.0 / (dx * dy * dz)) *
                ((x2 - x) * (y2 - y) * (z2 - z) * W(iCoord - 1, jCoord - 1, kCoord - 1) +
                    (x - x1) * (y2 - y) * (z2 - z) * W(iCoord, jCoord - 1, kCoord - 1) +
                    (x2 - x) * (y - y1) * (z2 - z) * W(iCoord - 1, jCoord, kCoord - 1) +
                    (x - x1) * (y - y1) * (z2 - z) * W(iCoord, jCoord, kCoord - 1) +
                    (x2 - x) * (y2 - y) * (z - z1) * W(iCoord - 1, jCoord - 1, kCoord) +
                    (x - x1) * (y2 - y) * (z - z1) * W(iCoord, jCoord - 1, kCoord) +
                    (x2 - x) * (y - y1) * (z - z1) * W(iCoord - 1, jCoord, kCoord) +
                    (x - x1) * (y - y1) * (z - z1) * W(iCoord, jCoord, kCoord));

            double new_z                      = z + particletracer->dt * w_n;
            particletracer->particlePool[i].z = new_z;

            // printf("\tOld X : %.2f, New X : %.2f, iCoord : %d\n\tOld Y : %.2f, New Y :
            // %.2f, jCoord : %d\n\n", x, new_x, iCoord, y, new_y, jCoord);
            // printf("\tU(iCoord - 1, jCoord - 1) : %.2f, U(iCoord, jCoord - 1) : %.2f,
            // U(iCoord - 1, jCoord) : %.2f, U(iCoord, jCoord) : %.2f\n", U(iCoord - 1,
            // jCoord - 1), U(iCoord, jCoord - 1), U(iCoord - 1, jCoord), U(iCoord,
            // jCoord)); printf("\tV(iCoord - 1, jCoord - 1) : %.2f, V(iCoord, jCoord - 1)
            // : %.2f, V(iCoord - 1, jCoord) : %.2f, V(iCoord, jCoord) : %.2f\n\n",
            // V(iCoord - 1, jCoord - 1), V(iCoord, jCoord - 1), V(iCoord - 1, jCoord),
            // V(iCoord, jCoord)); printf("\t U N : %.2f, V N : %.2f\n\n", u_n, v_n);
            // printf("\t j-1 * (imax + 2) + i-1 = %d with element from U : %.2f", (jCoord
            // - 1) * (200 + 2) + (iCoord - 1), u[(jCoord - 1) * (imax + 2) + (iCoord -
            // 1)]); printf("\nimax : %d, jmax : %d\n", imax, jmax);

            if (((new_x < 0.0) || (new_x > xlength) || (new_y < 0.0) ||
                    (new_y > ylength) || (new_z < 0.0) || (new_z > zlength))) {
                particletracer->particlePool[i].flag = false;
            }
            int i_new = new_x / dx, j_new = new_y / dy, k_new = new_z / dz;

            if (S(i_new, j_new, k_new) != NONE) {
                particletracer->particlePool[i].flag = false;
            }
        }
    }
}

void freeParticles(ParticleTracer* particletracer)
{
    free(particletracer->particlePool);
    free(particletracer->linSpaceLine);
}

void writeParticles(ParticleTracer* particletracer)
{
    VtkOptions opts = { .particletracer = particletracer };

    char filename[50];
    snprintf(filename, 50, "vtk_files/particles%d.vtk", ts);
    vtkOpenPT(&opts, filename, ts);
    vtkParticle(&opts, "particle");
    vtkClose(&opts);
    ++ts;
}

void initParticleTracer(ParticleTracer* particletracer, Parameter* params)
{
    particletracer->numberOfParticles = params->numberOfParticles;
    particletracer->startTime         = params->startTime;
    particletracer->injectTimePeriod  = params->injectTimePeriod;
    particletracer->writeTimePeriod   = params->writeTimePeriod;

    particletracer->dt = params->dt;
    particletracer->dx = params->xlength / params->imax;
    particletracer->dy = params->ylength / params->jmax;
    particletracer->dz = params->zlength / params->kmax;

    particletracer->xlength = params->xlength;
    particletracer->ylength = params->ylength;
    particletracer->zlength = params->zlength;

    particletracer->x1 = params->x1;
    particletracer->y1 = params->y1;
    particletracer->z1 = params->z1;
    particletracer->x2 = params->x2;
    particletracer->y2 = params->y2;
    particletracer->z2 = params->z2;

    particletracer->lastInjectTime = params->startTime;
    particletracer->lastUpdateTime = params->startTime;
    particletracer->lastWriteTime  = params->startTime;

    particletracer->pointer        = 0;
    particletracer->totalParticles = 0;

    particletracer->imax = params->imax;
    particletracer->jmax = params->jmax;
    particletracer->kmax = params->kmax;

    particletracer->estimatedNumParticles = ((params->te - params->startTime) + 2) *
                                            params->numberOfParticles;

    particletracer->particlePool = malloc(
        sizeof(Particle) * particletracer->estimatedNumParticles);
    particletracer->linSpaceLine = malloc(
        sizeof(Particle) * particletracer->numberOfParticles);

    for (int i = 0; i < particletracer->numberOfParticles; ++i) {
        // double spacing                      = (double)i /
        // (double)(particletracer->numberOfParticles - 1);
        // particletracer->linSpaceLine[i].x   = spacing * particletracer->x1 + (1.0 -
        // spacing) * particletracer->x2; particletracer->linSpaceLine[i].y   = spacing *
        // particletracer->y1 + (1.0 - spacing) * particletracer->y2;
        // particletracer->linSpaceLine[i].z   = spacing * particletracer->z1 + (1.0 -
        // spacing) * particletracer->z2;

        particletracer->linSpaceLine[i].x = particletracer->x1;
        particletracer->linSpaceLine[i].y = (double)rand() / RAND_MAX *
                                            particletracer->ylength;
        particletracer->linSpaceLine[i].z = (double)rand() / RAND_MAX *
                                            particletracer->zlength;

        particletracer->linSpaceLine[i].flag = true;
    }
}

void printParticleTracerParameters(ParticleTracer* particletracer)
{
    printf("Particle Tracing data:\n");
    printf("\tNumber of particles : %d being injected for every period of %.2f\n",
        particletracer->numberOfParticles,
        particletracer->injectTimePeriod);
    printf("\tstartTime : %.2f\n", particletracer->startTime);
    printf("\t(Line along which the particles are to be injected) \n\tx1 : %.2f, y1 : "
           "%.2f, z1 : %.2f, x2 : %.2f, y2 : %.2f, z2 : %.2f\n",
        particletracer->x1,
        particletracer->y1,
        particletracer->z1,
        particletracer->x2,
        particletracer->y2,
        particletracer->z2);
    printf("\tPointer : %d, TotalParticles : %d\n",
        particletracer->pointer,
        particletracer->totalParticles);
    printf("\tdt : %.2f, dx : %.2f, dy : %.2f,  dz : %.2f\n",
        particletracer->dt,
        particletracer->dx,
        particletracer->dy,
        particletracer->dz);
}

void trace(ParticleTracer* particletracer,
    double* u,
    double* v,
    double* w,
    int* seg,
    double time)
{
    if (time >= particletracer->startTime) {
        // printParticles(particletracer);
        if ((time - particletracer->lastInjectTime) >= particletracer->injectTimePeriod) {
            injectParticles(particletracer, seg);
            particletracer->lastInjectTime = time;
        }
        if ((time - particletracer->lastWriteTime) >= particletracer->writeTimePeriod) {
            writeParticles(particletracer);
            particletracer->lastWriteTime = time;
        }
        advanceParticles(particletracer, u, v, w, seg, time);
        compress(particletracer);
        particletracer->lastUpdateTime = time;
    }
}

void compress(ParticleTracer* particletracer)
{
    Particle* memPool = particletracer->particlePool;
    Particle tempPool[particletracer->totalParticles];
    int totalParticles = 0;

    for (int i = 0; i < particletracer->totalParticles; ++i) {
        if (memPool[i].flag == 1) {
            tempPool[totalParticles].x    = memPool[i].x;
            tempPool[totalParticles].y    = memPool[i].y;
            tempPool[totalParticles].z    = memPool[i].z;
            tempPool[totalParticles].flag = memPool[i].flag;
            ++totalParticles;
        }
    }

    particletracer->totalParticles = totalParticles;
    particletracer->pointer        = totalParticles + 1;

    memcpy(particletracer->particlePool, tempPool, totalParticles * sizeof(Particle));
}