/*
** $Id: fakeimage.c,v 1.4 2004-06-03 20:03:58 price Exp $
** Program to make some fake images.
**
** Paul A. Price
** 2003 September 3
** INDNJC.
**
** Copyright (C) 2003 Paul A. Price
**
** This program is free software; you can redistribute it and/or modify
** it under the terms of the GNU General Public License as published by
** the Free Software Foundation; either version 2 of the License, or
** (at your option) any later version.
**                                                                             
** This program is distributed in the hope that it will be useful,
** but WITHOUT ANY WARRANTY; without even the implied warranty of
** MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
** GNU General Public License for more details.
**                                                                             
** You should have received a copy of the GNU General Public License
** along with this program; if not, write to the Free Software
** Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
**
*/

#include <stdio.h>
#include <stdlib.h>
#include <getopt.h>
#include <fitsio.h>
#include <math.h>

/* For seeding the RNG */
#include <sys/timeb.h>

/* Memory allocation routines */
#include "memory.h"

/* GSL */
#include <gsl/gsl_sf_erf.h>		// For error functions


/* Conversion factor by which to multiply sigma to get FWHM.  Roughly 2.35 */
#define SIGMAtoFWHM (2.0*sqrt((double)2.0*(log((double)2.0))))

/* Errors - range of sigma to use */
#define MAXSIGMA 4.5
#define NUMSIGMA 100


/* Handy functions */
#define abs(a) (((a) > 0) ? (a) : (-a))
#define max(a,b) (((a) > (b)) ? (a) : (b))
#define min(a,b) (((a) > (b)) ? (b) : (a))

/* Function declarations */
void help(void);
void ckstatus(int status);
void addstar (float *data, int nx, int ny, int nxos, float x, float y, float flux, float seeing);
float linint(double x[], double y[], int num, double val);   /* Linear interpolation */
float gauss (float val);   /* Gaussian-type function */
float fringe (int x, int y, float periodx, float periody, float phasex, float phasey);

/* Global variables */
int verbose=0;   /* Verbosity switch */
long dims[2];   /* Dimensions of images */

