/*******************************************************************************
*
* McXtrace, x-ray tracing package
*         Copyright, All rights reserved
*         DTU Physics, Kgs. Lyngby, Denmark
*         Synchrotron SOLEIL, Saint-Aubin, France
*
* Component: PSD_monitor_coh
*
* %Identification
* Written by: Erik Knudsen
* Date: March 13, 2010
* Origin: Risoe
*
* Position-sensitive monitor with phase integration.
*
* %Description
* An (n times m) pixel PSD monitor taking phase into account.
* As the i:th ray hits a pixel (j,k) in the monitor the intensity in that
* pixel will be updated as a complex sum, i.e. <math>P_i = P_{i-1} + p_i exp{-\phi_i}</math>.
*
* By setting ratio<1 the effective pixel area becomes a
* fraction of the ideal (which is to divide the xwidth and yheight intervals into nx and ny abutting
* subintervals). This reduces the monitor effective area by ratio^2.
* If the centering flag is set - the monitor will treat all rays as if they hit a pixel
* center. This behaves as if ratio -> 0, but at no cost in statistics.
*
* Example: PSD_monitor_coh(xwidth=0.1, yheight=0.1,
*            nx=90, ny=90, filename="Output.psd")
*
* %Parameters
* INPUT PARAMETERS:
*
* xwidth:       [m] Width of detector. 
* yheight:      [m] Height of detector.
* nx:           [ ] Number of pixel columns.
* ny:           [ ] Number of pixel rows.
* filename:     [m] Name of file in which to store the detector images - the suffixes .abs and .arg will be added.
* restore_xray: [ ] If set, the monitor does not influence the xray state.
* ratio:        [ ] ratio between pixel area and effective pixel area.
* centering:    [ ] Treat all rays as if they hit the center of the pixel.
* nowritefile:  [1] If set, monitor will skip writing to disk
*
* CALCULATED PARAMETERS:
*
* PSD_N:    Array of xray event counts.
* PSD_phi:  Array of xray complex weight counts (the argument consituttes the phase).
* PSD_p2:   Array of second ms.
*
* %End
*******************************************************************************/

DEFINE COMPONENT PSD_monitor_coh

SETTING PARAMETERS (int nx=90, int ny=90, string filename=0, restore_xray=0, 
  xwidth=0.05, yheight=0.05, ratio=1, int centering=1, int nowritefile=0)

/*STATE PARAMETERS (x,y,z,kx,ky,kz,phi,Ex,Ey,Ez,p)*/
SHARE
%{
#include <complex.h>
%}

DECLARE
%{
    DArray2d PSD_N;
    DArray2d PSD_p2;
    DArray2d PSD_phi_real;
    DArray2d PSD_phi_imag;
    double xmin;
    double xmax;
    double ymin;
    double ymax;
%}

INITIALIZE
  %{
    int i,j;

    xmax = xwidth/2;  xmin = -xmax;
    ymax = yheight/2; ymin = -ymax;

    if ((xmin >= xmax) || (ymin >= ymax)) {
      fprintf(stderr,"ERROR (%s): Null detection area! Aborting.\n",NAME_CURRENT_COMP);
      exit(-1);
    }

    PSD_N = create_darr2d(nx,ny);
    PSD_p2 = create_darr2d(nx,ny);
    PSD_phi_real = create_darr2d(nx,ny);
    PSD_phi_imag = create_darr2d(nx,ny);

    // Use instance name for monitor output if no input was given
    if (!strcmp(filename,"\0")) sprintf(filename,"%s",NAME_CURRENT_COMP);
%}

