content: #!/usr/bin/env python
  
  import os
  import glob
  import numpy as np
  import fitsio
  import matplotlib.pyplot as plt
  from matplotlib import ticker
  from matplotlib.backends.backend_pdf import PdfPages
  from astropy.table import Table
  
  from bokpipe import *
  from bokpipe.bokoscan import _convertfitsreg
  
  def init_data_map(datadir,outdir,expTimes=None,files=None):
  	dataMap = {}
  	if not os.path.exists(outdir):
  		os.mkdir(outdir)
  	dataMap['outdir'] = outdir
  	if files is None:
  		dataMap['files'] = sorted(glob.glob(datadir+'*.fits') + 
  		                          glob.glob(datadir+'*.fits.gz') +
  		                          glob.glob(datadir+'*.fits.fz'))
  	else:
  		dataMap['files'] = files
  	dataMap['rawFiles'] = dataMap['files']
  	dataMap['oscan'] = bokio.FileNameMap(outdir)
  	dataMap['proc'] = bokio.FileNameMap(outdir,'_p')
  	dataMap['files'] = [ dataMap['oscan'](f) for f in dataMap['files'] ]
  	if expTimes is None:
  		dataMap['expTime'] = np.array([fitsio.read_header(f)['EXPTIME']
  		                                  for f in dataMap['files']])
  	else:
  		dataMap['expTime'] = expTimes
  	try:
  		# assume they are all the same
  		dataMap['dataSec'] = \
  		         _convertfitsreg(fitsio.read_header(
  		                             dataMap['files'][0],'IM4')['DATASEC'])
  	except IOError:
  		pass
  	return dataMap
  
  def process_data(dataMap,redo=True,withvar=True,oscanims=False,bias2d=False):
  	oscanSubtract = BokOverscanSubtract(output_map=dataMap['oscan'],
  	                                    overwrite=redo,
  		                                write_overscan_image=oscanims,
  		                    oscan_cols_file=dataMap['outdir']+'oscan_cols',
  		                    oscan_rows_file=dataMap['outdir']+'oscan_rows',
  		                                verbose=10)#method='median_value')
  	oscanSubtract.process_files(dataMap['rawFiles'])
  	if bias2d:
  		biasname = 'bias'
  		biasStack = bokproc.BokBiasStack(#reject=None,
  		                                 overwrite=redo,
  		                                 with_variance=withvar)
  		bias2dFile = os.path.join(dataMap['outdir'],biasname+'.fits')
  		biasStack.stack(dataMap['biasFiles'],bias2dFile)
  		#imProcess = bokproc.BokCCDProcess(bias2dFile,
  		#                                  output_map=dataMap['proc'])
  		#imProcess.process_files(flatFrames)
  
  def imstat(dataMap,outfn='stats'):
  	from astropy.stats import sigma_clip
  	from scipy.stats import mode,scoreatpercentile
  	array_stats = bokutil.array_stats
  	fnlen = len(os.path.basename(dataMap['files'][0]))
  	st = np.zeros(len(dataMap['flatSequence']),
  	              dtype=[('file','S%d'%fnlen),
  	                     ('expTime','f4'),
  	                     ('median','16f4'),
  	                     ('mean','16f4'),
  	                     ('mode','16f4'),
  	                     ('iqr25','16f4'),
  	                     ('iqr75','16f4'),
  	                     ('iqr10','16f4'),
  	                     ('iqr90','16f4')])
  	for _i,i in enumerate(dataMap['flatSequence']):
  		expTime = dataMap['expTime'][i]
  		fn = os.path.basename(dataMap['files'][i])
  		fits = fitsio.FITS(dataMap['files'][i])
  		print '%s %4.1f  ' % (fn,expTime),
  		st['file'][_i] = fn
  		st['expTime'][_i] = expTime
  		for j,extn in enumerate(['IM%d' % n for n in range(1,17)]):
  			modeVal,pix = array_stats(fits[extn].read()[dataMap['statsPix']],
  			                          method='mode',retArray=True)
  			st['mode'][_i,j] = modeVal
  			st['mean'][_i,j] = pix.mean()
  			st['median'][_i,j] = np.ma.median(pix)
  			st['iqr25'][_i,j] = scoreatpercentile(pix,25)
  			st['iqr75'][_i,j] = scoreatpercentile(pix,75)
  			st['iqr10'][_i,j] = scoreatpercentile(pix,10)
  			st['iqr90'][_i,j] = scoreatpercentile(pix,90)
  			print '%5d ' % (modeVal),
  		print
  	fitsio.write(outfn+'.fits',st,clobber=True)
  
  def scaled_histograms(dataMap,nims=None,outfn='pixhist'):
  	pdf = PdfPages(outfn+'.pdf')
  	for _i,i in enumerate(dataMap['flatSequence']):
  		if nims is not None and _i==nims:
  			break
  		expTime = dataMap['expTime'][i]
  		expScale = dataMap['refExpTime'] / expTime
  		print dataMap['files'][i]
  		fn = os.path.basename(dataMap['files'][i])
  		fits = fitsio.FITS(dataMap['files'][i])
  		fig = plt.figure(figsize=(8.0,10))
  		plt.subplots_adjust(0.08,0.08,0.92,0.92,0.3,0.35)
  		for j,extn in enumerate(['IM%d' % n for n in range(1,17)]):
  			ax = plt.subplot(8,2,j+1)
  			pix = fits[extn].read()[dataMap['statsPix']]
  			ax.hist(expScale*pix.flatten(),100,(0,40000),edgecolor='none')
  			ax.text(0.05,0.9,extn,va='top',size=9,transform=ax.transAxes)
  			ax.set_xlim(0,40000)
  			ax.xaxis.set_major_locator(ticker.MultipleLocator(10000))
  			ax.xaxis.set_minor_locator(ticker.MultipleLocator(2000))
  			ax.yaxis.set_major_locator(ticker.MultipleLocator(50000))
  		plt.figtext(0.5,0.99,fn+' exp=%.1f' % expTime,ha='center',va='top')
  		pdf.savefig(fig)
  		plt.close(fig)
  	pdf.close()
  
  def plot_sequence(dataMap,st,imNum,which='median'):
  	expScale = dataMap['refExpTime']/st['expTime']
  	seqno = 1 + np.arange(len(st))
  	ref = np.isclose(expScale,1.0)
  	j = imNum - 1
  	plt.figure(figsize=(8,6))
  	plt.subplots_adjust(0.11,0.08,0.96,0.95)
  	plt.errorbar(seqno[ref],expScale[ref]*st[which][ref,j],
  	                   [expScale[ref]*(st[which]-st['iqr10'])[ref,j],
  	                    expScale[ref]*(st['iqr90']-st[which])[ref,j]],
  	             fmt='bs-')
  	plt.errorbar(seqno[~ref],expScale[~ref]*st[which][~ref,j],
  	                   [expScale[~ref]*(st[which]-st['iqr10'])[~ref,j],
  	                    expScale[~ref]*(st['iqr90']-st[which])[~ref,j]],
  	             fmt='cs-')
  	#plt.scatter(seqno,expScale*st['mode'][:,j],marker='+',c='r')
  	#plt.scatter(seqno,expScale*st['mean'][:,j],marker='x',c='g')
  	plt.xlabel('sequence number')
  	plt.ylabel('counts scaled by exp time')
  	plt.title('IM%d'%imNum)
  	plt.xlim(0.5,len(st)+0.5)
  
  def fit_ref_exposures(dataMap,st,imNum,
                        which='median',method='spline',doplot=False):
  	from scipy.interpolate import UnivariateSpline
  	seqno = 1 + np.arange(len(st))
  	t = st['expTime']
  	ref = np.isclose(t,dataMap['refExpTime'])
  	j = imNum - 1
  	refCounts = st[which][ref,j][0]
  	if method=='linear':
  		_fit = np.polyfit(seqno[ref],refCounts/st[which][ref,j],1)
  		fit = lambda x: np.polyval(_fit,x)
  	elif method=='spline':
  		fit = UnivariateSpline(seqno[ref],refCounts/st[which][ref,j],
  		                       s=1e-5,k=3)
  	else:
  		raise ValueError
  	if doplot:
  		plt.figure()
  		plt.subplot(211)
  		plt.plot(seqno[ref],st[which][ref,j],'bs-')
  		plt.plot(seqno,refCounts/fit(seqno),c='r')
  		plt.subplot(212)
  		plt.plot(seqno[ref],(st[which][ref,j]-refCounts/fit(seqno[ref]))
  		                      /st[which][ref,j],'bs-')
  		plt.axhline(0,c='r')
  	return fit
  
  def plot_linearity_curves(dataMap,st,which='median',correct=True,isPTC=False,
                            refCor=None,fitmethod='spline',outfn='linearity',
  	                      onlyim=None):
  	seqno = 1 + np.arange(len(st))
  	t = st['expTime']
  	print seqno,t
  	refExpTime = dataMap['refExpTime']
  	ref = np.isclose(t,refExpTime)
  	refCorFit = None
  	ii = np.arange(len(st))
  	# only use the increasing sequence, not the reference exposures
  	ii = ii[~ref]
  	if isPTC:
  		# for PTCs skip every other image since they are done in pairs
  		ii = ii[::2]
  	# only fit to unsaturated frames
  	try:
  		firstsat = np.where(np.any(st[which][ii,:] > 55000,axis=1))[0][0]
  	except IndexError:
  		firstsat = -1
  	if onlyim is None:
  		pdf = PdfPages(outfn+'.pdf')
  	for imNum in range(1,17):
  		if onlyim is not None and imNum != onlyim:
  			continue
  		j = imNum - 1
  		# correct lamp variation
  		if correct:
  			if refCor is None:
  				fscl_fit = fit_ref_exposures(dataMap,st,imNum,which,
  				                             method=fitmethod)
  			else:
  				if refCorFit is None:
  					refCorFit = fit_ref_exposures(dataMap,st,imNum,which)
  				fscl_fit = refCorFit
  			fscl = fscl_fit(seqno)
  		else:
  			fscl = np.ones_like(seqno)
  		fit = np.polyfit(t[ii[:firstsat]],
  		                 fscl[ii[:firstsat]]*st[which][ii[:firstsat],j],1)
  		fitv = np.polyval(fit,t)
  		slope = fit[0] / (st[which][ref,j][0]/refExpTime)
  		#
  		pltindex = imNum % 4
  		if onlyim is None:
  			if pltindex == 1:
  				fig = plt.figure(figsize=(8,10))
  				plt.subplots_adjust(0.11,0.08,0.96,0.95,0.25,0.2)
  			ax = plt.subplot(4,2,2*(j%4)+1)
  		else:
  			fig = plt.figure(figsize=(6,2.5))
  			plt.subplots_adjust(0.11,0.23,0.99,0.98,0.35,0.2)
  			ax = plt.subplot(1,2,1)
  		plt.plot(t[ii],fscl[ii]*st[which][ii,j],'bs-')
  		plt.xlim(0.9*t.min(),t.max()+0.5)
  		plt.xscale('log')
  		plt.ylim(1e2,9e4)
  		plt.yscale('log')
  		plt.ylabel('counts [%s]' % which)
  		tt = np.logspace(-1,np.log10(1.3*t.max()),100)
  		plt.plot(tt,np.polyval(fit,tt),c='r')
  		plt.text(0.05,0.9,'IM%d'%imNum,va='top',transform=ax.transAxes)
  		plt.text(0.95,0.18,r'y = %.1f $\times$ t + %.1f' % tuple(fit),
  		         ha='right',va='top',size=9,transform=ax.transAxes)
  		plt.text(0.95,0.10,r'y = %.3f $\times$ counts + %.1f' % (slope,fit[1]),
  		         ha='right',va='top',size=9,transform=ax.transAxes)
  		if pltindex==0 or onlyim is not None:
  			plt.xlabel('exptime (s)')
  		#
  		if onlyim is None:
  			ax = plt.subplot(4,2,2*(j%4)+2)
  		else:
  			ax = plt.subplot(1,2,2)
  		plt.plot(t[ii],100*(fscl[ii]*st[which][ii,j]-fitv[ii])/fitv[ii],'bs-')
  		plt.axhline(0,c='r')
  		#ax.xaxis.set_major_locator(ticker.MultipleLocator(10))
  		#ax.xaxis.set_minor_locator(ticker.MultipleLocator(2))
  		ax.yaxis.set_major_locator(ticker.MultipleLocator(2))
  		ax.yaxis.set_minor_locator(ticker.MultipleLocator(0.5))
  		plt.ylim(-5,5)
  		plt.xlim(0.9*t.min(),t.max()+0.5)
  		plt.xscale('log')
  		if pltindex==0 or onlyim is not None:
  			plt.xlabel('exptime (s)')
  		plt.ylabel('residual \%')
  		if onlyim is None:
  			if pltindex == 0:
  				pdf.savefig(fig)
  				plt.close(fig)
  	if onlyim is None:
  		pdf.close()
  
  def rel_gain(dataMap,st,which='median',correct=True,fitmethod='spline',
               nskip=0):
  	seqno = 1 + np.arange(len(st))
  	t = st['expTime']
  	refExpTime = dataMap['refExpTime']
  	ref = np.isclose(t,refExpTime)
  	refCorFit = None
  	ii = np.arange(len(st))
  	ii = ii[~ref]
  	ii = ii[nskip:]
  	sky4 = st[which][ii,3]
  	fit_ii = ii[np.where((sky4>5000)&(sky4<25000))[0]]
  	plt.figure()
  	for imNum in range(1,17):
  		j = imNum - 1
  		# correct lamp variation
  		if correct:
  			if True: #refCor is None:
  				fscl_fit = fit_ref_exposures(dataMap,st,imNum,which,
  				                             method=fitmethod)
  			else:
  				if refCorFit is None:
  					refCorFit = fit_ref_exposures(dataMap,st,imNum,which)
  				fscl_fit = refCorFit
  			fscl = fscl_fit(seqno)
  		else:
  			fscl = np.ones_like(seqno)
  		fit = np.polyfit(t[fit_ii],fscl[fit_ii]*st[which][fit_ii,j],1)
  		fitv = np.polyval(fit,t)
  #		slope = fit[0] / (st[which][ref,j][0]/refExpTime)
  		xx = np.array(0,1.1*t.max())
  		plt.subplot(4,4,imNum)
  		if False:
  			plt.scatter(t[ii],fscl[ii]*st[which][ii,j])
  			plt.plot(xx,np.polyval(fit,xx),c='r')
  		else:
  			plt.scatter(t[ii],fscl[ii]*st[which][ii,j]/fitv[ii])
  			plt.axhline(1,c='r')
  		plt.ylim(0.7,1.3)
  		if True:
  			plt.xscale('log')
  		plt.xlim(0.9*t.min(),1.1*t.max())
  
  def get_first_saturated_frame(seq):
  	try:
  		firstsat = np.where(seq > 55000)[0][0]
  	except IndexError:
  		firstsat = -1
  	return firstsat
  
  def compare_oscan_levels(dataMap,st):
  	files = [ dataMap['files'][i] for i in dataMap['flatSequence'] ]
  	oscans = np.zeros((len(files),16))
  	for j in range(16):
  		oscans[:,j] = [ fitsio.read_header(f,'IM%d'%(j+1))['OSCANMED']
  	                      for f in files ]
  	seqno = 1 + np.arange(len(st))
  	plt.figure()
  	for j in range(8,16):
  		ax = plt.subplot(8,2,2*(j%8)+1)
  		i1 = get_first_saturated_frame(st['median'][:,j])
  		plt.scatter(st['median'][:i1,j],oscans[:i1,j],c='b')
  		plt.ylabel('IM%d'%(j+1))
  		ax = plt.subplot(8,2,2*(j%8)+2)
  		plt.scatter(seqno[:i1],oscans[:i1,j],c='b')
  
  def init_sep09bss_data_map():
  	datadir = os.environ.get('BASSDATA')+'/20150909/bss/20150908/'
  	exptimes = np.loadtxt(datadir+'../bss.20150909.log',usecols=(3,))
  	exptimes = exptimes[50:]
  	print exptimes
  	rdxdir = os.environ.get('GSCRATCH','tmp_sep')+'/bss_sep09/'
  	if not os.path.exists(rdxdir):
  		os.makedirs(rdxdir)
  	dataMap = init_data_map(datadir,rdxdir,
  	                        expTimes=exptimes,files=None)
  	dataMap['rawFiles'] = dataMap['rawFiles'][50:]
  	dataMap['files'] = dataMap['files'][50:]
  	dataMap['biasFiles'] = dataMap['files'][-5:]
  	#dataMap['flatSequence'] = range(50,68)
  	dataMap['flatSequence'] = range(18)
  	dataMap['statsPix'] = bokutil.stats_region('amp_corner_ccdcenter_small')
  	dataMap['refExpTime'] = 40.0
  	return dataMap
  
  def init_sep29ptc_data_map():
  	dataMap = init_data_map(
  	      "/home/ian/dev/rapala/bokpipe/scratch/sep29ptcs/ptc/",'sep29ptcs/')
  	dataMap['biasFiles'] = [dataMap['files'][0],]
  	dataMap['flatSequence'] = range(1,len(dataMap['files']))
  	dataMap['statsPix'] = np.s_[20:-20,100:-100]
  	dataMap['refExpTime'] = 10.0
  	return dataMap
  
  def init_oct02ptc_data_map():
  	dataMap = init_data_map(os.environ.get('GSCRATCH')+'/02oct15/ptc/',
  	                        os.environ.get('GSCRATCH')+'/02oct15/ptc_proc/')
  	dataMap['biasFiles'] = [dataMap['files'][0],]
  	dataMap['flatSequence'] = range(1,len(dataMap['files']))
  	dataMap['statsPix'] = bokutil.stats_region('amp_corner_ccdcenter_small')
  	dataMap['refExpTime'] = 10.0
  	return dataMap
  
  def init_oct20_data_map():
  	datadir = os.environ.get('BASSDATA')+'/20151020/'
  	exptimes = np.loadtxt(datadir+'images.log',usecols=(6,))
  	nuse = 53
  	exptimes = exptimes[:nuse]
  	print exptimes
  	dataMap = init_data_map(datadir,'tmp_oct20',expTimes=exptimes)
  	dataMap['rawFiles'] = dataMap['rawFiles'][:nuse]
  	dataMap['files'] = [ dataMap['oscan'](f) 
  	                       for f in dataMap['files'][:nuse] ]
  	dataMap['biasFiles'] = dataMap['files'][:20]
  	dataMap['flatSequence'] = range(20,nuse)
  	dataMap['statsPix'] = bokutil.stats_region('amp_corner_ccdcenter_small')
  	dataMap['refExpTime'] = 3.0
  	return dataMap
  
  def init_nov11g_data_map():
  	datadir = os.environ.get('BASSDATA')+'/Nov2015/'
  	log = Table.read(datadir+'bassLog_Nov2015.fits')
  	exptimes = log['expTime'][111:150]
  	files = [ datadir+f['utDir']+'/'+f['fileName']+'.fits'
  	              for f in log[111:150] ]
  	dataMap = init_data_map(datadir,'tmp_nov11g',
  	                        expTimes=exptimes,files=files)
  	dataMap['biasFiles'] = dataMap['files'][-10:]
  	dataMap['flatSequence'] = np.arange(len(dataMap['files'])-10)
  	dataMap['statsPix'] = bokutil.stats_region('amp_corner_ccdcenter_small')
  	dataMap['refExpTime'] = 3.0
  	return dataMap
  
  def init_nov14_data_map(filt):
  	datadir = os.environ.get('BASSDATA')+'/Nov2015/'
  	log = Table.read(datadir+'bassLog_Nov2015.fits')
  	if filt=='g':
  		frames = np.r_[np.s_[297:345],np.s_[247:257]]
  	else:
  		frames = np.r_[np.s_[345:393],np.s_[247:257]]
  	exptimes = log['expTime'][frames]
  	files = [ datadir+f['utDir']+'/'+f['fileName']+'.fits'
  	              for f in log[frames] ]
  	dataMap = init_data_map(datadir,'tmp_nov14'+filt,
  	                        expTimes=exptimes,files=files)
  	dataMap['biasFiles'] = dataMap['files'][-10:]
  	dataMap['flatSequence'] = np.arange(len(dataMap['files'])-10)
  	dataMap['statsPix'] = bokutil.stats_region('amp_corner_ccdcenter_small')
  	dataMap['refExpTime'] = {'Ha':10.0,'g':3.0}[filt]
  	return dataMap
  
  def init_jan3_data_map(filt):
  	datadir = os.environ.get('BASSDATA')
  	log = Table.read('basslogs/log_ut20160103.fits')
  	if filt=='g':
  		frames = np.r_[np.s_[57:105],np.s_[160:170]]
  	else:
  		frames = np.r_[np.s_[105:160],np.s_[160:170]]
  	exptimes = log['expTime'][frames]
  	files = [ datadir+'/'+f['utDir'].strip()+'/'+f['fileName'].strip()+'.fits'
  	              for f in log[frames] ]
  	dataMap = init_data_map(datadir,'tmp_jan3'+filt,
  	                        expTimes=exptimes,files=files)
  	dataMap['biasFiles'] = dataMap['files'][-10:]
  	dataMap['flatSequence'] = np.arange(len(dataMap['files'])-10)
  	dataMap['statsPix'] = bokutil.stats_region('amp_corner_ccdcenter_small')
  	dataMap['refExpTime'] = {'Ha':10.0,'g':3.0}[filt]
  	return dataMap
  
  def init_data_map_fromfile(filename,outdir='tmp',nersc=True):
  	datadir = os.environ.get('BASSDATA')
  	if nersc:
  		datadir = os.path.join(datadir,'BOK_Raw')
  	log = np.loadtxt(filename,dtype=[('frameNum','i4'),('utDir','S8'),
  	                                 ('fileName','S35'),
  	                                 ('imType','S10'),('filter','S8'),
  	                                 ('expTime','f4')],skiprows=1)
  	exptimes = log['expTime']
  	files = [ datadir+'/'+f['utDir'].strip()+'/'+f['fileName'].strip()+'.fits'
  	              for f in log ]
  	if nersc:
  		files = [ f+'.fz' for f in files ]
  	dataMap = init_data_map(datadir,outdir,
  	                        expTimes=exptimes,files=files)
  	dataMap['biasFiles'] = np.array(dataMap['files'])[log['imType']=='zero']
  	dataMap['flatSequence'] = np.where(log['imType']=='flat')[0]
  	dataMap['statsPix'] = bokutil.stats_region('amp_corner_ccdcenter_small')
  	# assume it starts with reference
  	dataMap['refExpTime'] = exptimes[dataMap['flatSequence'][0]]
  	return dataMap
  
  if __name__=='__main__':
  	import sys
  	dataset = sys.argv[1] 
  	if dataset == 'sep09bss':
  		dataMap = init_sep09bss_data_map()
  	elif dataset == 'oct02':
  		dataMap = init_oct02ptc_data_map()
  	elif dataset == 'oct20':
  		dataMap = init_oct20_data_map()
  	elif dataset == 'nov11g':
  		dataMap = init_nov11g_data_map()
  	elif dataset == 'nov14g':
  		dataMap = init_nov14_data_map('g')
  	elif dataset == 'nov14Ha':
  		dataMap = init_nov14_data_map('Ha')
  	elif dataset == 'jan3g':
  		dataMap = init_jan3_data_map('g')
  	elif dataset == 'jan3Ha':
  		dataMap = init_jan3_data_map('Ha')
  	else:
  		dataMap = init_data_map_fromfile(sys.argv[2],dataset)
  	print 'processing ',dataset
  	if not os.path.exists('stats_'+dataset+'.fits'):
  		process_data(dataMap,bias2d=True)
  		imstat(dataMap,outfn='stats_'+dataset)
  	st = fitsio.read('stats_'+dataset+'.fits')
  	plot_linearity_curves(dataMap,st,outfn='linearity_'+dataset)
  	if True:
  		plot_linearity_curves(dataMap,st,outfn='linearity_'+dataset,
  		                      onlyim=4)
  		plt.savefig('linearity_IM4_%s.png'%dataset)
  		plot_sequence(dataMap,st,4)
  		plt.savefig('linsequence_IM4_%s.png'%dataset)
  
  