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

#define LINEPIX 2048
#define SATULATE 60000
#define SATULATE2 5000
#define MARK_SATULATE -1000000
#define MARK_BAD      -1100000
#define BAD_PIXEL     0
#define COSMIC_RAY    100
#define ZSHIFT 1000000

void printerror( int status);

int main(int argc,char *argv[])
{
  fitsfile *fp;
  int i,j,status,nfound,anynull,bitpix=-32;
  long fpixel, nelements,npixels,naxis=2;
  long naxes[2]={LINEPIX, LINEPIX},nbuffer;
  float nullval;
  static unsigned short buffer[LINEPIX];
  static int zs[LINEPIX][LINEPIX],xzs[LINEPIX][LINEPIX],cr[LINEPIX][LINEPIX];
  static double x2zs[LINEPIX][LINEPIX],zzs[LINEPIX][LINEPIX];
  int ihdu,ihdumax,itype,sx,sx2,sx3,sigma,zex,zre,ncr,sg=1;
  char buf[256];
  double a,b,c,sx4,sigma2,sigma3,stddev;
  FILE *fpt;

  if( argc != 3 ){
    fprintf( stderr, "Usage: %s input output\n", argv[0]);
    exit(1);
  }

  if((fpt=fopen(argv[2],"rb"))!=NULL){
    fprintf(stderr,"Output file (%s) exists.\n",argv[2]);
    fclose(fpt);
    exit(0);
  }

  fprintf(stderr,"Read %s  ... ",argv[1]);
  npixels=naxes[0]*naxes[1];
  status=0;
  if(fits_open_file(&fp, argv[1], READONLY, &status))
    printerror(status);
  if(fits_read_key(fp, TSTRING, "DET-NSMP", buf, NULL, &status)){
    status=0;
    if(fits_read_key(fp, TSTRING, "AC_NEXP", buf, NULL, &status))
      printerror(status);
    sg=-1;
  }
  sscanf(buf,"%d",&ihdumax);ihdumax-=5;
  //fprintf(stderr,"ihdumax=%d\n",ihdumax);

  for(i=0;i<LINEPIX;i++){
    for(j=0;j<LINEPIX;j++){
      zs[i][j]=0;
      xzs[i][j]=0;
      x2zs[i][j]=0.;
      zzs[i][j]=0.;
      cr[i][j]=0;
    }
  }

  sx=0;sx2=0;sx3=0;sx4=0.;sigma=0;sigma2=0.;sigma3=0.;
  for(ihdu=0;ihdu<ihdumax;ihdu++){
    if(fits_movabs_hdu(fp, ihdu+2+5, &itype, &status) )
      printerror(status);
    if(fits_read_keys_lng(fp, "NAXIS", 1, 2, naxes, &nfound, &status))
      printerror(status);
    if(naxes[0]!=LINEPIX || naxes[1]!=LINEPIX){
      fprintf(stderr,"Wrong format (%d %ld %ld).\n",ihdu,naxes[0],naxes[1]);
      exit(0);
    }
    fpixel = 1;
    nullval = 0;
    nbuffer = naxes[0];
    ncr=0;
    for(i=0;i<naxes[1];i++){
      if (fits_read_img(fp, TUSHORT, fpixel, nbuffer, &nullval, buffer, &anynull, &status))
	printerror(status);
      for(j=0;j<naxes[0];j++){
	//if(j==1133 && i==1147) fprintf(stderr,"%d %d %d %d\n",ihdu,buffer[j],zs[i][j],xzs[i][j]);
	if(xzs[i][j]>MARK_SATULATE){
	  a=(((double)xzs[i][j]*ihdu-(double)zs[i][j]*sx)*sigma2-(x2zs[i][j]*ihdu-(double)zs[i][j]*sx2)*sigma)/(sigma2*sigma2-sigma3*sigma);
	  b=((double)xzs[i][j]*ihdu-(double)zs[i][j]*sx-a*sigma2)/sigma;
	  c=(zs[i][j]-a*sx*sx/ihdu-b*sx)/ihdu;
	  if((sg==1&&buffer[j]<SATULATE)||(sg==-1&&buffer[j]>SATULATE2)){
	    zex=(int)(a*ihdu*ihdu+b*ihdu+c);
	    zre=(int)buffer[j];
	    stddev=sqrt((zzs[i][j]-2*a*x2zs[i][j]-2*b*xzs[i][j]-2*c*zs[i][j]+a*a*sx4+2*a*b*sx3+(2*a*c+b*b)*sx2+2*b*c*sx)/ihdu+c*c);
	    //if(j==1133 && i==1147) fprintf(stderr,"%d %d %d %g %g %d %d %d %g %d %g %g %g %g %d %d %d %d %d %d\n",ihdu,zs[i][j],xzs[i][j],x2zs[i][j],zzs[i][j],sx,sx2,sx3,sx4,sigma,sigma2,sigma3,a,b,(int)c,zre,zex,cr[i][j],zex-zre,(int)stddev);
	    if(ihdu>3 && cr[i][j]<=ZSHIFT/2){
	      if(ihdu>4 && sg*(zre-zex)>COSMIC_RAY && sg*(zre-zex)-cr[i][j]>COSMIC_RAY && sg*(zre-zex)>10*stddev && sg*(zre-zex)-cr[i][j]>10*stddev){
                //fprintf(stderr,"%d %d %d %d %d %d %d %d %d %d\n",ihdu,j,i,(int)(b*ihdu),(int)c,zre,zex,cr[i][j],-(zre-zex),(int)stddev);
		cr[i][j]=sg*(zre-zex)+ZSHIFT;ncr++;
	      }else{
		cr[i][j]=sg*(zre-zex);
	      }
	    }
	    //if(i==0&&j==0)getchar();
	    if(cr[i][j]>ZSHIFT/2) zre=(int)buffer[j]-(cr[i][j]-ZSHIFT)*sg;
	    else zre=(int)buffer[j];
	    zs[i][j]+=zre;
	    xzs[i][j]+=ihdu*zre;
	    x2zs[i][j]+=ihdu*ihdu*zre;
	    zzs[i][j]+=(double)zre*zre;
	  }else{
	    if(sigma==0){
	      zs[i][j]=BAD_PIXEL;
	      xzs[i][j]=MARK_BAD;
	    }else{
	      zs[i][j]=(int)(b*(ihdumax-1));
	      xzs[i][j]=MARK_SATULATE;
	    }
	  }
	}
      }
      fpixel+=naxes[0];
    }
    sx+=ihdu;
    sx2+=ihdu*ihdu;
    sx3+=ihdu*ihdu*ihdu;
    sx4+=ihdu*ihdu*ihdu*ihdu;
    sigma=(double)(ihdu+1)*sx2-(double)sx*sx;
    sigma2=(double)(ihdu+1)*sx3-(double)sx*sx2;
    sigma3=(double)(ihdu+1)*sx4-(double)sx2*sx2;
    //fprintf(stderr,"%d(%d) ",ihdu,ncr);
  }
  if(sigma==0) sigma=1;
  if(sigma2==0) sigma2=1;
  if(sigma3==0) sigma3=1;
  ncr=0;
  for(i=0;i<naxes[1];i++){
    for(j=0;j<naxes[0];j++){
      a=(((double)xzs[i][j]*ihdumax-(double)zs[i][j]*sx)*sigma2-(x2zs[i][j]*ihdumax-(double)zs[i][j]*sx2)*sigma)/(sigma2*sigma2-sigma3*sigma);
      b=((double)xzs[i][j]*ihdumax-(double)zs[i][j]*sx-a*sigma2)/sigma;
      //if(j==1025 && i==1022) fprintf(stderr,"%d %d %d %g %d %d %d %g %d %g %g",ihdumax,zs[i][j],xzs[i][j],x2zs[i][j],sx,sx2,sx3,sx4,sigma,sigma2,sigma3);
      if(xzs[i][j]!=MARK_BAD && xzs[i][j]!=MARK_SATULATE){
	zs[i][j]=-(int)(b*(ihdumax-1));
	xzs[i][j]=-(int)(a*(ihdumax-1)*(ihdumax-1));
      }
      if(cr[i][j]>ZSHIFT/2){
	cr[i][j]-=ZSHIFT;
	ncr++;
      }
      else cr[i][j]=0;
      //if(j==1025 && i==1022) fprintf(stderr," %d %d %d\n",zs[i][j],cr[i][j],xzs[i][j]);
    }
  }
  if(fits_close_file(fp, &status))
    printerror(status);
  fprintf(stderr,"OK\n");
  fprintf(stderr,"CR rejected : %d pixels\n",ncr);

  fprintf(stderr,"Write %s ... ",argv[2]);
  status=0;
  if (fits_create_file(&fp, argv[2], &status))
    printerror(status);
  if (fits_create_img(fp, bitpix, naxis, naxes, &status))
    printerror(status);
  fpixel=1;
  fprintf(stderr,"%d x %d format ",LINEPIX,LINEPIX);
  nelements=LINEPIX;
  for(i=0;i<LINEPIX;i++){
    if(fits_write_img(fp, TINT, fpixel, nelements, zs[i], &status))
      printerror(status);
    fpixel+=LINEPIX;
  }
  if (fits_close_file(fp, &status))
    printerror(status);
  fprintf(stderr,"OK\n");

  remove("check.fits");
  fprintf(stderr,"Write check.fits ... ");
  status=0;
  if (fits_create_file(&fp, "check.fits", &status))
    printerror(status);
  if (fits_create_img(fp, bitpix, naxis, naxes, &status))
    printerror(status);
  fpixel=1;
  fprintf(stderr,"%d x %d format ",LINEPIX,LINEPIX);
  nelements=LINEPIX;
  for(i=0;i<LINEPIX;i++){
    if(fits_write_img(fp, TINT, fpixel, nelements, xzs[i], &status))
      printerror(status);
    fpixel+=LINEPIX;
  }
  if (fits_close_file(fp, &status))
    printerror(status);
  fprintf(stderr,"OK\n");

  remove("cr.fits");
  fprintf(stderr,"Write cr.fits ... ");
  status=0;
  if (fits_create_file(&fp, "cr.fits", &status))
    printerror(status);
  if (fits_create_img(fp, bitpix, naxis, naxes, &status))
    printerror(status);
  fpixel=1;
  fprintf(stderr,"%d x %d format ",LINEPIX,LINEPIX);
  nelements=LINEPIX;
  for(i=0;i<LINEPIX;i++){
    if(fits_write_img(fp, TINT, fpixel, nelements, cr[i], &status))
      printerror(status);
    fpixel+=LINEPIX;
  }
  if (fits_close_file(fp, &status))
    printerror(status);
  fprintf(stderr,"OK\n");
  return(0);
}

void printerror( int status)
{
  /*****************************************************/
  /* Print out cfitsio error messages and exit program */
  /*****************************************************/

  char status_str[FLEN_STATUS], errmsg[FLEN_ERRMSG];

  if (status)
    fprintf(stderr, "\n*** Error occurred during program execution ***\n");

  fits_get_errstatus(status, status_str);   /* get the error description */
  fprintf(stderr, "\nstatus = %d: %s\n", status, status_str);

  /* get first message; null if stack is empty */
  if ( fits_read_errmsg(errmsg) )
    {
      fprintf(stderr, "\nError message stack:\n");
      fprintf(stderr, " %s\n", errmsg);

      while ( fits_read_errmsg(errmsg) )  /* get remaining messages */
	fprintf(stderr, " %s\n", errmsg);
    }

  exit( status );       /* terminate the program, returning error status */
}