/* ///////////////////////////////////////////////////////////////////// */
/*!
 \file
 \brief Solve 1D heat equation.
 
 Solve the 1D heat equation using a 1st order explicit method
 on a parallel domain.
 
 \author A. Mignone (mignone@to.infn.it)
 \date   March 12, 2020
 */
/* ///////////////////////////////////////////////////////////////////// */
#include <stdio.h>
#include <stdlib.h>
#include <math.h>

#define PARALLEL

#ifdef PARALLEL
#include <mpi.h>
#endif

#define NX_GLOB   64  /* Global number of interior points */
#define NGHOST     1

void Write (double *, double *, int, int);

int main(int argc, char ** argv)
{
    int    i, k, beg, end;
    int    nx_loc;               /* Local grid size */
    int    dstL = -1, dstR=-1;   /* Rank of left and right neighbour procs */
    int    rank=0, size=1;
    double t, tstop, dt, cfl = 0.5;
    double *u0;
    double *u1;
    double xbeg =  0.0;
    double xend = +1.0;
    double xglob[NX_GLOB + 2*NGHOST];   // Global grid array
    double *xloc;
    double dx;    /* Mesh spacing */
#ifdef PARALLEL
    double *send_buf;
    double *recv_buf;
#endif
    FILE *fp;
    
    /* --------------------------------------------------------
     0. Initialize parallel environment & get neighbour
     proc rank
     -------------------------------------------------------- */
    
#ifdef PARALLEL
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);
    
    dstL   = rank - 1;
    dstR   = rank + 1;
    if (dstL < 0)     dstL = MPI_PROC_NULL;
    if (dstR >= size) dstR = MPI_PROC_NULL;
#endif
    
    /* --------------------------------------------------------
     1. Generate global & local grids
     -------------------------------------------------------- */
    
#ifdef PARALLEL
    nx_loc = NX_GLOB/size;
    beg    = NGHOST;
    end    = beg + nx_loc - 1;
    
    dx = (xend - xbeg)/(NX_GLOB+1);
    for (i = 0; i < NX_GLOB + 2*NGHOST; i++){
        xglob[i] = xbeg + (i-beg+1)*dx;
    }
    xloc = xglob + nx_loc*rank;  /* Use pointer arithmetic */
#else
    nx_loc = NX_GLOB;
    beg    = NGHOST;
    end    = beg + nx_loc - 1;
    
    dx = (xend - xbeg)/(NX_GLOB+1);
    for (i = 0; i < NX_GLOB + 2*NGHOST; i++){
        xglob[i] = xbeg + (i-beg+1)*dx;
    }
    xloc = xglob;  /* Use pointer arithmetic */
#endif
    
    /* --------------------------------------------------------
     2. Allocate memory on local grids
     -------------------------------------------------------- */
    
    u0 = (double *) malloc((nx_loc + 2*NGHOST)*sizeof(double));
    u1 = (double *) malloc((nx_loc + 2*NGHOST)*sizeof(double));
    
#ifdef PARALLEL
    {
        int proc, go;
        for (proc = 0; proc < size; proc++){
            go = proc;
            MPI_Bcast(&go, 1, MPI_INT, 0, MPI_COMM_WORLD);
            if (rank == go) {
                printf ("[Rank %d]\n",rank);
                printf ("  dstL = %d, dstR = %d\n",dstL, dstR);
                printf ("  beg, end = %d, %d; x = [%f, %f]\n",
                        beg, end, xloc[beg],xloc[end]);
            }
            MPI_Barrier(MPI_COMM_WORLD);
        }
    }
