//---------------------------------------------------------------------------
/*              *************************************************
                *                    PHOTZIP                    *
                *    A 16bit FITS Images Compression Utility    *
                *         Lior Shamir & Robert Nemiroff         *
                *       Michigan Technological University       *
                *                Houghton, MI 49931             *
                *************************************************
*/

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

/* this should be linked to the header file of cfitsio */
#include "../fits/fitsio.h"

//---------------------------------------------------------------------------

//unsigned short **data;    /* "data" should store the fits image's pixels values */
long width,height;          /* dimensions of the fits image */
unsigned short median=0;

typedef struct DATUM
{  unsigned short value;
   unsigned short background;
   long order;
}datum;

datum **data;

typedef struct PIXEL
{  unsigned short value;
   unsigned short x,y;
}pixel;

pixel *pixel_array=NULL;

unsigned short S[4028],B[4028]; /* use global memory for avoiding using the stack */
char header[200][256];  /* use this for the header */

int image_type,anynull,simple,bitpix,extend,keysnum;
long pcount,gcount;

unsigned short Find_Nth_Item(unsigned short *array, unsigned short num, unsigned short item)
{  unsigned short pivot,index,s_index=0,b_index=0,eq=0;
   if (num==1) return(array[0]);
   pivot=(unsigned short)(random()%num);
   for (index=0;index<num;index++)
   {  if (array[index]<array[pivot]) S[s_index++]=array[index];
      if (array[index]>=array[pivot]) B[b_index++]=array[index];
      if (array[index]==array[pivot]) eq++;
   }
   if (eq==num) return(array[0]);
   if (s_index==item) return(array[pivot]);
   if (s_index<=item)
   {  memmove(array,B,b_index*sizeof(unsigned short));  /* use global memory and not the stack */
      return(Find_Nth_Item(array,b_index,item-s_index));
   }
   else
   {  memmove(array,S,s_index*sizeof(unsigned short));  /* use global memory and not the stack */
      return(Find_Nth_Item(array,s_index,item));
   }
}


/* sort function - a function for using quick sort C function for finding
   the midean value
*/
int sort_function(const void *a, const void *b)
{
   return((*((unsigned short *)(a))>*((unsigned short *)(b)))*2-1);
}

int sort_function1(const void *a, const void *b)
{
  return((((pixel *)a)->value>((pixel *)b)->value)*2-1);
}

/*********************************************************************
/* ***************************** LoadFits ****************************
   purpose: load a FITS file.
   input: fits_file_path: the file name of the output file.
          sort: if 1 - the pixels in the frame are sorted for getting the
                midean value. If "t" parameter is specified, then the sorting
                is required, otherwise the pixel values are not sorted.
   returned value: 1 if succeeded. 0 if failed.
*/

int LoadFits(char *fits_file_path,int sort)
{  unsigned short val;
   fitsfile *fptr;       /* pointer to the FITS file; defined in fitsio.h */
   int status, ii, jj,x;
   unsigned short *short_array;
   long  nelements, exposure;
   int keycount;
   int naxis = 2;
   long naxes[2];

   status=0;    /* very important - the status must be initialized before calling to fitsio routines */
   fits_open_image(&fptr,fits_file_path,0,&status);
   if (status) return(0);
   fits_read_imghdr(fptr,MAX_COMPRESS_DIM,&simple,&bitpix, &naxis,naxes,&pcount,&gcount, &extend, &status);
   /* validate the attributes of the input file */
   if (status) return(0);
   /* read the header */
   fits_get_hdrspace(fptr,&keysnum,NULL,&status);
   if (status) return(0);
   for (keycount=1;keycount<=keysnum+1;keycount++)
   {  char keyword[256];
      int a;
      fits_read_record(fptr,keycount,keyword,&status);
      a=0;
      while (keyword[a]) 
      {  header[keycount][a]=keyword[a];
         a++;
      }
      header[keycount][a]='\0';
   }
   /* read the data */
   if (naxis!=2)
   {  printf("Input FITS doesn't have exactly two axis\n");
      return(0);
   }
   if (bitpix!=16)
   {  printf("Input FITS is not a 16 bit/pixel image\n");
      return(0);
   }

   /* allocate memory for the image */
   width=naxes[0];
   height=naxes[1];
   data=new datum *[width];
   if (!data) return(0); /* memory allocation failed */
   for (x=0;x<width;x++)
   {  data[x]=new datum[height];
      if (!data[x]) return (0);   /* memory allocation failed */
   }

   /* read the image */
   short_array=new unsigned short[width*height];
   pixel_array=new pixel[width*height];
   fits_read_img(fptr, TUSHORT, 1, width*height,NULL, short_array, NULL, &status);
   fits_close_file(fptr,&status);

   /* load the image data into the data structure */
   for (jj = 0;jj < height;jj++)
     for (ii = 0;ii < width;ii++)
     {  val=short_array[(jj)*width + ii];
        data[ii][height-jj-1].value=val;
        pixel_array[(jj)*width + ii].value=val;
        pixel_array[(jj)*width + ii].x=ii;
        pixel_array[(jj)*width + ii].y=height-jj-1;
     }

     /* find the median value */
     if (sort)
     {  qsort(pixel_array,width*height,sizeof(pixel),sort_function1);
        median=pixel_array[(int)(width*height/2)].value;
     }

   for (ii=0;ii<width*height;ii++)
     data[pixel_array[ii].x][pixel_array[ii].y].order=ii;
   delete short_array;
   return(1);
}
//---------------------------------------------------------------------------

