#! /usr/bin/env python

"""sepfilt.py - Separate Bayer filtered image into separate filters

    Usage examples:
            (a) sepfilt.py
            (b) sepfilt.py image1.fit
            (c) sepfilt.py -f image1.fit
            (d) sepfilt.py image*.fit
            (e) sepfilt.py image1.fit image2.fits img*.fit

    (a) operates on the default filename format of IMG???.FIT,
    (b,c,d,e) each operate only on the filename(s) specified,
    (c) also demonstrates the -f option, which allows existing files
        to be overwritten.

    Four files are produced for each input file, appended with f1, f2, f3, f4.
    These correspond to r, b, and two g filters, but the assignment is
    currently uncertain.
"""

from glob import glob
import pyfits
import sys, os.path
import getopt
import warnings

def sepfilt(fglob="IMG???.FIT", clobber=False):
    fnames = glob(fglob)
    if len(fnames) == 0:
        print "No matching files found"
    for fname in fnames:
        sepfiltone(fname, clobber=clobber)


def sepfiltone(fname, clobber=False):
    # Function for separating a Bayer filtered CCD image
    # into it's component colours.
    #
    # This currently performs the simplest method of simply assuming
    # each pixel samples the same point.
    #
    # The arrangement of the filter grid and hence assignment to R, G, B
    # still needs to be checked!
    #
    # A more sophisticated approach would do some interpolation and
    # preserve some of the original resolution.  The importance of
    # this depends on the size of the PSF in pixels.  So long as
    # information is not shared between the different filters, if the
    # PSF smooths the signal over a reasonable number of pixels and
    # integrated measurements should be accurate.
    #
    # The main difficulty with interpolating is that the interpolated
    # pixels will have different error properties to the original
    # pixels, however, I think that is probably not a particularly
    # troublesome issue compared with the simple technique which uses
    # systematically inaccurate positions for the samples.
    #
    print 'Processing %s'%fname
    warnings.resetwarnings()
    warnings.filterwarnings("ignore", ".*may have been truncated.*", UserWarning, append="True")
    p = pyfits.open(fname)
    d = p[0].data
    h = p[0].header
    #h = None
    fname = fname.lower() # IRAF doesn't like uppercase extensions
    for n, f in enumerate((d[::2,::2], d[1::2,::2], d[::2,1::2], d[1::2,1::2])):
        safewrite(fname, "f%i"%(n+1), f, h, clobber)


def safewrite(fname, postfix, data, header, clobber):
    fout = fname[::-1].replace(".", "."+postfix[::-1], 1)[::-1]
    if clobber or not os.path.exists(fout):
        pyfits.writeto(fout, data, header=header, clobber=clobber)
    else:
        print "Not overwriting %s (use -f option to force overwriting)"%fout


class Usage(Exception):
    def __init__(self, msg):
        self.msg = msg

def main(argv=None):
    if argv is None:
        argv = sys.argv
    try:
        try:
            opts, args = getopt.getopt(argv[1:], "hf", ["help", "force"])
        except getopt.error, msg:
            raise Usage(msg)
        clobber = False
        for o, a in opts:
            if o in ("-h", "--help"):
                print __doc__
                return 1
            if o in ("-f", "--force"):
                clobber = True
        if len(args) == 0:
            sepfilt(clobber=clobber)
        else:
            for a in args:
                sepfilt(a, clobber=clobber)
    except Usage, err:
        print >>sys.stderr, err.msg
        print >>sys.stderr, "for help use --help"
        return 2

if __name__ == "__main__":
    sys.exit(main())