#endif
    
    /* --------------------------------------------------------
     3. Set initial condition
     -------------------------------------------------------- */
    
    for (i = beg; i <= end; i++){
        u0[i] = sin(M_PI*xloc[i]);
    }
    
    /* --------------------------------------------------------
     4. Advance solution
     -------------------------------------------------------- */
    
    t     = 0.0;
    tstop = 0.1;
    dt    = cfl*dx*dx;
    k     = 0;
    
    Write (xloc, u0, beg, end);
    while (t < tstop){
        
        if (rank == 0){
            printf ("step #%d; t = %8.3e\n",k,t);
        }
        
        /* -- 4a. Set physical boundary conditions -- */
        
        if (dstL == MPI_PROC_NULL){
            u0[beg-1] = 0.0;
        }
        if (dstR == MPI_PROC_NULL){
            u0[end+1] = 0.0;
        }
        
        /* -- 4b. Set inter-process boundary conditions -- */
        
#ifdef PARALLEL
        send_buf = u0 + end - (NGHOST - 1);  // Address of rightmost interior point
        recv_buf = u0 + 0;                   // Address of leftmost ghost zone
        MPI_Sendrecv (send_buf, NGHOST, MPI_DOUBLE, dstR, 0,
                      recv_buf, NGHOST, MPI_DOUBLE, dstL, 0,
                      MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        
        send_buf = u0 + beg;     // Address of leftmost interior point
        recv_buf = u0 + end + 1; // Address of first ghost zone on the right
        MPI_Sendrecv (send_buf, NGHOST, MPI_DOUBLE, dstL, 0,
                      recv_buf, NGHOST, MPI_DOUBLE, dstR, 0,
                      MPI_COMM_WORLD, MPI_STATUS_IGNORE);
#endif
        
        /* -- 4c. Advance solution by one time step -- */
        
        for (i = beg; i <= end; i++){
            u1[i] = u0[i] + dt/(dx*dx)*(u0[i-1] - 2.0*u0[i] + u0[i+1]);
        }
        t += dt;
        k++;
        
        /* -- 4d. Copy arrays for next time level -- */
        
        for (i = beg; i <= end; i++) u0[i] = u1[i];
    }
    Write (xloc, u0, beg, end);
    
#ifdef PARALLEL
    MPI_Finalize();
#endif
    return 0;
}

/* ********************************************************************* */
void Write (double *x, double *u, int beg, int end)
/*
 *********************************************************************** */
{
    int    i;
    int    rank;
    static int n = 0;  /* File number */
    FILE *fp;
    char fname[32];
    
    /* --------------------------------------------------------
     1. Serial output
     -------------------------------------------------------- */
    
#ifndef PARALLEL
    sprintf (fname,"heat_eq%02d.dat",n);
    fp = fopen (fname,"w");
    for (i = beg; i <= end; i++) fprintf (fp, "%12.6e  %12.6e\n", x[i], u[i]);
    fclose(fp);
#endif
    
    /* --------------------------------------------------------
     2. Parallel output
     -------------------------------------------------------- */
    
#ifdef PARALLEL
    
    /* -- 2a. Process #0 gathers data and does the writing -- */
    
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    int nx_loc = end - beg + 1;
    static double *recv_buf;
    if (recv_buf == NULL) {
        recv_buf = (double *) malloc((NX_GLOB + 2*NGHOST)*sizeof(double));
    }
    
    MPI_Gather (u + beg, nx_loc, MPI_DOUBLE,
                recv_buf + beg, nx_loc, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    if (rank == 0){
        sprintf (fname,"heat_eq%02d.dat",n);
        fp = fopen (fname,"w");
        for (i = beg; i < beg+NX_GLOB; i++) {
            fprintf (fp, "%f  %f\n", x[i], recv_buf[i]);
        }   
        fclose(fp);
    }
    
    /* -- 2b. Shared file pointer -- */
    
    /* -- 2c. Individual file pointer -- */
    
#endif
    
    n++;
}



/*
 MAPLE Script:
 
 restart;
 u := A*exp(-D*mu^2*t)*sin(mu*x + B) + C;
 eq := diff(u,t) - D*diff(diff(u,x),x);
 simplify(eq);
 
 
 */
