#!/usr/bin/env python

#
# 2014-04-02
#   readh : debug 
#           it failed for the narrow band twilight flats 
#           it has been fixed by adopting pyfits.getheader
#
# 2014-04-01
#   bkcheck : debug %%s -> %s 
#
# 2014-02-07 
#   bkcheck() : the unstable messages are displayed only for the J-band
#   printhelp() is created
#   main : (1) checking of arguments is improved so that errorneous inputs are rejected
#          (2) the parameters are recorded in mktwflist.param
#          (3) non-default parameters are displayed 
#

import os
import sys
import re
import commands
import pyfits


def readh(date, okey):

    #
    # reading the headers of all the FITS files in ../[date]/rawdata
    # and extracting twiligh flat frames based on the OBJECT name
    # and dividing them into morning and evening based on the LT
    #

    fflag = 0
    eflag = 0
    mflag = 0

    foe = open('eflat.tmp', 'w')
    fom = open('mflat.tmp', 'w')
    flist = os.listdir('../'+date+'/rawdata')
    for file in flist:
        if re.search('fits', file):
            rfile = '../'+date+'/rawdata/'+file
            if file[-2:] == 'ic' or file[-2:] == 'fz':
                hdunum = 1
            else:
                hdunum = 0
            hdr = pyfits.getheader(rfile, hdunum)
            object = hdr['OBJECT']
            timelthh = int(hdr['TIME_LT'][:2])

            if object == okey:
                fflag += 1
                if timelthh > 15 and timelthh < 22:
                    print >> foe, file
                    eflag += 1
                if timelthh > 2 and timelthh < 9:
                    print >> fom, file
                    mflag += 1

    foe.close()
    fom.close()

    return fflag, eflag, mflag



def bkcheck(band, mode, date, numarr, lvlarr):

    flag = 0
    fw = open('mktwfwarning.txt', 'aw')
    fo = open(band+'background.txt', 'aw')
    print >> fo, date, numarr[0], lvlarr[0] 
    for i in range(1, len(numarr)):
        print >> fo, date, numarr[i], lvlarr[i], 
        if lvlarr[i] - lvlarr[i-1] < 0 :
            flag += 1
            if band == 'j':
                print 'background level unstable at %s %s ' % (date, numarr[i-1]) 
                print >> fw, 'background level unstable at %s %s ' % (date, numarr[i-1])
            print >> fo, '*'
        else:
            print >> fo, ''
    print >> fo, ''
    fo.close()
    fw.close()



def pickpairs(band, mode, date, numarr, lvlarr, skydiff):

    #
    # reversing the order to pick pairs from frames with a higher background level
    #
    numarr.reverse()
    lvlarr.reverse()

    num = 0
    fo = open(band+'flatlist.0', 'aw')
    for i in range(len(numarr)):
        for j in range(i+1, len(numarr)):
            if lvlarr[i] - lvlarr[j] > skydiff:
                num += 1
                print >> fo, date, numarr[i], numarr[j], lvlarr[i], lvlarr[j]
                break
    fo.close()

    return num


def mklist(band, mode, date, skydiff, maxgood):

    # 
    # fdict = {'frame number(str)':backgroundlevel(float)}
    #
    fdict = {}
    f = open('twflat.tmp')
    for line in f:
        v = line[:-1].split()
        num = v[0][8:12]
        if v[0][0] == band:
            fdict[num] = float(v[1])
    f.close()

    if mode == 'e':
        vbool = True
    elif mode == 'm':
        vbool = False

    numarr = [] # frame number
    lvlarr = [] # background level

    #
    # the frames are sorted so that a twilight free frame comes first 
    # saturated frames are excluded 
    #
    for k, v in sorted(fdict.items(), reverse=vbool):
        if v < maxgood:
            numarr.append(k)
            lvlarr.append(v)
        else:
            break
        
    # check if the background level is monotonically increasing or decreasing 
    bkcheck(band, mode, date, numarr, lvlarr)    

    # pick pairs of frames and write to band+'flatlist'
    num = pickpairs(band, mode, date, numarr, lvlarr, skydiff)

    return num 


def listofaday(date, okey, skydiff, maxgood, bandarr, tnum):

    SCRIPTDIR = os.path.realpath(os.path.dirname(sys.argv[0]))
    cprog = 'twflatlev'

    #
    # reading the FITS header of the rawdata of the day
    # and counting the number of twilight flat frames 
    # and creating eflat.tmp and mflat.tmp
    #
    fflag, eflag, mflag = readh(date, okey)

    #
    # fflag : the total number of twilight flat frames
    # eflag : the total number of the evening twilight flat frames
    # mflag : the total number of the morning twilight flat frames
    #

    if fflag == 0:
        print 'No flat frames taken on %s' % date
        os.remove('eflat.tmp')
        os.remove('mflat.tmp')
        return tnum

    if eflag > 0:
        os.system('%s/%s  %s eflat.tmp' % (SCRIPTDIR, cprog, date))
        for band in bandarr:
            num = mklist(band, 'e', date, skydiff, maxgood)
            tnum[band] += num

    if mflag > 0:
        os.system('%s/%s %s mflat.tmp' % (SCRIPTDIR, cprog, date))
        for band in bandarr:
            num = mklist(band, 'm', date, skydiff, maxgood)
            tnum[band] += num

    os.remove('eflat.tmp')
    os.remove('mflat.tmp')
    os.remove('twflat.tmp')

    return tnum


