#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <mpi.h>

double ind2pos(int i, int n, double L) {
    double ret_val;
    ret_val = ((i-1)-(n-1)/ 2.0)*L/(n-1);
    return ret_val;
}

int save_gnuplot(double* temp, int n, double L, char* filename, int ix_start, int iy_start,int ix_size, int iy_size,
                 MPI_Comm cartesianComm, int nprocs, int rank) {
// save the temperature distribution
// the ascii format is suitable for splot gnuplot function
    int ix,iy,iproc;
    FILE *fp;
    double* temp_glob, *buffer_mpi;
    int ix_start_r,iy_start_r,ix_size_r,iy_size_r;
    
    temp_glob = calloc(n*n, sizeof(*temp));

    if(rank == 0) {
        for(iy=1;iy<=iy_size;++iy)
            for(ix=1;ix<=ix_size;++ix)
                 temp_glob[(ix_start+ix-1)*n+iy+iy_start-1] = temp[ix*(iy_size+2)+(iy)];
        for(iproc = 1;iproc< nprocs; iproc++) { 
            MPI_Recv(&ix_start_r, 1, MPI_INT, iproc, 400, cartesianComm, MPI_STATUS_IGNORE);
            MPI_Recv(&iy_start_r, 1, MPI_INT, iproc, 401, cartesianComm, MPI_STATUS_IGNORE);
            MPI_Recv(&ix_size_r, 1, MPI_INT, iproc, 402, cartesianComm, MPI_STATUS_IGNORE);
            MPI_Recv(&iy_size_r, 1, MPI_INT, iproc, 403, cartesianComm, MPI_STATUS_IGNORE);
            buffer_mpi = calloc(ix_size_r*iy_size_r, sizeof(*temp));
            MPI_Recv(buffer_mpi, ix_size_r*iy_size_r, MPI_DOUBLE, iproc, 404, cartesianComm, MPI_STATUS_IGNORE);
            for(iy=1;iy<=iy_size_r;++iy)      
                for(ix=1;ix<=ix_size_r;++ix)
                     temp_glob[(ix_start_r+ix-1)*n+iy+iy_start_r-1] = buffer_mpi[(ix-1)*(iy_size_r)+(iy-1)];
            free(buffer_mpi);
        }
        fp = fopen(filename, "w");
        for(iy=1;iy<=n;++iy){      
            for(ix=1;ix<=n;++ix)
                fprintf(fp, "\t%f\t%f\t%g\n", ind2pos(ix,n,L), ind2pos(iy,n,L), temp_glob[n*(ix-1)+iy-1]);
            fprintf(fp, "\n");
        }
        fclose(fp);
    } else {
        MPI_Send(&ix_start, 1, MPI_INT, 0, 400, cartesianComm);
        MPI_Send(&iy_start, 1, MPI_INT, 0, 401, cartesianComm);
        MPI_Send(&ix_size,  1, MPI_INT, 0, 402, cartesianComm);
        MPI_Send(&iy_size,  1, MPI_INT, 0, 403, cartesianComm);
        buffer_mpi = calloc(ix_size*iy_size, sizeof(*temp));
        for(iy=1;iy<=iy_size;++iy)
            for(ix=1;ix<=ix_size;++ix)
                buffer_mpi[(ix-1)*(iy_size)+(iy-1)] = temp[ix*(iy_size + 2)+iy];
        MPI_Send(buffer_mpi, ix_size*iy_size, MPI_DOUBLE, 0, 404, cartesianComm);
        free(buffer_mpi);
    }

    free(temp_glob);
    
    return 0;
}

int init_field(double *temp, int n, int L, int ix_start, int iy_start, int ix_size, int iy_size) {
    // initialize the T field
    int ix,iy;
    double x, y;
    const double sigma = 0.1;
    const double tmax = 100.;
    
    for(iy=0;iy<=iy_size+1;++iy)
        for(ix=0;ix<=ix_size+1;++ix){
            x=ind2pos(ix+ix_start, n, L);
            y=ind2pos(iy+iy_start, n, L);
            temp[ix*(iy_size + 2)+iy] = tmax*exp((-((x*x)+(y*y)))/(2.0*(sigma*sigma)));
        }
    return 0;
}