#!/usr/bin/env python

#
# sirpol pipeline script
#

import os
import sys
import re
import shutil
import math
import glob
import irsf
import pyfits
import sirrecom
import polrecom
import gmaptran
from pyraf import iraf


def mkpar(band):

    if os.path.exists('sirpol.param'):
        os.remove('sirpol.param')

    #
    # check directory structure
    #
    if not os.path.exists('../rawdata'):
        print 'no ../rawdata found'
        sys.exit()

    #
    # read yymmdd
    #
    yymmdd = irsf.yymmdd()

    #
    # check prepared flat fits files
    #
    twfname = irsf.chkflat(band)

    f = open('sirpol.param', 'w')
    print >> f, 'band %s   # [j|h|k|all]' % band
    print >> f, 'yymmdd %s   # yymmdd' % yymmdd 
    print >> f, 'twflat %s   # flat name, [j|h|k]+twflat+[.fits]' % twfname
    print >> f, 'stop 0   # to stop pipeline after 1:dark+flat 2:sky-subtraction'
    print >> f, 'crremove 0   # cosmic ray removal 0:no 1::min-max'
    print >> f, 'linear 0   # linearity correction 0:no 1:yes'
    print >> f, 'darkopt 0   # use prepared dark frames? 0:no 1:yes'
    print >> f, 'comsky 0   # combine sky frames? 0:no 1:yes'
    print >> f, 'allcom 1   # combine all the sets with the same object name? 0:no 1:yes'
    print >> f, 'polmode 1 # 1:LP, 2:CP2, 3:CP4'
    print >> f, 'cfitsio 1   # use cfitsio library? 0:no 1:yes'
    print >> f, 'fiim 1 # make I image? 0:no 1:yes'
    print >> f, 'fqim 1 # make Q image? 0:no 1:yes'
    print >> f, 'fuim 1 # make U image? 0:no 1:yes'
    print >> f, 'fvim 0 # make V image? 0:no 1:yes'
    print >> f, 'fpim 0 # make PI image? 0:no 1:yes'
    print >> f, 'f00im 1 # make 0deg image? 0:no 1:yes'
    print >> f, 'f22im 1 # make 22.5deg image? 0:no 1:yes'
    print >> f, 'f45im 1 # make 45deg image? 0:no 1:yes'
    print >> f, 'f67im 1 # make 67.5deg image? 0:no 1:yes'
    print >> f, 'f90im 0 # make 90deg image? 0:no 1:yes'
    print >> f, 'f135im 0 # make 135deg image? 0:no 1:yes'
    print >> f, 'keepd 0   # keep dark.fits files? 0:no 1:yes'
    print >> f, 'keeps 0   # keep sky.fits files? 0:no 1:yes'
    print >> f, 'keepm 0   # keep mask.fits files? 0:no 1:yes'
    f.close()

    # print >> f, 'slope 0   # slope correction 0:no 1:yes'
    # print >> f, 'fqcim 0 # make QC image? 0:no 1:yes'
    # print >> f, 'fucim 0 # make UC image? 0:no 1:yes'



def mklist():

    mstatus = 0

    #
    # check if the necessary file exist
    #
    if not os.access('./obslog', os.R_OK):
        print 'no ./obslog found'
        sys.exit()

    # obslog.fnum, obslog.oname, and obslog.itime are frame number, object name, and integration time                                            
    # for each frame listed in an obslog file except for dark and flat                                                            
    obslog = irsf.readobslog()

    if not os.access('./sirpol.param', os.R_OK):
        print 'no ./sirpol.param found'
        sys.exit()

    #
    # get user parameters 
    #
    para = irsf.readparam()
    if para.flag == 1:
        print 'set sirpol.param properly'
        sys.exit()

    #
    # dark list 
    #
    if para.darkopt == '0' or para.darkopt == '1':
        mstatus = irsf.mkdarklist(para.darkopt, para.yymmdd, para.blist)
        if mstatus == 1:
            sys.exit()

    #
    # dark+flat
    #
    for band in para.blist:
        irsf.mkdarkflat(para.twflat, para.yymmdd, band, obslog)


    #
    # skylist
    #
    mstatus = irsf.object2sky(obslog)
    if mstatus == 1:
        fmkl = open('darklist')
        for line in fmkl:
            os.remove(line[:-1])
        fmkl.close()
        os.remove('darklist')
        
        for band in para.blist:
            fmkl = open(band+'darkflatlist')
            for line in fmkl:
                os.remove(line.split()[0])
            fmkl.close()
            os.remove(band+'darkflatlist')
        
        sys.exit()

    numpoldict = irsf.numpol(para.yymmdd, para.blist, para.polmode, obslog)
    irsf.polsky(para.polmode, numpoldict)


    #
    # dithlist
    #
    irsf.mkdithlist(para.allcom, obslog, para.comsky)




