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

#include "trace.h"
#define U(i, j) u[(j) * (imax + 2) + (i)]
#define V(i, j) v[(j) * (imax + 2) + (i)]

static int ts = 0;

static void printState(Tracing* t)
{
    printf("Cursor: %d Total particles: %d\n", t->cursor, t->totalParticles);
}

static void advanceParticles(
    Tracing* t, double delt, double* restrict u, double* restrict v)
{
    double delx = t->grid.dx;
    double dely = t->grid.dy;

    double* m = t->memorypool;
    int* p    = t->particles;
    int imax  = t->grid.imax;
    int jmax  = t->grid.jmax;

    for (int i = 0; i < t->totalParticles; i++) {
        int particleId = p[i];

        double x = m[particleId * NCOORD + X];
        double y = m[particleId * NCOORD + Y];
        // printf("P%d - X %f Y %f\n", i, x, y);

        // Interpolate U
        int iCoord = (int)(x / delx) + 1;
        int jCoord = (int)((y + 0.5 * dely) / dely) + 1;

        double x1 = (double)(iCoord - 1) * delx;
        double y1 = ((double)(jCoord - 1) - 0.5) * dely;
        double x2 = (double)iCoord * delx;
        double y2 = ((double)jCoord - 0.5) * dely;

        // printf("U - iCoord %d jCoord %d\n", iCoord, jCoord);

        double un = (1.0 / (delx * dely)) *
                    ((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 xn                  = x + delt * un;
        m[particleId * NCOORD + X] = xn;

        // Interpolate V
        iCoord = (int)((x + 0.5 * delx) / delx) + 1;
        jCoord = (int)(y / dely) + 1;

        x1 = ((double)(iCoord - 1) - 0.5) * delx;
        y1 = (double)(jCoord - 1) * dely;
        x2 = ((double)iCoord - 0.5) * delx;
        y2 = (double)jCoord * dely;

        // printf("V - iCoord %d jCoord %d\n", iCoord, jCoord);

        double vn = (1.0 / (delx * dely)) *
                    ((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 yn                  = y + delt * vn;
        m[particleId * NCOORD + Y] = yn;
    }

    double xlength = t->grid.xlength;
    double ylength = t->grid.ylength;
    int cntNew     = 0;
    int tmp[t->totalParticles];

    // Check for particles to remove
    for (int i = 0; i < t->totalParticles; i++) {
        int particleId = p[i];

        double x = m[particleId * NCOORD + X];
        double y = m[particleId * NCOORD + Y];

        if (!((x < 0.0) || (x > xlength) || (y < 0.0) || (y > ylength))) {
            tmp[cntNew++] = i;
        }
    }

    t->totalParticles = cntNew;
    memcpy(t->particles, tmp, cntNew * sizeof(int));
}

static void injectParticles(Tracing* t)
{
    double* line = t->line;
    double* m    = t->memorypool;

    for (int i = 0; i < t->numParticles; i++) {
        t->particles[t->totalParticles] = t->cursor;
        m[(t->cursor) * NCOORD + X]     = line[i * NCOORD + X];
        m[(t->cursor) * NCOORD + Y]     = line[i * NCOORD + Y];
        t->cursor++;
        t->totalParticles++;
    }
}

static void writeParticles(Tracing* t)
{
    FILE* fp;
    double* m = t->memorypool;
    int* p    = t->particles;

    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 < t->totalParticles; i++) {
        int particleId = p[i];

        double x = m[particleId * NCOORD + X];
        double y = m[particleId * NCOORD + Y];
        fprintf(fp, "%f %f\n", x, y);
    }
    fclose(fp);
}

void trace(Tracing* t, double* restrict u, double* restrict v, double time)
{
    if (time >= t->traceStart) {
        if ((time - t->lastUpdate[INJECT]) > t->traceInject) {
            injectParticles(t);
            t->lastUpdate[INJECT] = time;
        }

        if ((time - t->lastUpdate[WRITE]) > t->traceWrite) {
            writeParticles(t);
            t->lastUpdate[WRITE] = time;
        }

        advanceParticles(t, time - t->lastUpdate[ADVANCE], u, v);
        t->lastUpdate[ADVANCE] = time;
    }
}

void initTrace(Tracing* t, Parameter* p)
{
    size_t numParticles   = p->nparticles;
    size_t totalParticles = (size_t)(p->te - p->traceStart) / (size_t)p->traceInject;
    totalParticles += 2;
    totalParticles *= numParticles;

    double x1 = p->lineX1;
    double y1 = p->lineY1;
    double x2 = p->lineX2;
    double y2 = p->lineY2;

    for (int i = 0; i < NUMTIMERS; i++) {
        t->lastUpdate[i] = p->traceStart;
    }
    t->grid.imax      = p->imax;
    t->grid.jmax      = p->jmax;
    t->grid.xlength   = p->xlength;
    t->grid.ylength   = p->ylength;
    t->grid.dx        = p->xlength / p->imax;
    t->grid.dy        = p->ylength / p->jmax;
    t->numParticles   = numParticles;
    t->totalParticles = 0;
    t->cursor         = 0;
    t->traceStart     = p->traceStart;
    t->traceWrite     = p->traceWrite;
    t->traceInject    = p->traceInject;
    t->particles      = (int*)malloc(totalParticles * sizeof(int));
    t->memorypool     = (double*)malloc(totalParticles * NCOORD * sizeof(double));
    t->line           = (double*)malloc(numParticles * NCOORD * sizeof(double));
    double* line      = t->line;

    for (int i = 0; i < numParticles; i++) {
        double spacing = (double)i / (double)(numParticles - 1);
        double x       = spacing * x1 + (1.0 - spacing) * x2;
        double y       = spacing * y1 + (1.0 - spacing) * y2;

        //printf("S: %f x: %f y: %f\n", spacing, x, y);
        line[i * NCOORD + X] = x;
        line[i * NCOORD + Y] = y;
    }
}

void freeTrace(Tracing* t) { free(t->line); }