#define _XOPEN_SOURCE 600
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <math.h>
#include <sys/time.h>
#include <sys/stat.h>
#include <errno.h>
#include <inttypes.h>

#include "argument_utils.h"


// Convert 'struct timeval' into seconds in double prec. floating point
#define WALLTIME(t) ((double)(t).tv_sec + 1e-6 * (double)(t).tv_usec)

// Option to change numerical precision
typedef int64_t int_t;
typedef double real_t;

// Simulation parameters: size, step count, and how often to save the state
int_t
    N = 256,
    M = 256,
    max_iteration = 4000,
    snapshot_freq = 20;

// Wave equation parameters, time step is derived from the space step
const real_t
    c  = 1.0,
    dx = 1.0,
    dy = 1.0;
real_t
    dt;

// Buffers for three time steps, indexed with 2 ghost points for the boundary
real_t
    *buffers[3] = { NULL, NULL, NULL };

#define U_prv(i,j) buffers[0][((i)+1)*(N+2)+(j)+1]
#define U(i,j)     buffers[1][((i)+1)*(N+2)+(j)+1]
#define U_nxt(i,j) buffers[2][((i)+1)*(N+2)+(j)+1]


// Rotate the time step buffers.
void move_buffer_window ( void )
{
    real_t *temp = buffers[0];
    buffers[0] = buffers[1];
    buffers[1] = buffers[2];
    buffers[2] = temp;
}


// Save the present time step in a numbered file under 'data/'
void domain_save ( int_t step )
{
    char filename[256];

    // Ensure output directory exists (ignore error if it already exists)
    if (mkdir("data", 0755) != 0 && errno != EEXIST) {
        perror("mkdir data");
        exit(EXIT_FAILURE);
    }

    snprintf(filename, sizeof(filename), "data/%05" PRId64 ".dat", step);

    FILE *out = fopen(filename, "wb");
    if (out == NULL) {
        perror("fopen output file");
        fprintf(stderr, "Failed to open '%s' for writing.\n", filename);
        exit(EXIT_FAILURE);
    }

    for ( int_t i = 0; i < M; ++i ) {
        size_t written = fwrite ( &U(i,0), sizeof(real_t), (size_t)N, out );
        if ( written != (size_t)N ) {
            perror("fwrite");
            fclose(out);
            exit(EXIT_FAILURE);
        }
    }

    if ( fclose(out) != 0 ) {
        perror("fclose");
        exit(EXIT_FAILURE);
    }
}



// Neumann (reflective) boundary condition
void boundary_condition ( void )
{
    // Vertical ghost layers (bottom and top)
    for (int_t i=0; i<M; i++) {
        U(i,-1) = U(i,1);       // bottom ghost row <- mirror of row 1
        U(i,N)  = U(i,N-2);     // top ghost row <- mirror of row N-2
    }

    // Horizontal ghost layers (left and right)
    for (int_t j=0; j<N; j++) {
        U(-1,j) = U(1,j);       // left ghost col <- mirror of col 1
        U(M,j)  = U(M-2,j);     // right ghost col <- mirror of col M-2
    }

    // Corner ghost cells (use ghost indices)
    U(-1,-1) = U(1,1);             // bottom-left
    U(-1,N)  = U(1,N-2);           // top-left
    U(M,-1)  = U(M-2,1);           // bottom-right
    U(M,N)   = U(M-2,N-2);         // top-right
}


// Set up our three buffers, and fill two with an initial perturbation
void domain_initialize ( void )
{
    buffers[0] = malloc ( (M+2)*(N+2)*sizeof(real_t) );
    buffers[1] = malloc ( (M+2)*(N+2)*sizeof(real_t) );
    buffers[2] = malloc ( (M+2)*(N+2)*sizeof(real_t) );

    if ( !buffers[0] || !buffers[1] || !buffers[2] ) {
        perror("malloc");
        exit(EXIT_FAILURE);
    }

    // initialize interior (physical cells)
    for ( int_t i=0; i<M; i++ )
    {
        for ( int_t j=0; j<N; j++ )
        {
            real_t dx_i = (i - M/2.0);
            real_t dy_j = (j - N/2.0);
            real_t delta = sqrt( (dx_i*dx_i) / (real_t)M + (dy_j*dy_j) / (real_t)N );
            U_prv(i,j) = U(i,j) = exp ( -4.0*delta*delta );
        }
    }

    // Set the time step (CFL) with a safety factor
    dt = 1.0 / ( c * sqrt( 1.0/(dx*dx) + 1.0/(dy*dy) ) );
    dt *= 0.9; // safety factor

    // Fill ghost cells once so they are valid before the first step
    boundary_condition();
}


// Get rid of all the memory allocations
void domain_finalize ( void )
{
    free ( buffers[0] );
    free ( buffers[1] );
    free ( buffers[2] );
}


// Integration formula (Eq. 9 from the pdf document)
void time_step ( void )
{
    double coef_x  = (c * dt) * (c * dt) / (dx * dx);
    double coef_y  = (c * dt) * (c * dt) / (dy * dy);
    double coef_xy = (c * dt) * (c * dt) / (6.0 * dx * dy);

    // Only iterate over physical domain cells i = 0..M-1, j = 0..N-1
    for (int i = 0; i < M; i++)
    {
        for (int j = 0; j < N; j++)
        {
            U_nxt(i,j) = 2.0*U(i,j) - U_prv(i,j)
                + coef_x  * (U(i+1,j) - 2.0*U(i,j) + U(i-1,j))
                + coef_y  * (U(i,j+1) - 2.0*U(i,j) + U(i,j-1))
                + coef_xy * (U(i+1,j+1) + U(i+1,j-1) +
                             U(i-1,j+1) + U(i-1,j-1) - 4.0*U(i,j));
        }
    }
}


// Main time integration.
void simulate( void )
{
    // Go through each time step
    for ( int_t iteration=0; iteration<=max_iteration; iteration++ )
    {
        if ( (iteration % snapshot_freq)==0 )
        {
            domain_save ( iteration / snapshot_freq );
        }

        // Derive step t+1 from steps t and t-1
        boundary_condition();
        time_step();

        // Rotate the time step buffers
        move_buffer_window();
    }
}


int main ( int argc, char **argv )
{
    OPTIONS *options = parse_args( argc, argv );
    if ( !options )
    {
        fprintf( stderr, "Argument parsing failed\n" );
        exit( EXIT_FAILURE );
    }

    M = options->M;
    N = options->N;
    max_iteration = options->max_iteration;
    snapshot_freq = options->snapshot_frequency;

    if (M < 2 || N < 2)
    {
        fprintf(stderr, "M and N must be >= 2\n");
        exit(EXIT_FAILURE);
    }


    // Set up the initial state of the domain
    domain_initialize();

    struct timeval t_start, t_end;

    gettimeofday ( &t_start, NULL );
    simulate();
    gettimeofday ( &t_end, NULL );

    printf ( "Total elapsed time: %lf seconds\n",
        WALLTIME(t_end) - WALLTIME(t_start)
    );

    // Clean up and shut down
    domain_finalize();
    exit ( EXIT_SUCCESS );
}