# Procedure to fix known bad columns in CCD2 images.  2016 Oct 2 E. Gates

import numpy as np
from astropy.io import fits

#read in lists of data
datain=[line.rstrip('\n') for line in open('datafiles_ff.list')]
dataout=[line.rstrip('\n') for line in open('datafiles_fix.list')]

n=len(datain)
# size of box for area around bad pixel to be averaged
s=2

#read in one image to get image size for bad pixel mask
data,header=fits.getdata(datain[0],header=True)

#make bad pixel mask
mask=np.ma.make_mask(data,copy=True,shrink=True,dtype=np.bool)
mask[:,:]=False
mask[:,255:257]=True
mask[:,783:785]=True
mask[:,1001:1003]=True

#loop for all the data bad pixel correction 
for k in range(0,n):
  data,header=fits.getdata(datain[k],header=True)
  mdata=np.ma.masked_array(data,mask=mask,fill_value=np.nan)
  dataFixed=data.copy()
  for i in range(0,mdata.shape[0]):
    for j in range(0,mdata.shape[1]):
      if np.math.isnan(mdata[i,j]):
        x1=i-s
        x2=i+s+1
        y1=j-s
        y2=j+s+1
        if x1<0:
          x1=0
        if x2>mdata.shape[0]:
          x2=mdata.shape[0]
        if y1<0:
          y1=0
        if y2>mdata.shape[1]:
          y2=mdata.shape[1]
        dataFixed[i,j]=np.mean(mdata[x1:x2,y1:y2])
  header['HISTORY']='Bad columns replaced'
  fits.writeto(dataout[k],dataFixed,header,overwrite=True)