def pipeline():

#    if re.search('bash', os.environ['SHELL']) or re.search('zsh', os.environ['SHELL']):
#        os.system('ulimit -s unlimited > /dev/null 2>&1')
#    elif re.search('csh', os.environ['SHELL']) or re.search('ksh', os.environ['SHELL']):
#        os.system('unlimit stacksize')

    os.system('date')

    iraf.set(writepars=0)

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

    obslog = irsf.readobslog()

    para = irsf.readparam()
    if para.flag == 1:
        print 'set sirpol.param properly'
        sys.exit()

    os.mkdir('ffiles')

    #
    # make dark frames
    #
    cprog = 'mkdarkavg'
    if para.darkopt == '0':
        print 'dark frames ... ', 
        sys.stdout.flush()
        f = open('darklist')
        for lname in f:
            fitsname = lname[:-1].replace('list', 'fits')
            os.system('%s/%s %s %s' % (SCRIPTDIR, cprog, lname[:-1], fitsname))
            os.remove(lname[:-1])
        f.close()
        os.remove('darklist')
        os.remove('mkdark.tmp')
        print 'have been created'


    #
    # dark subtraction and flat-field division 
    #
    print 'dark subtraction and flat division ... ', 
    sys.stdout.flush()

    if para.cfitsio == '0' :
        if para.linear == '0':
            cprog = 'darkflatncf'
        else:
            cprog = 'darkflatncf2'
    else:
        if para.linear == '0':
            cprog = 'darkflatfz'
        else:
            cprog = 'darkflatfz2'

    for band in para.blist:
        f = open(band+'darkflatlist')
        for line in f:
            v = line[:-1].split()
            os.system('%s/%s %s %sdark%s.fits %s%s.fits' % (SCRIPTDIR, cprog, v[0], band, v[1], band, para.twflat))
            os.remove(v[0])
        f.close()
        os.remove(band+'darkflatlist')

    print 'done'

    # stop 1
    if para.stop == '0.5':
        sys.exit()
    
    #
    # bad pixel mask
    #
    # zero shifted median combined sky are normalized 
    #
    print 'making sky bias frames and bad pixel masks ... ', 
    sys.stdout.flush()

    iraf.images()
    iraf.images.unlearn()
    iraf.imutil()
    iraf.imutil.unlearn()
    iraf.immatch()
    iraf.immatch.unlearn()

    iraf.imdelete.unlearn()
    iraf.imstat.unlearn()
    iraf.imarith.unlearn()
    iraf.imcombine.unlearn()

    if para.polmode == '1':
        pangarr = ['a00', 'a22', 'a45', 'a67']
    elif para.polmode == '2':
        pangarr = ['a000', 'a090']
    elif para.polmode == '3':
        pangarr = ['a000', 'a090', 'a045', 'a135']

    skyarr = []
    f = open('objectskylist')
    for line in f:
        v = line[:-1].split()
        if not v[3] in skyarr:
            skyarr.append(v[3])
            for pang in pangarr:
                for band in para.blist:
                    irsf.zerosky(band, v[3]+pang) # band+v[3]+pang+'.fits' and 'n'+band+v[3]+pang+'.fits' will be created
    f.close()

    for band in para.blist:
        irsf.skybglist(band)

    #
    # then combined to make a bad pixel mask frame for each band
    #

    iraf.unlearn('imcombine')
    iraf.imcombine.combine = 'median'
    iraf.imcombine.reject = 'sigclip'
    iraf.imcombine.mclip = 'yes'
    iraf.imcombine.blank = -1
    iraf.imcombine.logfile = ''
    cprog = 'mkbadmask'
    for band in para.blist:
        #flist = ''
        fout = open('normsky.list', 'w')
        for zsky in skyarr:
            for pang in pangarr:
                print >> fout, 'n'+band+zsky+pang+'.fits'
                #flist += 'n'+band+zsky+pang+','
        fout.close()
                
        #flist = flist[:-1]
        #iraf.imcombine(flist, band+'normsky', Stdout=0)
        iraf.imcombine('@normsky.list', band+'normsky', Stdout=0)
        os.system('%s/%s %snormsky.fits %smask.fits 1' % (SCRIPTDIR, cprog, band, band)) 
        
        for zsky in skyarr:
            for pang in pangarr:
                os.remove('n'+band+zsky+pang+'.fits')
        os.remove(band+'normsky.fits')

    
    print 'done.'

    #
    # mask [j|h|k]fxxxx.fits 
    #
    cprog = 'domask'
    f = open('obj2sky')
    for line in f:
        v = line[:-1].split()
        for band in para.blist:
            os.system('%s/%s ./ffiles/%sf%s.fits %smask.fits m' % (SCRIPTDIR, cprog, band, v[0], band))
    f.close()
    print 'done'

    # stop 1                                                                                                          
    if para.stop == '1':
        sys.exit()

    
    #
    # sky subtraction
    #
    print 'sky subtraction ... ', 
    sys.stdout.flush()
    cprog = 'domask'
    for band in para.blist:
        print band, 
        sys.stdout.flush()
        for zsky in skyarr:
            for pang in pangarr:
                zskyfits = band + zsky + pang + '.fits'
                os.system('%s/%s %s %smask.fits p' % (SCRIPTDIR, cprog, zskyfits, band))
                med, dummy = irsf.getstat(zskyfits)
                iraf.imarith(zskyfits, '-', med, zskyfits)
            
        f = open('obj2sky')
        for line in f:
            v = line[:-1].split()
            iraf.imarith('./ffiles/'+band+'f'+v[0], '-', band+v[2], './ffiles/'+band+'f'+v[0])
        f.close()

    os.remove('obj2sky')
    print 'done'


    # stop 2
    if para.stop == '2':
        sys.exit()


    #
    #  recombine
    #
    searchrad = 20
    # to do : automatic adjustment of searchrad when no star is matched
    minsep = 9
    mlim = 0.2
    ufwhm = 20
    uellip = 1
    urms = 1
    lstarnum = 0
    reject = 'sigclip'
    theta = 0
    ifwhm = 3.5

    numpoldict = irsf.numpol(para.yymmdd, para.blist, para.polmode, obslog)

    if para.crremove == '1':
        reject = 'minmax'

    os.symlink('ffiles', '0.ffiles')
    f = open('dithlistlist')
    for line in f:

        v = line[:-1].split()
        os.mkdir(v[1])
        iraf.chdir(v[1])
        itime = v[2]
        print ''
        print '*****************'
        print v[1]
        print '*****************'
        for band in para.blist:
            print band+'-band' 
            polrecom.mkpolrecom(v[0], band, numpoldict)
            print 'find stars for geomap/geotran'
            sirrecom.seqphot(band, itime, ifwhm, mlim, 'pol')

            print 'star match and geomap'
            gmaptran.pregeo(band, ufwhm, uellip, minsep, theta)
            gflag = gmaptran.getgmap(band, ufwhm, uellip, searchrad, minsep, urms, lstarnum)

            if gflag == 0:
                print 'geotran in progress'
                gmaptran.dotran(band, urms, lstarnum)
                os.remove('gtran.coo')

                polrecom.mkcomlist(band, urms, lstarnum)

                print 'making i,q,u, images for each cycle ', 
                polrecom.mkqui()
                print 'done'
                
                print 'combining'
                polrecom.polrecom(band, reject)
                print 'done'
                print ''

            os.remove(band+'fwhm.txt')
            os.remove(band+'polrecom.list')
            os.remove(band+'tran.txt')

        polrecom.clean()
        iraf.chdir()
        os.remove(v[0])

    f.close()
    print 'done'

    irsf.clean(para.keepd, para.keeps, para.keepm)

    os.system('date')

    