int main(int argc, char **argv)
{
    /* Input variables */
    const char *noisefile, *outfile;   /* Input and output filenames */
    double background;   /* Sky background, ADU */
    double gain;   /* Detector gain, units e-/ADU */
    double readnoise;   /* Detector read noise, units e- */
    long overscan;   /* Number of overscan columns */
    double zero;   /* Bias/Zero level for image */
    double zeroslope;   /* Bias/Zero slope for image */
    int numcrs;   /* Number of CRs */
    int numstars;   /* Number of stars */
    double seeing;   /* Seeing FWHM (pixels) */
    char *starfile;   /* Name of file to use for stars */
    int debug;   /* Output debug information?  (ie CR position and flux) */
    int numbadpix;   /* Number of bad pixels */

    /* Calculation variables */
    fitsfile *fitsfp;   /* FITS files: input, output */
    int status = 0;   /* fitsio status */
    int i,x,y;   /* Counters */
    float *outdata;   /* Output pixels */
    float *noisedata;   /* Noise - Sigma to add to ideal image */
    int *mask;   /* Mask of CRs */
    float starx,stary;   /* Star centre */
    float flux;   /* (peak) flux of star */
    FILE *starfilefp;   /* File pointer for star file */
    int usestarfile;   /* Boolean indicating whether we get the star data from the file or not */
    int usenoisefile;   /* Boolean indicating whether we apply the noise or not */
    long noisedims[2];   /* Dimensions of noise map */
    double *gint, *gsigma;   /* Gaussian */
    char datasec[FLEN_KEYWORD],biassec[FLEN_KEYWORD]; // Data section and bias section for FITS
    struct timeb thetime;   /* For seeding the random number generator */
    float fringeamp, fringeperiod, fringephase;	// Fringe amplitude, period and phase

    /* Variables for getopt */
    int opt;   /* Option, from getopt */
    extern char *optarg;   /* Argument accompanying switch */
    extern int optind;   /* getopt variables */

    /* Seed the random number generator */
    ftime(&thetime);
    srand((unsigned int)thetime.millitm);

    /* Set defaults for variables */
    debug = 0;
    gain = 1.0;
    readnoise = 10.0;
    background = 10000.0;
    dims[0] = dims[1] = 1024;
    numcrs = 0;
    numbadpix = 0;
    seeing = 2.5;
    numstars = 0;
    usestarfile=0;
    noisefile=NULL;
    usenoisefile = 0;
    starfile=NULL;
    overscan = 0;
    zero = 0.;
    zeroslope=0.;
    fringeamp = 0.0;
    fringeperiod = (float)dims[0] / 20.0;

    while ((opt = getopt(argc, argv, "hvdkg:r:n:m:c:s:f:b:x:y:z:Z:o:p:F:")) != -1)
	switch (opt)
	{
	  case 'h':
	    help();
	    exit(0);
	  case 'v':
	    verbose = 1;
	    break;
	  case 'd':
	    debug = 1;
	    verbose = 0;
	    break;
	  case 'k':
	    usenoisefile = -1;
	    break;
	  case 'g':
	    if (sscanf(optarg, "%lf", &gain) != 1)
	    {
		help();
		exit(1);
	    }
	    break;
	  case 'r':
	    if (sscanf(optarg, "%lf", &readnoise) != 1)
	    {
		help();
		exit(1);
	    }
	    break;
	  case 'x':
	    if (sscanf(optarg, "%ld", &dims[0]) != 1)
	    {
		help();
		exit(1);
	    }
	    break;
	  case 'y':
	    if (sscanf(optarg, "%ld", &dims[1]) != 1)
	    {
		help();
		exit(1);
	    }
	    break;
	  case 'z':
	    if (sscanf(optarg, "%lf", &zero) != 1)
	    {
		help();
		exit(1);
	    }
	    break;
	  case 'Z':
	    if (sscanf(optarg, "%lf", &zeroslope) != 1)
	    {
		help();
		exit(1);
	    }
	    break;
	  case 'o':
	    if (sscanf(optarg, "%ld", &overscan) != 1)
	    {
		help();
		exit(1);
	    }
	    break;
	  case 'c':
	    if (sscanf(optarg, "%d", &numcrs) != 1)
	    {
		help();
		exit(1);
	    }
	    break;
	  case 's':
	    if (sscanf(optarg, "%d", &numstars) != 1)
	    {
		help();
		exit(1);
	    }
	    usestarfile=0;
	    break;
	  case 'p':
	    if (sscanf(optarg, "%d", &numbadpix) != 1)
	    {
		help();
		exit(1);
	    }
	    break;
	  case 'n':
	    noisefile = optarg;
	    usenoisefile = 1;
	    break;
	  case 'b':
	    if (sscanf(optarg, "%lf", &background) != 1)
	    {
		help();
		exit(1);
	    }
	    break;
	  case 'm':
	    starfile = optarg;
	    usestarfile=1;
	    break;
	  case 'f':
	    if (sscanf(optarg, "%lf", &seeing) != 1)
	    {
		help();
		exit(1);
	    }
	    break;
	  case 'F':
	    /*
	      Note: incrementing optind, so I can read more than
	      one parameter.
	    */
	    
	    if ((sscanf(argv[optind-1], "%f", &fringeamp) != 1) ||
		(sscanf(argv[optind++], "%f", &fringeperiod) != 1) ||
		(sscanf(argv[optind++], "%f", &fringephase) != 1 ))
	    {
		help();
		exit(EXIT_FAILURE);
	    }
	    break;
	  default:
	    help();
	    exit(1);
	}

    argc -= optind;
    argv += optind;

    if (argc != 1)
    {
	help();
	exit(1);
    }
    
    outfile = argv[0];


    /* Now, the program proper... */

    /* Do we read an input, or create an empty image? */
    if (usenoisefile == 1)
    {
	/* Open input file */
	if (verbose) printf ("Opening %s.... ", noisefile);
	if (fits_open_file(&fitsfp,noisefile,READONLY,&status)) ckstatus(status);

	/* Get dimensions */
	if (fits_get_img_size(fitsfp,2,noisedims,&status)) ckstatus(status);
	if (verbose) printf ("%ldx%ld\n",noisedims[0],noisedims[1]);

	if ((noisedims[0] != dims[0] + overscan) || (noisedims[1] != dims[1]))
	{
	    fprintf (stderr, "Noise map must have same dimensions as output image!\n");
	    fprintf (stderr, "Altering output image dimensions to %ldx%ld\n",noisedims[0],noisedims[1]);
	    dims[0] = noisedims[0];
	    dims[1] = noisedims[1];
	}

	/* Allocate storage */
	noisedata = farray((dims[0]+overscan)*dims[1]);

	/* Read the pixels */
	if (fits_read_img(fitsfp, TFLOAT, 1, (dims[0]+overscan)*dims[1], 0, noisedata, NULL, &status)) ckstatus(status);

	/* Close the file */
	if (fits_close_file(fitsfp, &status)) ckstatus(status);

    }
    else
    {
	/* Integrate gaussian so we can have a nice distribution */
	if (verbose) printf ("Creating noise map.....\n");
	gsigma = darray(NUMSIGMA);
	gint = darray(NUMSIGMA);
	for (i = 0; i < NUMSIGMA; i++)
	{
	    gsigma[i] = 2.0*(double)(MAXSIGMA*i)/(double)NUMSIGMA - MAXSIGMA;
	    gint[i] = gsl_sf_erf((double)gsigma[i]);
	}

	/* Allocate memory for output */
	noisedata = farray((dims[0]+overscan)*dims[1]);   /* Output Image */

	/* Create noise map */
	for (i=0;i<(dims[0]+overscan)*dims[1];i++) noisedata[i] = (float)linint(gint,gsigma,NUMSIGMA,(double)rand()/(double)RAND_MAX*(gint[NUMSIGMA-1]-gint[0])+gint[0]);
    }


    /* Create empty output file */
    if (verbose) printf ("Creating %ldx%ld output image, %s\n",dims[0]+overscan,dims[1],outfile);
    if (fits_create_file(&fitsfp, outfile, &status)) ckstatus(status);


    /* Initialise the output buffer to the background */
    outdata = farray((dims[0]+overscan)*dims[1]);
    mask = iarray((dims[0]+overscan)*dims[1]);
    if (verbose) printf("Initialising output buffer\n");
    for (y=0;y<dims[1];y++)
	for (x=0;x<dims[0]+overscan;x++)
	{
	    /* Bias */
	    outdata[(dims[0]+overscan)*y+x] = zero + zeroslope*(float)y;
	    /* CR pixel mask */
	    mask[(dims[0]+overscan)*y+x] = 0;
	    /* Add background to the true pixels */
	    if (x < dims[0]) outdata[(dims[0]+overscan)*y+x] += background;
	}
    
    /* Add stars */
    if (verbose) printf("Adding stars\n");

    if (usestarfile)
    {
	if ((starfilefp = fopen(starfile,"r")) == NULL)
	{
	    fprintf(stderr,"Can't open %s\n",starfile);
	    exit(1);
	}

	while (fscanf(starfilefp,"%f %f %f",&starx,&stary,&flux) == 3)
	    addstar(&outdata[0],dims[0],dims[1],overscan,starx,stary,flux,seeing);
	fclose(starfilefp);
    }
    else
    {
	/* Random distribution of Stars */
	for (i=1;i<=numstars;i++)
	{
	    starx = (rand()/(RAND_MAX+1.0)*dims[0]);
	    stary = (rand()/(RAND_MAX+1.0)*dims[1]);
	    flux = rand()/(RAND_MAX+1.0)*(32768.0-background);

	    addstar(&outdata[0],dims[0],dims[1],overscan,starx,stary,flux,seeing);
	    if (verbose) printf ("%f %f %f\n",starx,stary,flux);
	}
    }

    /* Add fringes */
    if (fringeamp != 0.0) {
	for (y = 0; y < dims[1]; y++)
	    for (x = 0; x < dims[0]; x++)
		outdata[(dims[0]+overscan)*y+x] += fringeamp * fringe(x,y,fringeperiod,fringeperiod,fringephase,fringephase);
    }
    
    /* Add Poisson noise */
    if (usenoisefile != -1)
    {
	if (verbose) printf ("Adding noise\n");
	for (y=0;y<dims[1];y++)
	    for (x=0;x<(dims[0]+overscan);x++)
		outdata[(dims[0]+overscan)*y+x] += noisedata[(dims[0]+overscan)*y+x]*sqrt((double)(gain*(outdata[(dims[0]+overscan)*y+x]-(zero + zeroslope*y)) + readnoise*readnoise))/gain;
    }

    /* Random distribution of bad pixels */
    if (verbose) printf ("Adding bad pixels\n");
    for (i = 1; i <= numbadpix; i++)
    {
	x = (int)(rand()/(RAND_MAX+1.0)*dims[0]);
	y = (int)(rand()/(RAND_MAX+1.0)*dims[1]);
	outdata[(dims[0]+overscan)*y+x] = 0.0;
    }

    /* Random distribution of CRs */
    if (verbose) printf ("Adding CRs\n");
    if (numcrs > dims[0]*dims[1]/2) fprintf (stderr,"Image too small (%ldx%ld) for %d CRs --- no CRs added.\n",dims[0],dims[1],numcrs);
    else
    {
	for (i=1;i<=numcrs;i++)
	{
	    x = (int)(rand()/(RAND_MAX+1.0)*dims[0]);
	    y = (int)(rand()/(RAND_MAX+1.0)*dims[1]);
	    if (mask[(dims[0]+overscan)*y+x]) i--;   /* Get another one */
	    else
	    {
		flux = rand()/(RAND_MAX+1.0)*(32768.0-background);
		outdata[(dims[0]+overscan)*y+x] += flux;
		mask[(dims[0]+overscan)*y+x] = 1;
		if (debug) printf ("%d %d %f\n",x,y,flux);
	    }
	}
    }

    /* Write output */
    if (verbose) printf ("Writing %ldx%ld output image...\n",dims[0]+overscan,dims[1]);
    dims[0] += overscan;
    if (fits_write_imghdr (fitsfp, -32, 2, dims, &status)) ckstatus(status);

    /* Write DATASEC, BIASSEC */
    (void)sprintf (datasec,"[%ld:%ld,%ld:%ld]",1L,dims[0]-overscan,1L,dims[1]);
    (void)sprintf (biassec,"[%ld:%ld,%ld:%ld]",dims[0]-overscan+1,dims[0],1L,dims[1]);
    if (fits_write_key(fitsfp, TSTRING, "DATASEC", datasec,"Data section",&status)) ckstatus(status);
    if (fits_write_key(fitsfp, TSTRING, "BIASSEC", biassec,"Bias section",&status)) ckstatus(status);

    /* Write the FITS file */
    if (fits_write_img(fitsfp, TFLOAT, 1, dims[0]*dims[1], outdata, &status)) ckstatus(status);

    /* Close files and clean up */
    if(fits_close_file(fitsfp, &status)) ckstatus(status);
    free ((float*)outdata);
    free ((int*)mask);

    if (verbose) printf ("Done!\n");

    exit(0);
}


