#include "geostat.h"
#include "file_array.h"
#include <math.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include "log.h"
#include "memory.h"

/*FAST FOURIER TRANSFORM MOVING AVERAGE METHOD */
/*Turns a Gaussian white noise vector into a   */
/*spatially correlated vector                  */
/*input:                                       */
/*variogram: structure defining the variogram  */
/*           model                             */
/*grid: structure defining the grid            */
/*n: vector with the number of cells along the */
/*   X, Y and Z axes for the underlying grid   */
/*   i = [0 1 2]                               */
/*   --> 0 0 0        : n will be computed and */
/*   updated as output                         */
/*   --> nx ny nz: these dimensions are used   */
/*realin: structure defining a realization -   */
/*        must be a Gaussian white noise       */
/*output:                                      */
/*realout: structure defining a realization -  */

void FFTMA2(struct vario_mod variogram, struct grid_mod grid, int n[3], struct realization_mod* realin, struct realization_mod* realout, long* seed, int cores) {
    double* used_ram_t0 = malloc(sizeof(double));
    getVirtualMemUsed(used_ram_t0);

    clock_t t = clock();

    log_info("RESULT = in progress");

    struct cpustat initial[cores];
    struct cpustat final[cores];

    for (int i = 0; i < cores; i++) {
        get_stats(&initial[i], i - 1);
    }
    
    int NTOT, i, j, k, NMAX, NDIM, ntot, nmax, NXYZ, nxyz;
    int solver;
    double temp;
    file_array_t *covar, *ireal, *realization;
    double *workr, *worki;

    /*covariance axis normalization*/
    axes(variogram.ap, variogram.scf, variogram.Nvario);

    /*pseudo-grid definition*/
    cgrid(variogram, grid, n, cores);

    /*constant definition*/
    NTOT = n[0] * n[1] * n[2];
    ntot = NTOT + 1;
    NMAX = n[0];
    NDIM = 3;
    for (i = 1; i < 3; i++) {
        if (n[i] > NMAX)
            NMAX = n[i];
        if (n[i] == 1)
            NDIM--;
    }
    nmax = NMAX + 1;
    NXYZ = grid.NX * grid.NY * grid.NZ;
    nxyz = NXYZ + 1;

    /*array initialization*/
    covar = file_array_create("covar.txt", ntot);
    ireal = file_array_create("ireal.txt", ntot);
    realization = file_array_create("realization.txt", ntot);

    workr = (double*)malloc(nmax * sizeof(double));
    testmemory(workr);

    worki = (double*)malloc(nmax * sizeof(double));
    testmemory(worki);

    /*covariance function creation*/
    covariance(covar, variogram, grid, n, cores);

    /*power spectrum*/
    fourt(covar, ireal, n, NDIM, 1, 0, workr, worki, cores);

    /*organization of the input Gaussian white noise*/
    solver = 0;
    prebuild_gwn(grid, n, realin, realization, solver, seed, cores);

    /*forward fourier transform of the GWN*/
    fourt(realization, ireal, n, NDIM, 1, 0, workr, worki, cores);

    /* build realization in spectral domain   */
    build_real(n, NTOT, covar, realization, ireal, cores);

    double* used_ram_tf = malloc(sizeof(double));
    getVirtualMemUsed(used_ram_tf);

    file_array_free(covar);
    remove("covar.txt");

    /*backward fourier transform*/
    fourt(realization, ireal, n, NDIM, 0, 1, workr, worki, cores);

    file_array_free(ireal);
    remove("ireal.txt");

    free(workr);
    free(worki);

    /*output realization*/
    clean_real(realin, n, grid, realization, realout, cores);

    t = clock() - t;
    double time_taken = ((double)t)/CLOCKS_PER_SEC; // calculate the elapsed time

    for (int i = 0; i < cores; i++) {
        get_stats(&final[i], i - 1);
    }

    for (int i = 0; i < cores; i++) {
        log_info("CPU %d: %lf%%", i, calculate_load(&initial[i], &final[i]));
    }

    log_info("RESULT = success, NTOT = %d, NMAX = %d, NDIM = %d, ntot = %d, nmax = %d, NXYZ = %d, nxyz = %d, ELAPSED = %f seconds, DIF USED VIRTUAL MEM = %5.1f MB", NTOT, NMAX, NDIM, ntot, nmax, NXYZ, nxyz, time_taken, *used_ram_tf - *used_ram_t0);

    free(used_ram_t0);
    free(used_ram_tf);
}