// interference.cc
// 
// Calculation of Diffraction Images

#include <cstdlib>
#include <iostream>
#include <fstream>
#include <complex>
#include <vector>
#include <string>
#include <sstream>
#include <algorithm>
#include <climits>
#include "mpi.h"

using namespace std;

const double two_pi=6.283185307179586;

// This structure maps the screen 
typedef struct {
  double x1, y1, x2, y2;
} screen;

// This structure maps the screen resolution
typedef struct {
  int x, y;
} resolution;

void read_input_file(const string &filename, double &dist, screen &S,
		     resolution &res, double &lambda,
		     vector<double> &x, vector<double> &y,
		     vector<double> &amplitude) {
  // extremely simple parser
  ifstream in(filename.c_str());
  do {
    string token;
    in >> token;
    if (in) {
      if (token=="dist")
	in >> dist;
      else if (token=="screen")
	in >> S.x1 >> S.y1 >> S.x2 >> S.y2;
      else if (token=="picture")
	in >> res.x >> res.y;
      else if (token=="lambda")
	in >> lambda;
      else if (token=="source") {
	double t_x, t_y, t_amplitude;
	in >> t_x >> t_y >> t_amplitude;
	x.push_back(t_x);
	y.push_back(t_y);
	amplitude.push_back(t_amplitude);
      };
      in.ignore(INT_MAX, '\n');
      // if #include <limits>
      // in.ignore(numeric_limits<int>::max(), '\n'); preferable
    }
  } while (in);
}

void write_output_file(const string &filename, resolution res,
		       const vector<double> &intensity,
		       const double max_intensity) {
  // output data to a pgm file (portable graymap file format), brightest
  // point on screen is colored white, all others gray depending on brightness
  ofstream out(filename.c_str());
  out << "P2" << endl  // magic number
      << res.x << " " << res.y << endl  // width and height 
      << 256 << endl;  // number of grayscales
  for (int i=0; i<res.x*res.y; ++i) 
    out << floor(intensity[i]/max_intensity*255.0+0.5) << endl;
}

void calc_intensity(const double dist, const screen S, const resolution res,
		    const double lambda, const vector<double> x, 
		    const vector<double> y, const vector<double> amplitude, 
		    vector<double> &intensity) {
  // calculate intensity for each screen coordinate (x, y) 
  intensity.resize(res.x*res.y);
  double dx=S.x2-S.x1;
  double dy=S.y2-S.y1;
  double dist2=dist*dist;
  for (int j=0; j<res.y; ++j) {
    double Y=S.y1+(j+0.5)*dy/res.y;
    for (int i=0; i<res.x; ++i) {
      double X=S.x1+(i+0.5)*dx/res.x;
      complex<double> comp_int=(0.0, 0.0);
      for (int k=0; k<x.size(); ++k) { // Sum for all sources 
	double d=sqrt((X-x[k])*(X-x[k])+(Y-y[k])*(Y-y[k])+dist2);
	comp_int+=amplitude[k]*
	  exp(complex<double>(0.0, two_pi*d/lambda))/d;
      }
      // intensity is the square of the magnitude of the complex amplitude
      intensity[j*res.x+i]=norm(comp_int);
    }
  }
}

int main(int argc, char *argv[]) {

  MPI::Init();  // initialize MPI

  int rank=MPI::COMM_WORLD.Get_rank();  // ascertain own rank
  int size=MPI::COMM_WORLD.Get_size();  // and number of processes

  // Assign default values to parameters 
  double dist=1.0;  // Distance from sources to screen 
  screen S={-0.05, -0.05,   // lower left corner of screen 
	    +0.05, +0.05};  // top right corner of screen 
  resolution res={250, 250};   // Number of coordinates in x and y direction
  double lambda=0.4e-6;  // Wave length of radiation
  vector<double> x, y, amplitude;
  int num_sources;

  if (rank==0) {  // Rank 0 reads the configuration file
    if (argc==2) 
      read_input_file(argv[1], dist, S, res, lambda, x, y, amplitude);
    num_sources=x.size();
    if (num_sources==0)  // no calculation possible without sources 
      MPI::COMM_WORLD.Abort(EXIT_FAILURE);
  }
  
  // Process 0 distributes parameters to all other processes
  MPI::COMM_WORLD.Bcast(&dist, 1, MPI::DOUBLE, 0);
  MPI::COMM_WORLD.Bcast(&S.x1, 1, MPI::DOUBLE, 0);
  MPI::COMM_WORLD.Bcast(&S.y1, 1, MPI::DOUBLE, 0);
  MPI::COMM_WORLD.Bcast(&S.x2, 1, MPI::DOUBLE, 0);
  MPI::COMM_WORLD.Bcast(&S.y2, 1, MPI::DOUBLE, 0);
  MPI::COMM_WORLD.Bcast(&res.x, 1, MPI::INT, 0);
  MPI::COMM_WORLD.Bcast(&res.y, 1, MPI::INT, 0);
  MPI::COMM_WORLD.Bcast(&lambda, 1, MPI::DOUBLE, 0);    
  MPI::COMM_WORLD.Bcast(&num_sources, 1, MPI::INT, 0);
  // all processes except process 0 need to allocate memory first
  if (rank!=0) {
    x.resize(num_sources);
    y.resize(num_sources);
    amplitude.resize(num_sources);
  }
  MPI::COMM_WORLD.Bcast(&x[0], num_sources, MPI::DOUBLE, 0);
  MPI::COMM_WORLD.Bcast(&y[0], num_sources, MPI::DOUBLE, 0);
  MPI::COMM_WORLD.Bcast(&amplitude[0], num_sources, MPI::DOUBLE, 0);

  // round y resolution so that value is divisible by number of processes
  if (res.y%size>0)
    res.y=res.y-res.y%size+size;

  // resolution to be calculated by each process
  resolution res_local={res.x, res.y/size};

  // area to be calculated by each process
  screen S_local={S.x1, rank*(S.y2-S.y1)/size+S.y1,
		  S.x2, (rank+1)*(S.y2-S.y1)/size+S.y1};
  
  // calculations start here
  vector<double> intensity_local;
  calc_intensity(dist, S_local, res_local, lambda, x, y, amplitude, 
		 intensity_local);

  // ascertain maximum intensity, first locally, then globally 
  double max_intensity, max_intensity_local=
    *max_element(intensity_local.begin(), intensity_local.end());
  MPI::COMM_WORLD.Reduce(&max_intensity_local, &max_intensity, 1, MPI::DOUBLE, 
			 MPI::MAX, 0);

  // Process 0 collects the results...
  vector<double> intensity;
  intensity.resize(res.x*res.y);
  MPI::COMM_WORLD.Gather(&intensity_local[0], res_local.x*res_local.y, 
			 MPI::DOUBLE,
			 &intensity[0], res_local.x*res_local.y,
			 MPI::DOUBLE, 0);
  // ... and writes them to a file 
  if (rank==0)
    write_output_file(string(argv[1])+".pgm", res, intensity, max_intensity);
  
  MPI::Finalize();  // terminate MPI

  return EXIT_SUCCESS;
}