/*********************************************************************
/* ***************************** SaveFits ****************************
   purpose: save data to a FITS image
   input: fits_file_path: the file name of the output file.
   returned value: 1 if succeeded. 0 if failed.
*/
int SaveFits(char *fits_file_path)
{  fitsfile *fptr;       /* pointer to the FITS file; defined in fitsio.h */
   int status, ii, jj;
   int image_type,anynull,simple,bitpix,extend;
   int keycount;
   unsigned short *short_array;
   long  nelements,naxis = 2;
   long naxes[2];
   unsigned short val;

   naxes[0]=width;   /* set the image's width  */
   naxes[1]=height;  /* set the image's height */
   status=0;
   remove(fits_file_path); /* make sure the file does not exist (or else it will fail creating the file) */
   fits_create_file(&fptr, fits_file_path, &status);  /* create a new file */
   if (status) return(0);
   fits_create_img(fptr, USHORT_IMG, naxis, naxes, &status); /* USHORT_IMG */
   if (status) return(0);
   /* write the header */
   for (keycount=1;keycount<=keysnum+1;keycount++)
   {  char keyword[256];
      int a=0;
      while (header[keycount][a])
      {  keyword[a]=header[keycount][a];
         a++;
      }
      keyword[a]='\0';
      fits_write_record(fptr,keyword,&status);
   }

   /* write the data */
   short_array=new unsigned short [height*width];
   /* save the image data into the fits file */
   for (jj = 0;jj < height;jj++)
     for (ii = 0;ii < width;ii++)
     {  val=data[ii][height-jj-1].value;
        short_array[(jj)*width + ii]=val;
     }
   nelements = naxes[0] * naxes[1];
   fits_write_img(fptr, TUSHORT, 1, nelements, short_array, &status);  /* TUSHORT */
   fits_close_file(fptr, &status);
   delete short_array;
   return(1);  /* successfuly saved */
}

/***********************************************************************
/* ************************* GetBackground *****************************
   purpose: Gets the background count for a specific pixel (usually a star pixel)
   input: x,y: The x,y coordinates of the pixel
          half_width_size: the half width of the background window
   comments: The background count is computed by finding the median value of the area
             determined by the half_width_size.
             this function also uses the "sort_function" function used by "LoadFits" for
             sorting
*/
unsigned short GetBackground(unsigned short x,unsigned short y,unsigned short half_width_size)
{  long xc,yc,bg,num=0;
   double sum=0;
   unsigned short pixel_counts[(half_width_size*2+1)*(half_width_size*2+1)];
   if (x<=half_width_size) x=half_width_size;
   if (x>=width-half_width_size) x=width-half_width_size-1;
   if (y<=half_width_size) y=half_width_size;
   if (y>=height-half_width_size) y=height-half_width_size-1;
   /* get the pixel counts around the star and put them in an array */
   for (yc=y-half_width_size;yc<=y+half_width_size;yc++)
     for (xc=x-half_width_size;xc<=x+half_width_size;xc++)
       pixel_counts[num++]=data[xc][yc].value;
   /* find the midean value */
   bg=Find_Nth_Item(pixel_counts,num,(long)(ceil(num/2)));
   return(bg);
}

/* *******************************************************************
/* ********************* GetPixelsBackground *************************
   purpose: compute a background value
void GetPixelsBackground(int half_width_size)
{  long x,y,xc,yc;
   unsigned short background;
   for (y=0;y<height;y=y+5)
   {  for (x=0;x<width;x=x+5)
      {  background=GetBackground(x+2,y+2,half_width_size);
         for (yc=y;yc<y+5 && yc<height;yc++)
           for (xc=x;xc<x+5 && xc<width;xc++)
             data[xc][yc].background=background;
      }
   }
}
*/
//---------------------------------------------------------------------------

/***********************************************************************
/* **************************** quantisize *****************************
   purpose: quantisize a pixel's value according to its sigma and max_diff
   input: value: the count value of the pixel.
          sigma: the sigma of the background of the pixel.
          max_diff: the maximum allowed difference between the pixel's value
                    and the pixel's value in the compressed/decompressed frame
                    (the b switch in the command line)
*/
unsigned short quantisize(unsigned short value,double sigma,double max_diff)
{  unsigned short quantisized_value;
   if (value<sigma || value<sigma*max_diff) return(0);

   quantisized_value=value & (~((unsigned short)(pow(2,ceil(log10(sigma*max_diff)/log10(2)))-1)));
   if (abs(value-quantisized_value)>pow(2,floor(log10(sigma*max_diff)/log10(2))))
     quantisized_value+=(unsigned short)pow(2,ceil(log10(sigma*max_diff)/log10(2)));
   return(quantisized_value);

}

