""" Code to filter image files

Context : SRP
Module  : SRPImageFilter.py
Author  : Stefano Covino
Date    : 16/05/2014
E-mail  : stefano.covino@brera.inaf.it
URL:    : http://www.merate.mi.astro.it/utenti/covino
Purpose : Manage filtering of FITS frames.

Usage   : SRPImageFilter [-h] -i arg1 [-m arg2] [-n] [-v]
            -i file of list if files to be processed.
            -m size of median filter.
            -n NAN filtering.
            The output files are produced applying a median filter of given size.

History : (02/08/2011) First version.
        : (04/09/2011) Better cosmetics.
        : (25/03/2014) Deal with non standard FITS headers.
        : (16/05/2014) Fltering NAN out.
"""



import os, os.path, string, warnings
from optparse import OptionParser
import SRP.SRPConstants as SRPConstants
import SRP.SRPFiles as SRPFiles
import numpy, pyfits
from SRP.SRPFits.AddHeaderComment import AddHeaderComment
from SRP.SRPFits.IsFits import IsFits
import scipy.ndimage.filters as SNF



parser = OptionParser(usage="usage: %prog [-h] -i arg1 [-m arg2] [-n] [-v]", version="%prog 1.1.0")
parser.add_option("-i", "--inputfilelist", action="store", nargs=1, type="string", dest="fitsfilelist", help="Input FITS file list")
parser.add_option("-m", "--median", action="store", nargs=1, type="int", dest="mediansize", help="Size of the median filter")
parser.add_option("-n", "--nan", action="store_true", dest="nan", help="NAN data filtered out")
parser.add_option("-v", "--verbose", action="store_true", dest="verbose", help="Fully describe operations")
(options, args) = parser.parse_args()

if options.fitsfilelist and (options.mediansize or options.nan):
    #
    if options.mediansize != None and options.mediansize <= 0:
        parser.error("Median filter size must be positive.")
    elif options.mediansize > 0:
        if options.verbose:
            print "Median filter size: %d" % options.mediansize
    #
    if options.verbose:
        if options.nan:
            print "NAN filtering active."
        else:
            print "NAN filtering not active."
    #
    if os.path.isfile(options.fitsfilelist) and not IsFits(options.fitsfilelist):
        f = SRPFiles.SRPFile(SRPConstants.SRPLocalDir,options.fitsfilelist,SRPFiles.ReadMode)
        f.SRPOpenFile()
        if options.verbose:
            print "Input FITS file list is: %s." % options.fitsfilelist
        froot,fext = os.path.splitext(options.fitsfilelist)
        o = SRPFiles.SRPFile(SRPConstants.SRPLocalDir,froot+SRPConstants.SRPScienceFile+fext,SRPFiles.WriteMode)
        o.SRPOpenFile()
        while True:
            dt = f.SRPReadFile()
            if dt != '':
                file = string.split(string.strip(dt))[0]
                if not os.path.isfile(file):
                    parser.error("Input FITS file %s not found" % file)
                if options.verbose:
                    print "FITS file selected: %s" % file
                if options.verbose:
                    print "Loading frame..."
                cb = pyfits.open(file)                
                cbdata = cb[0].data
                cbhead = cb[0].header
                cbshape = cb[0].data.shape
                #
                if options.verbose:
                    print "Filtering..."
                if options.mediansize:
                    if options.verbose:
                        print "Median filter..."
                    cbf = SNF.median_filter(cbdata,options.mediansize)
                else:
                    cbf = cbdata
                #
                if options.nan:
                    if options.verbose:
                        print "NAN filtering..."
                    cbfn = numpy.nan_to_num(cbf)
                else:
                    cbfn = cbf
                #
                root,ext = os.path.splitext(os.path.basename(file))
                if options.verbose:
                    print "Saving file: %s" % root+SRPConstants.SRPImaFltFile
                nflt = pyfits.PrimaryHDU(cbfn,cbhead)
                nfltlist = pyfits.HDUList([nflt])
                warnings.resetwarnings()
                warnings.filterwarnings('ignore', category=UserWarning, append=True)
                if options.verbose:
                    nfltlist.writeto(root+SRPConstants.SRPImaFltFITS,clobber=True,output_verify='warn')
                else:
                    nfltlist.writeto(root+SRPConstants.SRPImaFltFITS,clobber=True,output_verify='ignore')
                warnings.resetwarnings() 
                warnings.filterwarnings('always', category=UserWarning, append=True)   
                if options.mediansize:
                    AddHeaderComment(root+SRPConstants.SRPImaFltFITS,(("SRPComment: Median filtered frame (size %d)." % options.mediansize),))                
                if options.nan:
                    AddHeaderComment(root+SRPConstants.SRPImaFltFITS,(("SRPComment: NAN filtered out."),))                
                oentr = root+SRPConstants.SRPScienceFITS+SRPConstants.SRPTab+string.join(string.split(string.strip(dt))[1:])
                o.SRPWriteFile(oentr+os.linesep)
            else:
                break
        f.SRPCloseFile()
        o.SRPCloseFile()
    elif IsFits(options.fitsfilelist):
        if options.verbose:
            print "FITS file selected: %s" % options.fitsfilelist
        if options.verbose:
            print "Loading frame..."
        cb = pyfits.open(options.fitsfilelist)
        cbdata = cb[0].data
        cbhead = cb[0].header
        cbshape = cb[0].data.shape
        #
        if options.verbose:
            print "Filtering..."
        if options.mediansize:
            if options.verbose:
                print "Median filter..."
            cbf = SNF.median_filter(cbdata,options.mediansize)
        else:
            cbf = cbdata
        #
        if options.nan:
            if options.verbose:
                print "NAN filtering..."
            cbfn = numpy.nan_to_num(cbf)
        else:
            cbfn = cbf
        #
        root,ext = os.path.splitext(os.path.basename(options.fitsfilelist))
        if options.verbose:
            print "Saving file: %s" % root+SRPConstants.SRPImaFltFITS
        else:
            print root+SRPConstants.SRPImaFltFITS
        nflt = pyfits.PrimaryHDU(cbfn,cbhead)
        nfltlist = pyfits.HDUList([nflt])
        warnings.resetwarnings()
        warnings.filterwarnings('ignore', category=UserWarning, append=True)
        if options.verbose:
            nfltlist.writeto(root+SRPConstants.SRPImaFltFITS,clobber=True,output_verify='warn')
        else:
            nfltlist.writeto(root+SRPConstants.SRPImaFltFITS,clobber=True,output_verify='ignore')
        warnings.resetwarnings()
        warnings.filterwarnings('always', category=UserWarning, append=True)   
        if options.mediansize:
            AddHeaderComment(root+SRPConstants.SRPImaFltFITS,(("SRPComment: Median filtered frame (size %d)." % options.mediansize),))                
        if options.nan:
            AddHeaderComment(root+SRPConstants.SRPImaFltFITS,(("SRPComment: NAN filtered out."),))                
    else:
        parser.error("Input FITS file list %s not found" % options.fitsfilelist)
else:
    parser.print_help()
