/* Converting Astra output particle distribution to Elegant SDDS compliant input 
 * Version 1.00 by Victor Verzilov, Sincrotrone Trieste
 * 13 Sept 2002
 *
 * update by Mauro Trovo`
 * 17 Gen 2004
 *
 * update S.M. Lidia 
 * 3 Feb 2006
 * 
 * Update by M. Borland
 * 27 July 2009
 * Include betaz in computation of arrival time.
 * Include correction to x and y using dz*xp and dz*yp, respectively.
 */

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

main (int argc,char *argv[])
{
  int i,index,sFlag,dFlag=0;
  FILE *fin,*fout;
  float x_r,y_r,z_r,px_r,py_r,pz_r,p_r,clock_r;
  float x,y,z,px,py,pz,p,clock,mCharge;  
  float dpz;
  float MEC=510999.06;  // rest electron momentum eV/c
  float t0=1.0e-09;   // conversion coefficient from ns to s 
  float clight=2.99792458e08;   //speed of light in m/s
  float betaz, gamma;
  char *arg=NULL,*par=NULL,foutName[255],description[255];
  char *fmt[2]={"%6ld","%14.7e"};
  char *Usage="Usage:   astra2elegant inputfile \n \
          [-output=<filename>] \n \
          [-description=<description string>] \n ";

  /* Analysing arguments */

  if (argc<2){
    printf("astra2elegant:  too few arguments\n");
    printf("%s",Usage);
    exit(1);
  }

  if ((fin=fopen(argv[1],"r"))==NULL){
    printf("\n astra2elegant:  Cannot open file %s \n",argv[1]);
    exit(1);
  }
  if(argc==2){ 
    strcpy(foutName,argv[1]);
    strcat(foutName,".sdds");
  }
  else{
     for(i=2;i<argc;i++){
       if((arg=strchr(argv[i],'-'))==NULL || (par=strchr(arg,'='))==NULL\
	  ||(par-arg)<2){
	  printf("astra2elegant: invalid argument syntax\n");
	  printf("%s",Usage);
	  exit(1);
       }
       arg++;
       if(strncmp("output", arg, strcspn(arg,"="))==0){
	 if(strlen(par)<2){
	   printf("astra2elegant: empty output file name\n");
      	   printf("%s",Usage);
	   exit(1);
	 }
	 strcpy(foutName,++par);
       }
       else if(strncmp("description", arg, strcspn(arg,"="))==0){
	 if(strlen(par)<2){
	   printf("astra2elegant: empty description\n");
	   printf("%s",Usage);
	   exit(1);
	 }
       	   strcpy(description,++par);
	   dFlag=1;
       }	
       else {
	  printf("astra2elegant: invalid argument\n");
	  printf("%s",Usage);
	  exit(1);
	  }
     }
   }
  

    if ((fout=fopen(foutName,"w"))==NULL){
    printf("\n astra2elegant:  Cannot create output file %s \n\n",foutName);
    }
    /*  Writting an SDDS-style header  */
    fprintf(fout,"SDDS1 \n");
    if (dFlag==0)
    fprintf(fout,"&description text=\"%s\", &end\n",foutName);
    else
    fprintf(fout,"&description text=\"%s\", &end\n",description);
    fprintf(fout,"&column name=ID, description=\"particle index\",");
    fprintf(fout," format_string=%s, type=long,  &end\n",fmt[0]);
    fprintf(fout,"&column name=x, symbol=x, units=m, description=\"horizontal position\",");
    fprintf(fout," format_string=%s, type=double,  &end\n",fmt[1]);
    fprintf(fout,"&column name=xp, symbol=x', description=\"horizontal slope\",");
    fprintf(fout," format_string=%s, type=double,  &end\n",fmt[1]);
    fprintf(fout,"&column name=y, symbol=y, units=m, description=\"vertical position\",");
    fprintf(fout," format_string=%s, type=double,  &end\n",fmt[1]);
    fprintf(fout,"&column name=yp, symbol=y', description=\"vertical slope\",");
    fprintf(fout," format_string=%s, type=double,  &end\n",fmt[1]);
    fprintf(fout,"&column name=t, symbol=t, units=s, description=\"time\",");
    fprintf(fout," format_string=%s, type=double,  &end\n",fmt[1]);
    fprintf(fout,"&column name=p, symbol=p, units=m$be$nc, description=\"momentum\",");
    fprintf(fout," format_string=%s, type=double,  &end\n",fmt[1]);
    fprintf(fout,"&data mode=ascii, no_row_counts=1, &end\n");
    fprintf(fout,"! page number 1\n");

    i=1;
 
    /* Reading the reference particle  */

    fscanf(fin,"%E%E%E%E%E%E%E%E%d%d",&x_r,&y_r,&z_r,&px_r,&py_r,&pz_r,&clock_r,&mCharge,&index,&sFlag);
       if(sFlag!=5){
	 printf(" There isn't the reference particle in the input file\n");
	 exit(1);
       }
      p_r = sqrt ((double)(px_r*px_r+py_r*py_r+pz_r*pz_r));    /* total momentum */
 
    /*  Writing the reference particle  */

	fprintf(fout,fmt[0],i); fprintf(fout," ");
	fprintf(fout,fmt[1],x_r); fprintf(fout," "); fprintf(fout,fmt[1],px_r/p_r); fprintf(fout," ");
	fprintf(fout,fmt[1],y_r); fprintf(fout," "); fprintf(fout,fmt[1],py_r/p_r); fprintf(fout," ");
	fprintf(fout,fmt[1],0.0); fprintf(fout," "); fprintf(fout,fmt[1],p_r/MEC); fprintf(fout,"\n");

    /* Reading a successive row  */
  
	/* only longitudinal variables are relative to reference particle
	   transverse coordinates are absolute - SML */
    while(fscanf(fin,"%E%E%E%E%E%E%E%E%d%d",&x,&y,&z,&px,&py,&dpz,&clock,&mCharge,&index,&sFlag)!=EOF){
      i++;
      pz=pz_r+dpz;
      p=sqrt ( (double)(px*px+py*py+pz*pz)); /*  total momentum */

      if(p!=0 && (sFlag==3 || sFlag==5) ){

  /* Writing a successive row */

	fprintf(fout,fmt[0],i); fprintf(fout," ");
	fprintf(fout,fmt[1],x-z*px/pz); fprintf(fout," "); fprintf(fout,fmt[1],px/pz); fprintf(fout," ");
	fprintf(fout,fmt[1],y-y*px/pz); fprintf(fout," "); fprintf(fout,fmt[1],py/pz); fprintf(fout," ");

/* NOTE:                                                                                                */
/* The next row set the column "time" in the sdds file and there is the choice between the columns      */
/* of the Astra file: Z or T                                                                            */
/*                                                                                                      */
/* case A:                                                                                              */  
/*	fprintf(fout,fmt[1],clock*t0); fprintf(fout," "); fprintf(fout,fmt[1],p/MEC); fprintf(fout,"\n"); */
/* case B:                                                                                              */
        
        gamma = sqrt(p*p/(MEC*MEC)+1);
        betaz = (pz/MEC)/gamma;
	fprintf(fout,fmt[1],clock_r/1e9-z/(clight*betaz)); fprintf(fout," "); fprintf(fout,fmt[1],p/MEC); fprintf(fout,"\n");
      }
    }

  fclose(fout); 
  return 0;
}

/* Auxiliary functions */

char *str_Copy(char *src)
{
  char *dst;
  printf("src=%s length=%d\n",src,strlen(src));
  dst=malloc(strlen(src)*sizeof(char));
  printf("dst=%s\n",dst);
  strcpy(dst,src);
  printf("dst=%s\n",dst);
  return dst;
}

char *str_Concat(char *src1,char *src2)
{
  char *dst;

    dst=malloc((strlen(src1)+strlen(src2))*sizeof(char));
    strcpy(dst,src1);
    strcat(dst,src2);
    return dst;
}