/***********************************************************************
/* ************************** CompressFits *****************************
   purpose: quantisize a pixel's value according to its sigma and max_diff
   input: min_bright - the minimum sigmas over the backgrounf of a pixel to be
          quantisized (the d switch in the command line).
          half_width_size: the half width of the background window.
          (the s switch in the command line)
          max_diff - the maximum allowed difference between the pixel's value
          and the pixel's value in the compressed/decompressed frame
          (the b switch in the command line)
          threshold - a threshold for a pixel's value to be quantisized.
          (the t switch of the command line).
   returned value: 1 if succeeded, otherwise - 0.
*/
int CompressFits(double min_bright,double max_diff,int half_width_size,double threshold)
{  long x,y;
//   GetPixelsBackground(half_width_size);
   /* truncate all values which are under the midean value */
   for (y=0;y<height;y++)
   {  for (x=0;x<width;x++)
      {  double sigma;
         if ((x % 5) == 0 && (y % 5) == 0) data[x][y].background=GetBackground(x+2,y+2,half_width_size);
         else
         if ((x % 5) != 0 && (y % 5) == 0) data[x][y].background=data[x-1][y].background;
         else
         if ((x % 5) == 0 && (y % 5) != 0) data[x][y].background=data[x][y-1].background;
         else
         data[x][y].background=data[x-1][y].background;
         if (data[x][y].value<threshold*median)
           data[x][y].value=quantisize(data[x][y].value,sqrt(median),max_diff);
         else
         {  sigma=sqrt(data[x][y].background);
            if (data[x][y].value<data[x][y].background+min_bright*sigma)
              data[x][y].value=quantisize(data[x][y].value,sigma,max_diff);
         }
      }
   }
   return(1);  /* the function completed successfully */
}
//---------------------------------------------------------------------------

/* main */
int main(int argc, char* argv[])
{ unsigned short arg_index,last_switch=0,half_width_size=10,x;
  double min_bright=3,max_diff=1,threshold=0;
  for (arg_index=0;arg_index<=argc-1;arg_index++)
  {  /* check if -b was specified */
     if (strlen(argv[arg_index])>2 && argv[arg_index][0]=='-' && strchr(argv[arg_index],'b'))
     {  max_diff=atof(strchr(argv[arg_index],'b')+1);
        if (max_diff<0 || max_diff>20)
        {  printf("Invalid -b value '%s'.\n",strchr(argv[arg_index],'b')+1);
           return(0);
        }
        last_switch=arg_index;
     }
     /* check if -d was specified */
     if (strlen(argv[arg_index])>2 && argv[arg_index][0]=='-' && strchr(argv[arg_index],'d'))
     {  min_bright=atof(strchr(argv[arg_index],'d')+1);
        if (min_bright<0 || min_bright>20)
        {  printf("Invalid -d value '%s'.\n",strchr(argv[arg_index],'d')+1);
           return(0);
        }
        last_switch=arg_index;
     }
     /* check if -s was specified */
     if (strlen(argv[arg_index])>2 && argv[arg_index][0]=='-' && strchr(argv[arg_index],'s'))
     {  half_width_size=atoi(strchr(argv[arg_index],'s')+1);
        if (half_width_size<1)
        {  printf("Invalid -s value '%s'.\n",strchr(argv[arg_index],'s')+1);
           return(0);
        }
        last_switch=arg_index;
     }
     /* check if -t was specified */
     if (strlen(argv[arg_index])>2 && argv[arg_index][0]=='-' && strchr(argv[arg_index],'t'))
     {  threshold=atof(strchr(argv[arg_index],'t')+1);
        if (threshold<0)
        {  printf("Invalid -t value '%s'.\n",strchr(argv[arg_index],'t')+1);
           return(0);
        }
        last_switch=arg_index;
     }
  }
  if (argc<3 || last_switch<argc-3)
  { printf("usage: photzip [-dN|-bN|-sN|-tN] source target \n");
    return(0);
  }

  /* load the fits file */
  if (!LoadFits(argv[argc-2],threshold!=0))
  {  printf("Could not open file '%s'\n",argv[argc-2]);
     return(0);
  }

  /* run the photzip algorithm */
  CompressFits(min_bright,max_diff,half_width_size,threshold);

  /* save the output file to the disk */
  if (!SaveFits(argv[argc-1]))
  {  printf("Could not open file '%s'\n",argv[argc-1]);
     return(0);
  }

  /* free allocated memory */
  for (x=0;x<width;x++)
    delete data[x];
  delete data;
  if (pixel_array) delete pixel_array;
  return 0;
}
//---------------------------------------------------------------------------