/*
 * 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 "parameter.h"
#include "solver.h"
#include "particletracing.h"
#include "util.h"

#define U(i, j)  u[(j) * (imax + 2) + (i)]
#define V(i, j)  v[(j) * (imax + 2) + (i)]

static void print(Solver* solver, double* grid)
{
    int imax = solver->imax;

    for (int j = 0; j < solver->jmax + 2; j++) {
        printf("%02d: ", j);
        for (int i = 0; i < solver->imax + 2; i++) {
            printf("%3.2f  ", grid[j * (imax + 2) + i]);
        }
        printf("\n");
    }
    fflush(stdout);
}

static int ts = 0;
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)
{
    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].flag = true;
        ++(particletracer->pointer);
        ++(particletracer->totalParticles);
    }
}

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


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

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

    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;

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

            double x1 = (double)(iCoord - 1) * dx;
            double y1 = ((double)(jCoord - 1) - 0.5) * dy;
            double x2 = (double)iCoord * dx;
            double y2 = ((double)jCoord - 0.5) * dy;

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

            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;

            x1 = ((double)(iCoord - 1) - 0.5) * dx;
            y1 = (double)(jCoord - 1) * dy;
            x2 = ((double)iCoord - 0.5) * dx;
            y2 = (double)jCoord * dy;

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

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

            // 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)))
            {
                particletracer->particlePool[i].flag = false;
            }
        }
    }
}

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

void writeParticles(ParticleTracer* particletracer)
{
    FILE* fp;
    Particle* particlePool = particletracer->particlePool;

    char filename[50];
    snprintf(filename, 50, "particles_%d.dat", ts++);
    fp = fopen(filename, "w");

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

    for (int i = 0; i < particletracer->totalParticles; ++i)
    {
         if(particlePool[i].flag == true)
        {
            double x = particlePool[i].x;
            double y = particlePool[i].y;
            fprintf(fp, "%f %f\n", x, y);
        }
    }
    fclose(fp);
}

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->xlength           = params->xlength;
    particletracer->ylength           = params->ylength;
    
    particletracer->x1                = params->x1;
    particletracer->y1                = params->y1;
    particletracer->x2                = params->x2;
    particletracer->y2                = params->y2;

    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->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].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, x2 : %.2f, y2 : %.2f\n", particletracer->x1, particletracer->y1, particletracer->x2, particletracer->y2);
    printf("\tPointer : %d, TotalParticles : %d\n", particletracer->pointer, particletracer->totalParticles);
    printf("\tdt : %.2f, dx : %.2f, dy : %.2f\n", particletracer->dt, particletracer->dx, particletracer->dy);
}   

void trace(ParticleTracer* particletracer, Solver* solver, double* u, double* v, double time)
{
    if (time >= particletracer->startTime)
    {
        //printParticles(particletracer);
        if( (time - particletracer->lastInjectTime) >= particletracer->injectTimePeriod)
        {
            injectParticles(particletracer);
            particletracer->lastInjectTime = time;
        }
        if( (time - particletracer->lastWriteTime) >= particletracer->writeTimePeriod)
        {
            writeParticles(particletracer);
            particletracer->lastWriteTime = time;

        }
        advanceParticles(particletracer, solver, u, v, time);
        particletracer->lastUpdateTime = time;   
    }
}