TRACE
%{
    int i,j;
    double k,x_o,y_o,z_o;
    double p2;
    double complex dphi;
    double dphi_r,dphi_i;

    x_o=x;y_o=y;z_o=z;
    k=sqrt(kx*kx+ky*ky+kz*kz);
    PROP_Z0;
    if (x>xmin && x<xmax && y>ymin && y<ymax)
    {

      double dl,dlc,dnu,cx,cy;

      i = floor((x - xmin)*nx/(xmax - xmin));
      j = floor((y - ymin)*ny/(ymax - ymin));

      /*center coordinates of the pixel*/
      if(centering){
          dl=sqrt( (x-x_o)*(x-x_o) + (y-y_o)*(y-y_o) + z_o*z_o);

          cx=xwidth/nx*(i+0.5)-xwidth/2.0;
          cy=yheight/ny*(j+0.5)-yheight/2.0;
          /*what would the phase be if we transported it to the center of the pixel?*/
          dlc=sqrt((cx-x_o)*(cx-x_o) + (cy-y_o)*(cy-y_o) + z_o*z_o);
          dnu=(dl-dlc)*k*1e10;
      }else{
          dnu=0;
      }

      /*check if within the ratio of the pixel*/
      if ( fabs(xwidth/nx*(i+0.5)-xwidth/2.0-x)<xwidth/nx*ratio && fabs(yheight/ny*(j+0.5)-yheight/2.0-y)<yheight/ny*ratio ){
        dphi=p*cexp(I*(phi-dnu));
	dphi_r=creal(dphi);
	dphi_i=cimag(dphi);
        p2 = p*p;
#pragma acc atomic
        PSD_N[i][j]+=1;
#pragma acc atomic
        PSD_phi_real[i][j]+= dphi_r;
#pragma acc atomic
        PSD_phi_imag[i][j]+= dphi_i;
#pragma acc atomic
        PSD_p2[i][j] += p2;
          SCATTER;
      }
    }
    if (restore_xray) {
      RESTORE_XRAY(INDEX_CURRENT_COMP, x, y, z, kx, ky, kz, phi, t,  Ex, Ey, Ez, p);
    }
%}