def helpmsg():

    print ''
    print 'usage: sirpol.py para -band[j|h|k]'
    print ''
    print 'usage: sirpol.py set'
    print ''
    print 'usage: sirpol.py run'
    print ''
    print 'usage: sirpol.py comdith'
    print ''
    print 'usage: sirpol.py revsky'
    print ''
    print 'usage: sirpol.py --help'
    print 'to show this messege'
    print ''
    sys.exit()


if __name__ == "__main__":
    
    argvs = sys.argv
    argc = len(argvs)

    if argc < 2:
        helpmsg()
    elif argvs[1] == '--help':
        helpmsg()

    if argvs[1] != 'para' and argvs[1] != 'set' and argvs[1] != 'run' and argvs[1] != 'comdith' and  argvs[1] != 'revsky':
        print 'usage: sirpol.py [para|set|run|comdith|revsky] -option'
        sys.exit()

    if argvs[1] == 'para':
        band = 'all'
        if argc > 2:
            for i in range(2, argc, 1):
                if re.match('-band=', argvs[i]):
                    band = argvs[i][6:]
                    if band != 'j' and band != 'h' and band != 'k':
                        print 'specify a correct band [j|h|k]'
                        sys.exit()
        mkpar(band)

    elif argvs[1] == 'set':
        mklist()

    elif argvs[1] == 'run':
        pipeline()
        
    elif argvs[1] == 'comdith':
        irsf.comdithlist()

    elif argvs[1] == 'revsky':
        irsf.revskylist()
