#!/bin/env python

"""Write an IRAF script to do fully-automated reduction of a block of Gemini files"""

# to do:  recognize sequence interruptions (both to flag first resumed and to block separate sky and source epochs)

from pyfits import getheader
import sys,os,glob
from math import ceil

def unique(inlist):
    lis = inlist[:] #make a copy
    lis.sort()
    llen = len(lis)
    i = 0
    while i < llen-1:
        if lis[i+1] == lis[i]: 
            del lis[i+1]
            llen = llen - 1
        else:
            i = i + 1
    return lis

def prepend(stringlist, prefix):
    newlist = stringlist[:]
    for i in range(len(newlist)):
        newlist[i] = prefix + newlist[i]
    return newlist

def niriscript(files, supersky):
    
    skyblock = 15
    excludefirst = 1

    darkfiles = []
    darkexp = []
    darkexpinfo = []

    flatfiles = []
    flatfilters = []
    flatexpinfo = []
    flatshutter = []

    objfiles = []
    objfilters = []
    objnames = []
    objexpinfo = []

    for f in files:
        if f == '': continue
        
        h = getheader(f)
        hext = getheader(f,1)
        s = ''
        file = f[f.rfind('/')+1:]
        file = file[:file.rfind('.fits')]

        try:
            object = h['object'].strip()
            camera = h['camera'].strip()
            if camera != 'f6': 
               print f + ' is not f6 camera'
               continue
            sizestr = str(hext['naxis1'])+'x'+str(hext['naxis2'])
            if sizestr != '1024x1024':
               print f + ' is windowed'
               continue
            exp = "%.2f" % h['exptime']
            coadd = "%3i" % h['coadds']
            expinfo = coadd.ljust(2) + 'x' + exp.ljust(4)[0:4]
            filt =  h['filter1']
            if filt == 'Kprime_G0206': filt = "Kp"
            if filt == 'Kshort_G0205': filt = "Ks"
            if filt == 'K_G0204': filt = "K"
            if filt == 'H_G0203': filt = "H"
            if filt == 'J_G0202': filt = "J"

            try:
                shutter = h['gcalshut'].strip()
            except:
                shutter = ''
            if shutter == 'CLOSED': shutter = 'closed'
            if shutter == 'OPEN': shutter = 'open'
 
        except:
            #if this is not a standard file, ignore it
            print f + ' appears non-standard'
            continue

        if object == 'Dark':
            darkfiles.append(f)
            darkexp.append(exp)
            darkexpinfo.append(expinfo)
        elif object == 'GCALflat':
            flatfiles.append(f)
            flatfilters.append(filt)
            flatexpinfo.append(expinfo)
            flatshutter.append(shutter)
        else:
            objfiles.append(f)
            objfilters.append(filt)
            objnames.append(object)
            objexpinfo.append(expinfo)
        

        #print f, filt, object, shutter, coadd.ljust(2) + 'x' + exp.ljust(4)[0:4]

    print
    print
   
    darkproclist = []
    mindarkexp = min(darkexp)
    for i in range(len(darkfiles)):
        if darkexp[i] == mindarkexp:
           darkproclist.append(darkfiles[i])

    print 'Dark'
    print darkproclist

    filterlist = unique(objfilters)
    nfilter = len(filterlist)

    onlists = []
    offlists = []

    print
    print 'Flats:'
    for filt in filterlist:
        filtonlist = []
        filtofflist = []

        for i in range(len(flatfiles)):
            if flatfilters[i] == filt and flatshutter[i] == 'open': filtonlist.append(flatfiles[i])
            if flatfilters[i] == filt and flatshutter[i] == 'closed': filtofflist.append(flatfiles[i])

        onlists.append(filtonlist)
        offlists.append(filtofflist)

    for n in range(nfilter):
       print filterlist[n]
       print onlists[n]
       print offlists[n]


    objlist = unique(objnames)
    nobject = len(objlist)

    stacklists = []  
    stackobjs  = []
    stackfilts = []

    for obj in objlist:
        for filt in filterlist:
            inlist = []

            for i in range(len(objfiles)):
                if objfilters[i] == filt and objnames[i] == obj:
                    inlist.append(objfiles[i])
            if len(inlist) > 0:
                stacklists.append(inlist)
                stackobjs.append(obj)
                stackfilts.append(filt)

    nstack = len(stackobjs)

    if excludefirst > 0:
        for i in range(len(stackobjs)):
            stacklists[i] = stacklists[i][excludefirst:]

    print
    print 'Science:'
    for n in range(nstack):
        print stackobjs[n], stackfilts[n]
        print stacklists[n]

    scriptfilename = 'doniriredux.cl'
    scriptfile = open(scriptfilename,'w')

    # list darks.

    darklistfilename = 'dark.list'
    darklistfile = open(darklistfilename, 'w')
    darklistfile.writelines(l + '\n' for l in darkproclist)
    darklistfile.close()

    scriptfile.write('nprepare ("@' + darklistfilename + '")\n')

    ndarklistfilename = 'ndark.list'
    ndarklistfile = open(ndarklistfilename, 'w')
    ndarklistfile.writelines(l + '\n' for l in prepend(darkproclist,'n'))
    ndarklistfile.close()
    print

    # flat-fielding / creating bad pixel mask

    for j in range(nfilter):
        filt = filterlist[j]

        # nprepare
 
        callistfilename = 'allcal'+filt+'.list'
        callistfile = open(callistfilename,'w')
        #callistfile.writelines(l + '\n' for l in darkproclist) no need to do this repeatedly...
        callistfile.writelines(l + '\n' for l in onlists[j])
        callistfile.writelines(l + '\n' for l in offlists[j])
        callistfile.close()

        scriptfile.write('nprepare ("@' + callistfilename + '")\n')
   
        # create the flats.

        nonlistfilename = 'nflaton'+filt+'.list'
        nonlistfile = open(nonlistfilename,'w')
        nonlistfile.writelines(l + '\n' for l in prepend(onlists[j],'n'))
        nonlistfile.close()

        nofflistfilename = 'nflatoff'+filt+'.list'
        nofflistfile = open(nofflistfilename,'w')
        nofflistfile.writelines(l + '\n' for l in prepend(offlists[j],'n'))
        nofflistfile.close()
      
        flatfilename = 'flat_'+filt+'.fits'
        badpixfilename = 'bpm_'+filt+'.fits' # shouldn't be filter-dependent, but we do get one per filter.

        scriptfile.write('niflat @' + nonlistfilename + ' lampsoff=@'+nofflistfilename + ' darks=@'+ndarklistfilename + ' flatfile='+flatfilename+' bpmfile='+badpixfilename+'\n')
        #badpixmaskfiles.append('n'+onlists[j].split('.') [0] + '_bpm.pl')

        #scriptfile.write('del @' + ndarklistfilename + '\n') no longer necessary... why did I do this before?

   # we now have the flats and bad pixel masks.

    scriptfile.write("\n")

    for n in range(nstack):
        obj = stackobjs[n].replace(' ','')
        filt = stackfilts[n]
        objlistfilename = obj+filt+'.list'
        nobjlistfilename = 'n'+obj+filt+'.list'
        bnobjlistfilename = 'bn'+obj+filt+'.list'
        rbnobjlistfilename = 'rbn'+obj+filt+'.list'
        objlistfile = open(objlistfilename,'w')
        objlistfile.writelines(l + '\n' for l in stacklists[n])
        objlistfile.close()
        nobjlistfile = open(nobjlistfilename,'w')
        nobjlistfile.writelines(l + '\n' for l in prepend(stacklists[n],'n'))
        nobjlistfile.close()
        bnobjlistfile = open(bnobjlistfilename,'w')
        bnobjlistfile.writelines(l + '\n' for l in prepend(stacklists[n],'bn'))
        bnobjlistfile.close()
        skysubstack = []  # contains duplicates
        scisubstack = []  # doesn't contain duplicates
        if len(stacklists[n]) <= skyblock*1.5:
            skysubstack = [prepend(stacklists[n],'bn')]
            scisubstack = [prepend(stacklists[n],'bn')]
            nsubstack = 1
        else:
            nsubstack = int(ceil( len(stacklists[n])*1.0 / skyblock ))
            for i in range(nsubstack-1):
                skysubstack.append(prepend(stacklists[n],'bn')[i*skyblock:(i+1)*skyblock])
                scisubstack.append(prepend(stacklists[n],'bn')[i*skyblock:(i+1)*skyblock])
            skysubstack.append(prepend(stacklists[n],'bn')[-skyblock:])
            scisubstack.append(prepend(stacklists[n],'bn')[(nsubstack-1)*skyblock:])
        bnskyobjlistfilenames = []
        bnsciobjlistfilenames = []
        skyfilenames = []
        for i in range(nsubstack):
            bnskyobjlistfilename = 'sky'+str(i+1)+'_'+filt+'_'+obj+'.list'
            skyfilename =          'sky'+str(i+1)+'_'+filt+'_'+obj+'.fits'
            bnskyobjlistfilenames.append(bnskyobjlistfilename)
            skyfilenames.append(skyfilename)
            bnskyobjlistfile = open(bnskyobjlistfilename,'w')
            bnskyobjlistfile.writelines(l + '\n' for l in skysubstack[i])
            bnskyobjlistfile.close()
            bnsciobjlistfilename = 'skyproc'+str(i+1)+'_'+filt+'_'+obj+'.list'
            bnsciobjlistfilenames.append(bnsciobjlistfilename)
            bnsciobjlistfile = open(bnsciobjlistfilename,'w')
            bnsciobjlistfile.writelines(l + '\n' for l in scisubstack[i])
            bnsciobjlistfile.close()
        rbnobjlistfile = open(rbnobjlistfilename,'w')
        rbnobjlistfile.writelines(l + '\n' for l in prepend(stacklists[n],'rbn'))
        rbnobjlistfile.close()

        badpixfilename = 'bpm_'+filt+'.fits'
        #skyfilename = 'sky_'+filt+'_'+obj+'.fits'  #prepend(stacklists[n],'bn') [0] + '_sky.fits'
        #skyprefilename = 'skypre_'+filt+'_'+obj+'.fits'  #prepend(stacklists[n],'bn') [0] + '_sky.fits'
        flatfilename = 'flat_'+filt+'.fits'                      
        outfilename = obj+'_'+filt+'.fits'
        outprefilename = 'p'+obj+'_'+filt+'.fits'

        scriptfile.write('nprepare ("@' + objlistfilename + '", bpm="'+badpixfilename+'")\n')
        
        scriptfile.write('nresidual @' + nobjlistfilename + ' proptime=0.5\n')

        for i in range(nsubstack):
           scriptfile.write('nisky @' +bnskyobjlistfilenames[i]+ ' outimage='+skyfilenames[i] + '\n')
           scriptfile.write('nireduce @'+bnsciobjlistfilenames[i] + ' skyimage='+skyfilenames[i] + ' flatimage='+flatfilename+'\n')

        scriptfile.write('imcoadd @'+rbnobjlistfilename + ' geofitgeom=shift rotate- fl_over+ fl_scale- fl_fixpix+ fl_find+ fl_map+ fl_trn+ fl_med+ fl_add+ fl_avg+ badpix="'+badpixfilename+'" niter=1 outimage='+outfilename+'\n')

        if supersky > 0:

            scriptfile.write('nisupersky ' + outfilename + '\n')

            for i in range(nsubstack):
                scriptfile.write('mv ' + skyfilenames[i] + ' ' + 'p'+skyfilenames[i]+ '\n')
                scriptfile.write('nisky @' +bnskyobjlistfilenames[i]+ ' outimage='+skyfilenames[i] + '\n')

            scriptfile.write('del @' + rbnobjlistfilename+ '\n')

            for i in range(nsubstack):
                scriptfile.write('nireduce @'+bnsciobjlistfilenames[i] + ' skyimage='+skyfilenames[i] + ' flatimage='+flatfilename+'\n')
    
            scriptfile.write('mv ' + outfilename + ' ' + outprefilename+ '\n')

            scriptfile.write('imcoadd @'+rbnobjlistfilename + ' geofitgeom=shift rotate- fl_over+ fl_scale+ fl_fixpix+ fl_find+ fl_map+ fl_trn+ fl_med+ fl_add+ fl_avg+ badpix="'+badpixfilename+'" niter=1 outimage='+outfilename+'\n')

        scriptfile.write("\n")


    scriptfile.close()

    print 'Reduction script written to '+scriptfilename


###################
def main():

    argv = sys.argv

    supersky = 0
    input = ''
    input2 = ''
    fileselector = [] #default
    if (len(argv)>=2): 
        for v in argv[1:]:
            if v[0] != '-': fileselector.append(v)
            if v[0] == '-':
               if v == '-s': supersky = 1
            
    if len(fileselector) == 0: fileselector = ['*.fits']
    if len(fileselector) > 1: fileselector = fileselector[1:]
    
    files = []
    for fs in fileselector:
        files += glob.glob(fs)
    
    nbadfilenames = 0
    for i in range(len(files)):
        f = files[i]
        #print f
        if f.find('.fit') < 0: 
            if (len(argv)>=2): print f + ' is not a fits file.'
            files[i] = ''
            nbadfilenames = nbadfilenames + 1
                
    if len(files) - nbadfilenames <= 0:
        print 'Error: No valid .fits files specified.'
        return 1
    
    niriscript(files, supersky)
    
###################

if __name__=='__main__':
    main()