/* Functions follow */

void help(void)
{
    fprintf(stderr,
	    "fakeimage [-hvd] [-x NX] [-y NY] [-g GAIN] [-r READNOISE] [-b BACKGROUND] [-s NUM_STARS] [-f FWHM] [-c NUM_CRS] [-p NUM_BADPIX] [-z ZERO_LEVEL] [-o OVERSCAN] [-n NOISE_MAP] OUTPUT_IMAGE\n"
	    "\n"
	    "\tOUTPUT_IMAGE:   Output file (FITS file)\n"
	    "\t-i INPUT_IMAGE: Use input image (FITS file)\n"
	    "\t-x NX; -y NY:   Size of the created image\n"
	    "\t-g GAIN:        Detector gain (e-/ADU)\n"
	    "\t-r READNOISE:   Detector read noise (e-)\n"
	    "\t-b BACKGROUND:  Sky background (ADU)\n"
	    "\t-s NUM_STARS:   Number of stars on the image\n"
	    "\t-m STARFILE:    Use file (x,y,flux) for stars\n"
	    "\t-f FWHM:        Seeing FWHM\n"
            "\t-c NUM_CRS:     Number of CRs on the image\n"
	    "\t-p NUM_BADPIX:  Number of bad pixels on image\n"
	    "\t-n NOISE_MAP:   Name of an image to use for the noise\n"
	    "\t-z ZERO_LEVEL:  Bias/Zero level to add to image\n"
	    "\t-Z ZERO_SLOPE:  Bias/Zero slope as a function of row\n"
	    "\t-o OVERSCAN:    Number of overscan columns to add\n"
	    "\t-k:             Klean --- no noise added\n"
	    "\t-h:             Help\n"
	    "\t-v:             Verbose\n"
	    "\t-d:             Debug (ie output CR flux, position)\n"
	    "\n"
	);
}

