from pylab import *
from numpy import *
import pyfits as pyfits
import glob

def debias(data):
	 datadebiased = data[2:3072,53:2097] - median(data[:,0:48]) #the useful image and the bias area of the frame
	 datadebiased[datadebiased<0]=0 #in the very edge of one of the corners there is no illumination at all, so a bias substraction can make slightly negative values there, which we really do not want...
	 return datadebiased

minflatcounts = 5000 #the required minimum count level in a useful flatfield after bias substraction
maxflatcounts = 50000 #the permitted maximum count level in a useful flatfield after bias substraction

fig_size = [18,8]
params = {'backend':'ps' ,'axes.labelsize': 12,'text.fontsize': 8,'legend.fontsize': 8,'xtick.labelsize': 8,'ytick.labelsize': 8, 'text.usetex': False, 'figure.figsize': fig_size}
rcParams.update(params)
cmap=set_cmap('gnuplot')
clf()

print '\n','Processing of Flat images (creating Master Flats / reading in existing Master Flats)'

try:
	image = pyfits.open('ch1flatsmaster.fits')
	ch1flatsmaster = asarray(image[0].data)
	image = pyfits.open('ch2flatsmaster.fits')
	ch2flatsmaster = asarray(image[0].data)
	image = pyfits.open('ch3flatsmaster.fits')
	ch3flatsmaster = asarray(image[0].data)
	print 'Existing Master Flats were found, proceeding with them'
except:
	print 'No existing Master Flats were found, proceeding from scratch'
	skyflatfieldimages = glob.glob('*_FFS.fits')
	domeflatfieldimages = glob.glob('*_FFD.fits')
	flatfieldimages = append(skyflatfieldimages,domeflatfieldimages)
	flatfieldimages.sort()
	for flatframe in flatfieldimages:
		print '\n',flatframe
		image = pyfits.open(flatframe)
		ch1 = asarray(image[1].data)
		ch2 = asarray(image[2].data)
		ch3 = asarray(image[3].data)

		ch1 = debias(ch1)
		ch2 = debias(ch2)
		ch3 = debias(ch3)

		subplot(131)
		imshow(ch1,vmin=median(ch1)*0.9,vmax=median(ch1)*1.025)
		gca().invert_yaxis()
		xlabel('$r+i$')
		subplot(132)
		title(flatframe)
		imshow(ch2,vmin=median(ch2)*0.875,vmax=median(ch2)*1.05)
		gca().invert_yaxis()
		xlabel('$g$')
		subplot(133)
		imshow(ch3,vmin=median(ch3)*0.7,vmax=median(ch3)*1.4)
		gca().invert_yaxis()
		xlabel('$u$')
		figfilename = flatframe.split('.')[0]+'.png'
		subplots_adjust(top = 0.95, bottom = 0.05, left = 0.05, right = 0.95)
		savefig(figfilename,dpi=150)
		clf()

		print 'Median signal in Ch1: ',median(ch1)
		if median(ch1)>=minflatcounts and median(ch1)<=maxflatcounts:
			try:
				ch1flats = dstack([ch1flats,ch1/average(ch1)])
				print 'Flat added to Ch1 Flat collection',ch1flats.shape
			except:
				ch1flats = array(ch1/average(ch1))
				print 'Ch1 Flat collection created',ch1flats.shape
		print 'Median signal in Ch2: ',median(ch2)
		if median(ch2)>=minflatcounts and median(ch2)<=maxflatcounts:
			try:
				ch2flats = dstack([ch2flats,ch2/average(ch2)])
				print 'Flat added to Ch2 Flat collection',ch2flats.shape
			except:
				ch2flats = array(ch2/average(ch2))
				print 'Ch2 Flat collection created',ch2flats.shape
		print 'Median signal in Ch3: ',median(ch3)
		if median(ch3)>=minflatcounts and median(ch3)<=maxflatcounts:
			try:
				ch3flats = dstack([ch3flats,ch3/average(ch3)])
				print 'Flat added to Ch3 Flat collection',ch3flats.shape
			except:
				ch3flats = array(ch3/average(ch3))
				print 'Ch3 Flat collection created',ch3flats.shape

	print '\n','Calculating master flats'
	ch1flatsmaster = median(ch1flats,axis=2)
	ch2flatsmaster = median(ch2flats,axis=2)
	ch3flatsmaster = median(ch3flats,axis=2)

	subplot(131)
	imshow(ch1flatsmaster,vmin=median(ch1flatsmaster)*0.9,vmax=median(ch1flatsmaster)*1.025)
	gca().invert_yaxis()
	xlabel('$r+i$')
	subplot(132)
	title('Master Flats (Bias substracted)')
	imshow(ch2flatsmaster,vmin=median(ch2flatsmaster)*0.875,vmax=median(ch2flatsmaster)*1.05)
	gca().invert_yaxis()
	xlabel('$g$')
	subplot(133)
	imshow(ch3flatsmaster,vmin=median(ch3flatsmaster)*0.7,vmax=median(ch3flatsmaster)*1.4)
	gca().invert_yaxis()
	xlabel('$u$')
	figfilename = 'masterflats.png'
	subplots_adjust(top = 0.95, bottom = 0.05, left = 0.05, right = 0.95)
	savefig(figfilename,dpi=150)
	clf()

	pyfits.writeto('ch1flatsmaster.fits',ch1flatsmaster,clobber=True)
	pyfits.writeto('ch2flatsmaster.fits',ch2flatsmaster,clobber=True)
	pyfits.writeto('ch3flatsmaster.fits',ch3flatsmaster,clobber=True)

