# Python program to load a FITS image and display it

# import needed extensions
from numpy import *
import matplotlib.pyplot as plt # plotting package
import matplotlib.cm as cm # colormaps
from scipy import stats
import astropy.io.fits as pyfits

plt.ion()

# name of file containing the bias frame
biasfile = 'bias_18c_3.fit'
# read in the bias file
h = pyfits.open(biasfile)
# copy the image data into a numpy (numerical python) array
bias = h[0].data
h.close() # close the file

# list of dark files
darkfile = ['dark_18c_060s_1.fit', 'dark_18c_120s_1.fit', 'dark_18c_180s_1.fit', \
            'dark_18c_240s_1.fit', 'dark_18c_300s_1.fit']
# array of corresponding times
exposure = array([60.0, 120.0, 180.0, 240.0, 300.0])

# arrays to hold results
c_mean = 0.0*exposure
c_median = 0.0*exposure
c_rms = 0.0*exposure

# wait after each set of plots
plt.ion()
colmap = plt.get_cmap('gray') # load gray colormap

# process the files
for i in range(len(darkfile)):
  # read in the file
  h = pyfits.open(darkfile[i])
  print 'Read file ', darkfile[i]
  # copy the image data into a numpy (numerical python) array
  img = h[0].data
  h.close() # close the file
  # find the difference from the bias frames
  diff = img - bias
  # diff is a 2-d array, need to change to 1-d to make a histogram
  nx, ny = diff.shape # find the size of the array
  diffh = reshape(diff, nx*ny) # change the shape to be 1d
  # choose selection region
  f = 0.01 # ignore first and last fraction f of points
  s = sort(diffh)
  vmin = s[int(f*len(s))]
  vmax = s[int((1-f)*len(s))-1]
  print 'Excluding lowest and highest pixels, fraction = ', f
  print 'bounds = ', vmin, vmax
  # one figure with two plots
  plt.figure(i)
  plt.clf()
  #plt.title('Exposure (s) ='+str(exposure(i)))
  # plot the first image
  plt.subplot(121)
  plt.imshow(diff, cmap=colmap, vmin=vmin, vmax=vmax) # plot image using gray colorbar
  # plot the second image
  plt.subplot(122)
  # plot a histogram of the image values
  ht = plt.hist(diffh, bins=int(vmax-vmin), range=[vmin,vmax], histtype='stepfilled', color='k')
  nc = max(ht[0]) # maximum value in plotted histogram
  # select only values within ranges
  q = where((diffh >= vmin) & (diffh <= vmax))
  diffhcut = diffh[q]
  # find statistics
  c_mean[i] = mean(diffhcut)
  c_median[i] = median(diffhcut)
  c_rms[i] = std(diffhcut)
  print 'mean, median, rms = ', c_mean[i], c_median[i], c_rms[i]
  plt.plot([c_mean[i], c_mean[i]], [0, nc], '-g')
  plt.plot([c_median[i], c_median[i]], [0, nc], '--g')
  plt.show() # display the plots
  print
  #raw_input("Press Enter to continue...")

# plot median versus time
#m = c_mean
m = c_median
plt.figure(i+1)
plt.clf()
plt.xlabel('Time (s)')
plt.ylabel('Mean Counts')
#plt.ylabel('Median Counts')
plt.plot(exposure, m, 'd')
# do a linear fit to the data
slope, intercept, r_value, p_value, std_err = stats.linregress(exposure, m)
print 'slope = ', slope
print 'intercept = ', intercept
print 'correlation coefficient r =', r_value
# plot the fit
plt.plot(exposure, intercept + exposure*slope, '--r')
plt.show()