void ckstatus(int status)
{
    char msg[81];
    
    if (!status)
	return;
    while (fits_read_errmsg(msg))
	fprintf(stderr, "CFITSIO: %s\n", msg);
    exit(1);
}


/* Finds a value in an array --- note: zero-indexed */
int find(double value,			// Value to find
	 double *array,			// Array to find in
	 int n				// Number of values in array
    )
{
    int low = 0;			// Low-end bracket index
    int high = n - 1;			// High-end bracket index
    int index;				// Index to check

    /* Test the ends */
    if (array[high] < value) {
	return high - 1;
    }
    if (array[low] > value) {
	return low;
    }

    /* Bisect */
    do {
	index = (high - low) / 2 + low;
	if (array[index] > value) {
	    high = index;
	} else {
	    low = index;
	}
    } while (high > low + 1);

    return low;
}




float linint(double x[], double y[], int num, double val)
{
    /*
      Returns the value of the function y = f(x) at val,
      calculated by linear interpolation.
      There are num values in arrays x and y.
    */

    int index = find(val, x, num);	// Index of position in x
    double ans = y[index] + (y[index+1] - y[index]) * (val - x[index]) / (x[index+1] - x[index]);

    return (ans);
}

float gauss (float val)
{
    /*
      Returns a normal-shaped curve evaulated at val.
    */
    return (exp((double)(-0.5*val*val)));
}