def exslope(band):

    SCRIPTDIR = os.path.realpath(os.path.dirname(sys.argv[0]))
    cprog = 'exslopefz'

    snum = 0
    fo = open(band+'flatlist', 'w')
    f = open(band+'flatlist.0')
    for line in f:
        v = line[:-1].split()
        fits1 = '../'+v[0]+'/rawdata/'+band+v[0]+'_'+v[1]+'.fits'
        fits2 = '../'+v[0]+'/rawdata/'+band+v[0]+'_'+v[2]+'.fits'
        if os.path.exists(fits1) and os.path.exists(fits2):
            file1 = fits1
            file2 = fits2
        elif os.path.exists(fits1+'.fz') and os.path.exists(fits2+'.fz'):
            file1 = fits1+'.fz[1]'
            file2 = fits2+'.fz[1]'
        elif os.path.exists(fits1+'.ic') and os.path.exists(fits2+'.ic'):
            file1 = fits1+'.ic[1]'
            file2 = fits2+'.ic[1]'
        else:
            print 'Something wrong with compression for %s and %s' % (fits1, fits2)
            print 'ic or fz?'
            continue

        flag = commands.getoutput('%s/%s %s %s' % (SCRIPTDIR, cprog, file1, file2))
        if flag == '1':
            snum += 1
            print >> fo, line[:-1]

    f.close()
    fo.close()

    if snum == 0:
        print 'There is no good flat frame for %s-band.' % band
    else:
        print 'a total of %d pairs are selected for %s' % (snum, band)



def printhelp():

    print 'Usage # mktwflist.py date-list'
    print 'option'
    print '-noexgap'
    print '-skydiff=float (default 2000)'
    print '-maxgood=float (default 6000)'
    print '-keyword=string (default twflat)'
    print '-band=string (ex. -band=j -band=h,k)'

    
    
if __name__ == "__main__":

    argvs = sys.argv
    argc = len(argvs)

    skydiff = 2000.
    maxgood = 6000.  
    exflag = 0
    okey = 'twflat'
    bandarr = ['j', 'h', 'k']

    # checking the arguments 

    if argc < 2:
        printhelp()
        sys.exit()

    for item in argvs:
        if re.search('help', item):
            printhelp()
            sys.exit()

    for i in range(2, argc):
        if argvs[i] != '-noexgap' and not re.match('-skydiff=', argvs[i])  and not re.match('-maxgood=', argvs[i]) and not re.match('-keyword=', argvs[i]) and not re.match('-band=', argvs[i]):
            print '%s is an illegal argument' % argvs[i]
            printhelp()
            sys.exit()

    nondefault = ''
    oanum = 1
    for item in argvs:
        if item == '-noexgap':
            exflag = -99
            oanum += 1
            nondefault += ' -noexgap'
        elif item.startswith('-skydiff'):
            try:
                skydiff = float(item[9:])
            except:
                print 'Invalid : '+item
                sys.exit()
            nondefault += ' -skydiff='+str(skydiff)
            oanum += 1
        elif item.startswith('-maxgood'):
            try:
                maxgood = float(item[9:])
            except:
                print 'Invalid : '+item
                sys.exit()
            nondefault += ' -maxgood='+str(maxgood)
            oanum += 1
        elif item.startswith('-keyword'):
            try:
                okey = item[9:]
            except:
                print 'Invalid : '+item
                sys.exit()
            nondefault += ' -keyword='+okey
            oanum += 1
        elif item.startswith('-band'):
            if not re.search('[jhk]', item[6:]):
                print 'Invalid : '+item
                sys.exit()
            bandarr = item[6:].split(',')
            nondefault += ' '+item
            oanum += 1

    if oanum == argc:
        print 'Input the list'
        printhelp()
        sys.exit()

    # end of checking the arguments 

    if nondefault != '':
        print nondefault

    # keep the parameters in mktwflist.param
    fo = open('mktwflist.param', 'w')
    if exflag == -99:
        print >> fo, '-noexgap'
    print >> fo, 'skydiff='+str(skydiff)
    print >> fo, 'maxgood='+str(maxgood)
    print >> fo, 'keyword='+okey
    print >> fo, 'band=', bandarr
    fo.close()

    tnum = {}
    for band in bandarr:
        tnum[band] = 0

    f = open(argvs[1])
    for line in f:
        date = line[:-1]
        if re.search('\S', date):
            print date
            tnum = listofaday(date, okey, skydiff, maxgood, bandarr, tnum)
    f.close()

    #
    # excluding pairs with the gap at the central column
    #
    if exflag == 0:
        for band in bandarr:
            if tnum[band] == 0:
                print 'No frame for %s-band flat' % band 
            else:
                exslope(band)
    else:
        for band in bandarr:
            os.rename(band+'flatlist.0', band+'flatlist')
            print 'a total of %d pairs are selected for %s' % (tnum[band], band)

    if os.path.getsize('mktwfwarning.txt') == 0:
        os.remove('mktwfwarning.txt')

    print 'Done'