objectfiles = glob.glob('*_OBJ.fits')
objectfiles.sort()
for objfilename in objectfiles:

	try:
		testfilename = objfilename.split('.')[0]+'_ch1_proc.fits'
		image = pyfits.open(testfilename)
		testfilename = objfilename.split('.')[0]+'_ch2_proc.fits'
		image = pyfits.open(testfilename)
		testfilename = objfilename.split('.')[0]+'_ch3_proc.fits'
		image = pyfits.open(testfilename)
		print '\n','The science frame',objfilename,'is already reduced'
	except:
		print '\n','Reading in and initial processing (Bias substraction, trimming away prescan area) of the science frame',objfilename
		image = pyfits.open(objfilename)
		header = image[0].header
		objectname = image[0].header['OBJECT']
		objch1 = image[1].data
		objch2 = image[2].data
		objch3 = image[3].data

		objch1 = debias(objch1)
		objch2 = debias(objch2)
		objch3 = debias(objch3)

		subplot(131)
		imshow(np.arcsinh(objch1-median(objch1)),vmin=percentile(np.arcsinh(objch1-median(objch1)),1),vmax=percentile(np.arcsinh(objch1-median(objch1)),99.95))
		gca().invert_yaxis()
		xlabel('$r+i$')
		subplot(132)
		title(objfilename+' ['+objectname+'] - raw')
		imshow(np.arcsinh(objch2-median(objch2)),vmin=percentile(np.arcsinh(objch2-median(objch2)),1),vmax=percentile(np.arcsinh(objch2-median(objch2)),99.95))
		gca().invert_yaxis()
		xlabel('$g$')
		subplot(133)
		imshow(np.arcsinh(objch3-median(objch3)),vmin=percentile(np.arcsinh(objch3-median(objch3)),1),vmax=percentile(np.arcsinh(objch3-median(objch3)),99.95))
		gca().invert_yaxis()
		xlabel('$u$')
		figfilename = objfilename.split('.')[0]+'_raw.png'
		subplots_adjust(top = 0.95, bottom = 0.05, left = 0.05, right = 0.95)
		savefig(figfilename,dpi=150)
		clf()

		print '\n','Flatfielding science frame'
		objch1flatfielded = objch1/ch1flatsmaster
		objch2flatfielded = objch2/ch2flatsmaster
		objch3flatfielded = objch3/(ch3flatsmaster+1.75)
        #objch3flatfielded = objch3/ch3flatsmaster

		subplot(131)
		imshow(np.arcsinh(objch1flatfielded-median(objch1flatfielded)),vmin=percentile(np.arcsinh(objch1flatfielded-median(objch1flatfielded)),1),vmax=percentile(np.arcsinh(objch1flatfielded-median(objch1flatfielded)),99.95))
		gca().invert_yaxis()
		xlabel('$r+i$')
		subplot(132)
		title(objfilename+' ['+objectname+'] - reduced')
		imshow(np.arcsinh(objch2flatfielded-median(objch2flatfielded)),vmin=percentile(np.arcsinh(objch2flatfielded-median(objch2flatfielded)),1),vmax=percentile(np.arcsinh(objch2flatfielded-median(objch2flatfielded)),99.95))
		gca().invert_yaxis()
		xlabel('$g$')
		subplot(133)
		imshow(np.arcsinh(objch3flatfielded-median(objch3flatfielded)),vmin=percentile(np.arcsinh(objch3flatfielded-median(objch3flatfielded)),1),vmax=percentile(np.arcsinh(objch3flatfielded-median(objch3flatfielded)),99.95))
		gca().invert_yaxis()
		xlabel('$u$')
		figfilename = objfilename.split('.')[0]+'_proc.png'
		subplots_adjust(top = 0.95, bottom = 0.05, left = 0.05, right = 0.95)
		savefig(figfilename,dpi=150)
		clf()

		objch1filename_flatfielded = objfilename.split('.')[0]+'_ch1_proc.fits'
		objch2filename_flatfielded = objfilename.split('.')[0]+'_ch2_proc.fits'
		objch3filename_flatfielded = objfilename.split('.')[0]+'_ch3_proc.fits'

		pyfits.writeto(objch1filename_flatfielded,objch1flatfielded,clobber=True)
		pyfits.writeto(objch2filename_flatfielded,objch2flatfielded,clobber=True)
		pyfits.writeto(objch3filename_flatfielded,objch3flatfielded,clobber=True)