/*
 * 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)  u[(j) * (imax + 2) + (i)]
#define V(i, j)  v[(j) * (imax + 2) + (i)]

static int ts = 0;

#define IDIM 0
#define JDIM 1

#define XOFFSET 0
#define YOFFSET 1
#define XOFFSETEND 2
#define YOFFSETEND 3

static double sumX(double* sizes, int size, int offset, int coord, int init)
{
    double sum = 0;

    for (int i = init; i < size && coord > 0; i+=offset, --coord) 
    {
        sum += sizes[i];
    }

    return sum;
}

static double sumY(double* sizes, int size, int offset, int coord, int init)
{
    double sum = 0;

    for (int i = init * offset; i < size && coord > 0; ++i, --coord) 
    {
        sum += sizes[i];
    }

    return sum;
}

void printParticles(ParticleTracer* particletracer)
{
    for(int i = 0; i < particletracer->totalParticles; ++i)
    {   
        printf("Rank : %d Particle position X : %.2f, Y : %.2f, flag : %d, total pt : %d, pointer : %d, xOffset : %.2f, yOffset : %.2f, xOffsetEnd : %.2f, yOffsetEnd : %.2f\n", 
            particletracer->rank, 
            particletracer->particlePool[i].x,
            particletracer->particlePool[i].y,
            particletracer->particlePool[i].flag, 
            particletracer->totalParticles, 
            particletracer->pointer, 
            particletracer->xOffset, particletracer->yOffset, 
            particletracer->xOffsetEnd, particletracer->yOffsetEnd);
    }
}
void injectParticles(ParticleTracer* particletracer)
{
    double x, y;
    for(int i = 0; i < particletracer->numberOfParticles; ++i)
    {
        x = particletracer->linSpaceLine[i].x;
        y = particletracer->linSpaceLine[i].y;
        if(x >= particletracer->xOffset && y >= particletracer->yOffset && x < particletracer->xOffsetEnd && y < particletracer->yOffsetEnd )
        {
            particletracer->particlePool[particletracer->pointer].x = x;
            particletracer->particlePool[particletracer->pointer].y = y;
            particletracer->particlePool[particletracer->pointer].flag = true;
            ++(particletracer->pointer);
            ++(particletracer->totalParticles);
        }
    }
}

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

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

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

    Particle buff[particletracer->size][30];

    for(int i = 0; i < particletracer->size; ++i)
    {   
        for(int j = 0; j < 30; ++j)
        { 
            buff[i][j].x = 0.0;
            buff[i][j].y = 0.0;
            buff[i][j].flag = false;
        }
    }
    int particleBufIndex[particletracer->size];

    memset(particleBufIndex, 0, sizeof(particleBufIndex));


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

            double x   = xTemp - particletracer->xOffset;
            double y   = yTemp - particletracer->yOffset;

            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->xOffset) + 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->yOffset) + 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 < particletracer->xOffset) || (new_x >= particletracer->xOffsetEnd) ||
                 (new_y < particletracer->yOffset) || (new_y >= particletracer->yOffsetEnd)))
            {
                //New logic to transfer particles to neighbouring ranks or discard the particle.

                for(int i = 0; i < particletracer->size; ++i)
                {
                    if((new_x >= particletracer->offset[i + particletracer->size * XOFFSET]) && 
                       (new_x < particletracer->offset[i + particletracer->size * XOFFSETEND]) && 
                       (new_y >= particletracer->offset[i + particletracer->size * YOFFSET]) && 
                       (new_y < particletracer->offset[i + particletracer->size * YOFFSETEND]) &&
                       i != particletracer->rank)
                    {
                        buff[i][particleBufIndex[i]].x = new_x;
                        buff[i][particleBufIndex[i]].y = new_y;
                        buff[i][particleBufIndex[i]].flag = true;
                        ++particleBufIndex[i];
                    }
                }
                particletracer->particlePool[i].flag = false;
            }
        }
    }

    for(int i = 0; i < particletracer->size; ++i)
    {   
        if(i != particletracer->rank)
        {
            MPI_Send(buff[i], 30, particletracer->mpi_particle, i, 0, particletracer->comm);
        }
    }
    for(int i = 0; i < particletracer->size; ++i)
    {   
        if(i != particletracer->rank)
        {
            MPI_Recv(buff[i], 30, particletracer->mpi_particle, i, 0, particletracer->comm, MPI_STATUS_IGNORE);
        }
    }
    for(int i = 0; i < particletracer->size; ++i)
    {   
        if(i != particletracer->rank)
        {
            for(int j = 0; j < 30; ++j)
            {
                if(buff[i][j].flag == true)
                {
                    particletracer->particlePool[particletracer->pointer].x = buff[i][j].x;
                    particletracer->particlePool[particletracer->pointer].y = buff[i][j].y;
                    particletracer->particlePool[particletracer->pointer].flag = true;
                    ++(particletracer->pointer);
                    ++(particletracer->totalParticles);
                }
            }
        }
    }
}

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

void writeParticles(ParticleTracer* particletracer)
{
    
    if(particletracer->rank == 0)
    {
        int overallTotalParticles;
        MPI_Reduce(&particletracer->totalParticles, &overallTotalParticles, 1, MPI_INT, MPI_SUM, 0, particletracer->comm);

        printf("Total collected particles at rank 0 : %d", overallTotalParticles);
        

        VtkOptions opts = { .particletracer = particletracer, .overallTotalParticles = overallTotalParticles };

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

        
        FILE* fp;
        Particle* particlePool = particletracer->particlePool;

        snprintf(filename, 50, "vis_files/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)
        {
            double x = particlePool[i].x;
            double y = particlePool[i].y;
            fprintf(fp, "%f %f\n", x, y);
        }
        fclose(fp);
    }

    ++ts;
}

void initParticleTracer(ParticleTracer* particletracer, Parameter* params, Solver* solver)
{
    int dims[NDIMS]    = { 0, 0 };
    int periods[NDIMS] = { 0, 0 };

    /* initializing local properties from 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->imaxLocal         = solver->imaxLocal;
    particletracer->jmaxLocal         = solver->jmaxLocal;

    particletracer->estimatedNumParticles = (particletracer->imaxLocal * particletracer->jmaxLocal);

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

    for(int i = 0; i < particletracer->estimatedNumParticles; ++i)
    {   
        particletracer->particlePool[i].x = 0.0;
        particletracer->particlePool[i].y = 0.0;
        particletracer->particlePool[i].flag = false;
    }

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

    /* duplicating communication from solver */
    MPI_Comm_dup(solver->comm, &particletracer->comm);

    particletracer->rank              = solver->rank;
    particletracer->size              = solver->size;

    particletracer->offset            = (double *)malloc(sizeof(double) * 4 * particletracer->size);


    memcpy(particletracer->dims, solver->dims, sizeof(solver->dims));

    memcpy(particletracer->coords, solver->coords, sizeof(solver->coords));

    memcpy(particletracer->iNeighbours, solver->iNeighbours, sizeof(solver->iNeighbours));

    memcpy(particletracer->jNeighbours, solver->jNeighbours, sizeof(solver->jNeighbours));


    particletracer->xLocal           = particletracer->imaxLocal * particletracer->dx;
    particletracer->yLocal           = particletracer->jmaxLocal * particletracer->dy;


    double xLocal[particletracer->size];
    double yLocal[particletracer->size];

    double offset[4][particletracer->size];
    
    MPI_Allgather(&particletracer->xLocal, 1, MPI_DOUBLE, xLocal, 1, MPI_DOUBLE, particletracer->comm);
    MPI_Allgather(&particletracer->yLocal, 1, MPI_DOUBLE, yLocal, 1, MPI_DOUBLE, particletracer->comm);

    particletracer->xOffset     = sumX(xLocal, particletracer->size, particletracer->dims[JDIM], particletracer->coords[IDIM], particletracer->coords[JDIM]);
    particletracer->yOffset     = sumY(yLocal, particletracer->size, particletracer->dims[JDIM], particletracer->coords[JDIM], particletracer->coords[IDIM]);
    particletracer->xOffsetEnd  = particletracer->xOffset + particletracer->xLocal;
    particletracer->yOffsetEnd  = particletracer->yOffset + particletracer->yLocal;

    MPI_Allgather(&particletracer->xOffset, 1, MPI_DOUBLE, offset[0], 1, MPI_DOUBLE, particletracer->comm);
    MPI_Allgather(&particletracer->yOffset, 1, MPI_DOUBLE, offset[1], 1, MPI_DOUBLE, particletracer->comm);
    MPI_Allgather(&particletracer->xOffsetEnd, 1, MPI_DOUBLE, offset[2], 1, MPI_DOUBLE, particletracer->comm);
    MPI_Allgather(&particletracer->yOffsetEnd, 1, MPI_DOUBLE, offset[3], 1, MPI_DOUBLE, particletracer->comm);

    memcpy(particletracer->offset, offset, sizeof(offset));

    
    // if(particletracer->rank == 0)
    // {   
    //     for(int i = 0;i < particletracer->size; ++i)
    //     {
    //         printf("Rank : %d and its xOffset : %.2f, yOffset : %.2f, xLocal : %.2f, yLocal : %.2f, xOffsetEnd : %.2f, yOffsetEnd : %.2f\n", 
    //         i, particletracer->offset[i + particletracer->size * XOFFSET],
    //         particletracer->offset[i + particletracer->size * YOFFSET],  xLocal[i], yLocal[i], 
    //         particletracer->offset[i + particletracer->size * XOFFSETEND], 
    //         particletracer->offset[i + particletracer->size * YOFFSETEND]);
    //     }
    // }
    
    

    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;
    }

    // Create the mpi_particle datatype
    MPI_Datatype mpi_particle;
    int lengths[3] = { 1, 1, 1 };
 
    MPI_Aint displacements[3];
    Particle dummy_particle;
    MPI_Aint base_address;
    MPI_Get_address(&dummy_particle, &base_address);
    MPI_Get_address(&dummy_particle.x, &displacements[0]);
    MPI_Get_address(&dummy_particle.y, &displacements[1]);
    MPI_Get_address(&dummy_particle.flag, &displacements[2]);
    displacements[0] = MPI_Aint_diff(displacements[0], base_address);
    displacements[1] = MPI_Aint_diff(displacements[1], base_address);
    displacements[2] = MPI_Aint_diff(displacements[2], base_address);
 
    MPI_Datatype types[3] = { MPI_DOUBLE, MPI_DOUBLE, MPI_C_BOOL };
    MPI_Type_create_struct(3, lengths, displacements, types, &particletracer->mpi_particle);
    MPI_Type_commit(&particletracer->mpi_particle);

}

void printParticleTracerParameters(ParticleTracer* particletracer)
{
    printf("Particle Tracing data:\n");
    printf("Rank : %d\n", particletracer->rank);
    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);
    // printf("\tcoord[0] : %d, coord[1] : %d\n", particletracer->coords[IDIM], particletracer->coords[JDIM]);
    // printf("\txOffset : %.2f, yOffset : %.2f\n", particletracer->xOffset, particletracer->yOffset);
    // printf("\txLocal : %.2f, yLocal : %.2f\n", particletracer->xLocal, particletracer->yLocal);
}   

void trace(ParticleTracer* particletracer, 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);
            printParticles(particletracer);
            particletracer->lastWriteTime = time;

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

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

    for(int i = 0; i < particletracer->totalParticles; ++i)
    {   
        tempPool[i].x = 0.0;
        tempPool[i].y = 0.0;
        tempPool[i].flag = false;
    }
    int totalParticles = 0;

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

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

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