SAVE
%{
  if(!nowritefile){
    char filename_abs[255],filename_arg[255];

    double *PSD_p_abs, *PSD_p_arg;
    int i,j;

    snprintf(filename_abs,255,"%s.abs",filename);
    snprintf(filename_arg,255,"%s.arg",filename);

    /*This should be done differently if MPI*/ 
    PSD_p_abs=calloc(nx*ny,sizeof(double));
    PSD_p_arg=calloc(nx*ny,sizeof(double));
    if(PSD_p_abs==NULL || PSD_p_arg==NULL){
        fprintf(stderr,"Error (%s): Buffer memory allocation error\nAborting\n.",NAME_CURRENT_COMP);exit(-1);
        exit(-1);
    }

    for (i=0;i<nx;i++){
      for (j=0;j<ny;j++){
        PSD_p_abs[i*ny +j]=cabs(PSD_phi_real[i][j] + I*PSD_phi_imag[i][j]);
        PSD_p_arg[i*ny +j]=carg(PSD_phi_real[i][j] + I*PSD_phi_imag[i][j]);
      }
    }

#ifdef USE_MPI
    double *PSD_p_absbuf,*PSD_p_argbuf;
    double *N_buf;
    if (mpi_node_rank==0){
        PSD_p_absbuf=calloc(mpi_node_count*ny,sizeof(double));
        PSD_p_argbuf=calloc(mpi_node_count*ny, sizeof(double));
        N_buf=calloc(mpi_node_count*ny,sizeof(double));
    }else{
        PSD_p_absbuf=NULL;
        PSD_p_absbuf=NULL;
        N_buf=NULL;
    }

    for (i=0;i<nx;i++){
        MPI_Gather(&(PSD_p_abs[i*ny]), ny, MPI_DOUBLE, PSD_p_absbuf, ny, MPI_DOUBLE, 0,MPI_COMM_WORLD);
        MPI_Gather(&(PSD_p_arg[i*ny]), ny, MPI_DOUBLE, PSD_p_argbuf, ny, MPI_DOUBLE, 0,MPI_COMM_WORLD);
        MPI_Gather(&(PSD_N[i][0]),ny, MPI_DOUBLE, N_buf, ny, MPI_DOUBLE, 0,MPI_COMM_WORLD);
        if (mpi_node_rank==0){
          for (j=0;j<ny;j++){
              double complex tmp_phi=0;
              unsigned long tmp_N=0;
              int k;
              for (k=0;k<mpi_node_count;k++){
                  tmp_phi+=PSD_p_absbuf[k*ny+j]*cexp(I*PSD_p_argbuf[k*ny+j]);
                  tmp_N+=(unsigned long)N_buf[k*ny+j];
              }
              PSD_p_abs[i*ny+j]=cabs(tmp_phi);
              PSD_p_arg[i*ny+j]=carg(tmp_phi);
              PSD_N[i][j]= tmp_N;
          }
        }else{
            /*set all the other nodes' data to zero. The DETECTOR_OUT macros will do an unnecessary gather,
              and so we don't want to add things twice.*/
          for (j=0;j<ny;j++){
            PSD_p_abs[i*ny+j]=0;
            PSD_p_arg[i*ny+j]=0;
            PSD_N[i][j]=0;
          }
        }
    }
    /*doing it this way is inefficient since we send a lot of zeros across the network - but it should work*/ 
#endif
      
    char fname1[256];
    snprintf(fname1, 256, "%s_arg", NAME_CURRENT_COMP);

    // we use mcdetector_out_nD in order to update the [comp_name]_I symbol for the 2nd output
    if(nx==1 && ny==1){
        DETECTOR_OUT_0D("Coh. Intensity monitor (magnitude)", 
          (double) PSD_N[0][0], PSD_p_abs[0], PSD_p2[0][0]);
        
        
        mcdetector_out_0D("Coh. Phase monitor (argument)",
          (double) PSD_N[0][0], PSD_p_arg[0], PSD_p2[0][0], fname1,POS_A_CURRENT_COMP,ROT_A_CURRENT_COMP);
    }else if(nx==1){
        DETECTOR_OUT_1D(
                "Coh. PSD_monitor (magnitude)","Y Position[m]", "Intensity", "Y",
                ymin,ymax,ny,PSD_N[0],PSD_p_abs,PSD_p2[0],filename_abs);
                
        mcdetector_out_1D(
                "Coh. PSD_monitor (argument)","Y Position[m]", "Phase", "Y",
                ymin,ymax,ny,PSD_N[0],PSD_p_arg,PSD_p2[0],filename_arg,
                fname1,POS_A_CURRENT_COMP,ROT_A_CURRENT_COMP);
    }else if (ny==1){
        DETECTOR_OUT_1D(
                "Coh. PSD_monitor (magnitude)","X Position[m]", "Intensity", "X",
                xmin,xmax,nx,PSD_N[0],PSD_p_abs,PSD_p2[0],filename_abs);
        mcdetector_out_1D(
                "Coh. PSD_monitor (argument)","X Position[m]", "Phase", "X",
                xmin,xmax,nx,PSD_N[0],PSD_p_arg,PSD_p2[0],filename_arg,
                fname1,POS_A_CURRENT_COMP,ROT_A_CURRENT_COMP);
    }else{
        /*don't do the internal gather in DETECTOR_OUT_2D*/
        //MPI_NOGATHER;
        DETECTOR_OUT_2D(
                "Coh. PSD_monitor (magnitude)",
                "X position [m]","Y position [m]",
                xmin, xmax, ymin, ymax, nx, ny,
                PSD_N[0],PSD_p_abs,PSD_p2[0],
                filename_abs);

        /*don't do the internal gather in DETECTOR_OUT_2D*/
        //MPI_NOGATHER;
        mcdetector_out_2D(
                "Coh. PSD_monitor (argument)",
                "X position [m]","Y position [m]",
                xmin, xmax, ymin, ymax, nx, ny,
                (PSD_N[0]),PSD_p_arg,PSD_p2[0],
                filename_arg,
                fname1,POS_A_CURRENT_COMP,ROT_A_CURRENT_COMP);
    }
  }
%}

FINALLY
%{
  destroy_darr2d(PSD_N);
  destroy_darr2d(PSD_p2);
  destroy_darr2d(PSD_phi_real);
  destroy_darr2d(PSD_phi_imag);
%}

MCDISPLAY
%{
  
  multiline(5, (double)xmin, (double)ymin, 0.0,
               (double)xmax, (double)ymin, 0.0,
               (double)xmax, (double)ymax, 0.0,
               (double)xmin, (double)ymax, 0.0,
               (double)xmin, (double)ymin, 0.0);
%}

END