void addstar (float *data, int nx, int ny, int nxos, float x, float y, float flux, float seeing)
/*
  Adds a Gaussian to the data

  data: The pixels
  nx, ny: Dimensions
  nxos: Size of overscan (don't put star here)
  x, y: Centre of the star
  flux: peak flux of the star
  seeing: FWHM of the star
*/
{
    int j,k;   /* Counters */
    int minj,maxj,mink,maxk;   /* Ranges */

    /* Get range for star */
    minj = max(0,(int) (x-seeing/SIGMAtoFWHM*MAXSIGMA));
    maxj = min(nx-1,(int) (x+seeing/SIGMAtoFWHM*MAXSIGMA));
    mink = max(0,(int) (y-seeing/SIGMAtoFWHM*MAXSIGMA));
    maxk = min(ny-1,(int) (y+seeing/SIGMAtoFWHM*MAXSIGMA));

    /* Put the star in. */
    for (j=minj;j<=maxj;j++)
	for (k=mink;k<=maxk;k++)
	    data[(nx+nxos)*k+j] += flux/sqrt(2.0*M_PI)/(seeing/SIGMAtoFWHM)*exp(-1.0 * ((x-j)*(x-j) + (y-k)*(y-k))/2.0/(seeing/SIGMAtoFWHM)/(seeing/SIGMAtoFWHM));
    
}

/* Return a "fringe" value for a pixel */
float fringe (int x, int y,		// Pixel coordinates
	      float periodx, float periody, // Period in x and y
	      float phasex, float phasey// Phase for x and y
    )
{
    float xval = 2.0*M_PI/periodx * (float)x + phasex;
    float yval = 2.0*M_PI/periody * (float)y + phasey;

    return cos(xval)*cos(yval);
}


