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

#define KE 201

int main(int argc, char **argv) {
  MPI_Status status;
  int a, b;
  int size;
  int n, k, kc=KE/2, NSTEPS=140;
  double *ex, *hy, *rex, *rhy;
  double ex_low_m1 = 0.0, ex_low_m2 = 0.0;
  double ex_high_m1 = 0.0, ex_high_m2 = 0.0;
  int myrank, isize;
  double t_str, elapsed_time;
  double T = 0.0, t0 = 40.0, spread = 12.0, pulse;
  FILE *fp;

  MPI_Init(&argc, &argv);

  MPI_Comm_rank(MPI_COMM_WORLD, &myrank);
  MPI_Comm_size(MPI_COMM_WORLD, &isize);
  t_str = MPI_Wtime();

  // このノードが受け持つ空間サイズの決定
  size = ((KE-2)%isize==0 ? (KE-2)/isize : (KE-2)/isize+1) + 2;
  ex = malloc(sizeof(double)*size);
  hy = malloc(sizeof(double)*size);

  for(k=0; k<size; k++) {
    ex[k] = 0.0;
    hy[k] = 0.0;
  }

  for(n=1; n<=NSTEPS; n++) {
    T = (double)n;
    for(k=1; k<size-1; k++) {
      ex[k] = ex[k]+0.5*(hy[k-1]-hy[k]);
    }

    if((size-2)*myrank<=kc && kc<(size-2)*(myrank+1)) {
      pulse = (t0-T)/spread;
      pulse = -0.5 * pulse * pulse;
      pulse = exp(pulse);
      ex[kc-(size-2)*myrank+1] = pulse;
    }
    
    MPI_Sendrecv(&ex[size-2], 1, MPI_DOUBLE, (myrank+1)%isize, 0,
		 &ex[0], 1, MPI_DOUBLE, (myrank-1+isize)%isize, 0,
		 MPI_COMM_WORLD, &status);
    MPI_Sendrecv(&ex[1], 1, MPI_DOUBLE, (myrank-1+isize)%isize, 1,
		 &ex[size-1], 1, MPI_DOUBLE, (myrank+1)%isize, 1,
		 MPI_COMM_WORLD, &status);

    MPI_Sendrecv(&hy[size-2], 1, MPI_DOUBLE, (myrank+1)%isize, 2,
		 &hy[0], 1, MPI_DOUBLE, (myrank-1+isize)%isize, 2,
		 MPI_COMM_WORLD, &status);
    MPI_Sendrecv(&hy[1], 1, MPI_DOUBLE, (myrank-1+isize)%isize, 3,
		 &hy[size-1], 1, MPI_DOUBLE, (myrank+1)%isize, 3,
		 MPI_COMM_WORLD, &status);
    
    if(myrank==0) {
      ex[0] = ex_low_m2;
      ex_low_m2 = ex_low_m1;
      ex_low_m1 = ex[1];
    }

    if(myrank==isize-1) {
      ex[KE-1] = ex_high_m2;
      ex_high_m2 = ex_high_m1;
      ex_high_m1 = ex[KE-2];
    }

    for(k=0; k<size-1; k++) {
      hy[k] = hy[k] + 0.5 * (ex[k]-ex[k+1]);
    }
    
  }

  if(myrank==0) {
    rex = malloc(sizeof(double)*KE);
    rhy = malloc(sizeof(double)*KE);
  }

  MPI_Gather(ex+1, size-2, MPI_DOUBLE, rex, size-2, MPI_DOUBLE, 0, MPI_COMM_WORLD);
  MPI_Gather(hy+1, size-2, MPI_DOUBLE, rhy, size-2, MPI_DOUBLE, 0, MPI_COMM_WORLD);

  if(myrank==0) {
    fp = fopen("Ex.dat", "w");
    for(k=0; k<KE; k++)
      fprintf(fp, "%d %6.2f \n", k, rex[k]);
    fclose(fp);
  
    fp = fopen("Hy.dat", "w");
    for(k=0; k<KE; k++)
      fprintf(fp, "%d %6.2f \n", k, rhy[k]);
    fclose(fp);
  }

  free(ex);
  free(hy);

  if(myrank==0) {
    free(rex);
    free(rhy);
  }
  
  elapsed_time = MPI_Wtime() - t_str;
  printf("Elapsed time : %f (msec)\n", 1.0e3*elapsed_time);

  return MPI_Finalize();
}