#!/usr/bin/env python

#
# sirius pipeline script
#

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

def mkpar(band):

    if os.path.exists('sirius.param'):
        os.remove('sirius.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('sirius.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 method'
    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, 'cfitsio 1   # use cfitsio library? 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, 'crremove 0   # cosmic ray removal 0:no 1:laplacian-edge 2:min-max'


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('./sirius.param', os.R_OK):
        print 'no ./sirius.param found'
        sys.exit()

    #
    # get user parameters 
    #
    para = irsf.readparam()
    if para.flag == 1:
        print 'set sirius.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:
        if os.access('darklist', os.R_OK):    #  when dark frames are prepared. darklist does not exist. 2013.4.19 
            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()


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



def pipeline():

    # iraf.prcacheOff()

    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')

    iraf.set(writepars=0)

    os.system('date')

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

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

    # mfile is removed
    mfkeep = 0

    os.mkdir('ffiles')

    #
    # make dark frames
    #
    if para.darkopt == '0':
        print 'dark frames ... ', 
        sys.stdout.flush()
        cprog = 'mkdarkavg'
        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 0.5
    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()

    skyarr = []
    f = open('objectskylist')
    for line in f:
        v = line[:-1].split()
        if not v[3] in skyarr:
            skyarr.append(v[3])    
            for band in para.blist:
                irsf.zerosky(band, v[3]) # band+v[3]+'.fits' and 'n'+band+v[3]+'.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.imcombine.unlearn()
    iraf.imcombine.combine = 'median'
    iraf.imcombine.reject = 'sigclip'
    iraf.imcombine.mclip = 'yes'
    iraf.imcombine.blank = -1
    cprog = 'mkbadmask'
    for band in para.blist:
        flist = ''
        for zsky in skyarr:
            flist += 'n'+band+zsky+','
        flist = flist[:-1]
        iraf.imcombine(flist, band+'normsky', Stdout=1)
        os.system('%s/%s %snormsky.fits %smask.fits 1' % (SCRIPTDIR, cprog, band, band)) 

        for zsky in skyarr:
            os.remove('n'+band+zsky+'.fits')
        os.remove(band+'normsky.fits')

    print 'done.'

    print 'masking ... ', 
    sys.stdout.flush()

    #
    # 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:
            zskyfits = band + zsky + '.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

    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' 

            imgrecom.mkimgrecom(v[0], band)
            print 'find stars for geomap/geotran'
            sirrecom.seqphot(band, itime, ifwhm, mlim, 'img')

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

            print 'geotran in progress'
            gmaptran.dotran(band, urms, lstarnum)
            os.remove('gtran.coo')

            print 'combining'
            imgrecom.recomb(band, urms, lstarnum, reject)
            print 'done'
            print ''

            os.remove(band+'imgrecom.list')

            for file in glob.glob('ap*.mag'):
                os.remove(file)

            if mfkeep == 0:
                for file in glob.glob('*.mag'):
                    os.remove(file)
            else:
                if not os.access('mfile', os.F_OK):
                    os.mkdir('mfile')
                for file in glob.glob('*.mag'):
                    shutil.move(file, 'mfile')

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

    f.close()

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

    print 'done'

    os.system('date')


def helpmsg():

    print ''
    print 'usage: sirius.py para -band=[j|h|k]'
    print ''
    print 'usage: sirius.py set'
    print ''
    print 'usage: sirius.py run'
    print ''
    print 'usage: sirius.py comdithlist (or comdith for short)'
    print ''
    print 'usage: sirius.py revskylist (or revsky for short)'
    print ''
    print 'usage: sirius.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] != 'comdithlist' and argvs[1] != 'comdith' and argvs[1] != 'revskylist' and argvs[1] != 'revsky':
        print 'usage: sirius.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] == 'comdithlist' or argvs[1] == 'comdith':
        irsf.comdithlist()

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




