#!/usr/bin/env python

# 2014-02-17  'EXTEND' in the primary header of recombined FITS files to be 'T' in recomb() 

import os
import sys
import re
import shutil
import glob
import gmaptran
import sirrecom
import pyfits
import math
from pyraf import iraf
from iraf import obsutil
from iraf import imred
from iraf import ccdred

def initrecom(band):

    if os.path.exists('../0.ffiles'):
        print 'rename ../0.ffiles to another'
        sys.exit()

    curdir = os.getcwd()
    darr = curdir.split('/')
    updir = curdir[:-1*len(darr[len(darr)-1])]

    print updir

    dlist = os.listdir('..')
    fdlist = []
    for v in dlist:
        if v == 'ffiles':
            fdlist.append('0.ffiles')
            os.symlink(updir+'ffiles', updir+'0.ffiles')
        elif re.search('ffiles$', v):
            fdlist.append(v)

    if len(fdlist) > 1:
        fdlist.sort()

    if os.access(band+'imgrecom.list', os.F_OK):
        os.remove(band+'imgrecom.list')

    listflag=0

    for v in fdlist:

        onum = []
        objname = {}
        raoff = {}
        decoff = {}
        itime = {}
        flist = os.listdir('../'+v)
        for file in flist:
            if re.match(band+"f", file) and re.search(".fits$", file):

                hdu = pyfits.open('../'+v+'/'+file)
                tobjname = hdu[0].header['OBJECT'] 
                titime = hdu[0].header['EXPOS'] 
                traoff = hdu[0].header['RA_OFF'] 
                tdecoff = hdu[0].header['DEC_OFF'] 
                hdu.close()

                fnum = file[2:6]
                objname[fnum] = tobjname
                raoff[fnum] = traoff
                decoff[fnum] = tdecoff
                itime[fnum] = titime
                onum.append(fnum)

        onum.sort()
        if len(onum) == 0:
            print band+'band frames not found in ../'+v
        else:
            if v == '0.ffiles':
                print 'ffiles'
            else:
                print v
            listflag = listflag + 1
            fo = open(band+'imgrecom.list', 'aw')
            for v1 in onum:
                print >> fo, v, v1, objname[v1], raoff[v1], decoff[v1], itime[v1]
            fo.close()

    if listflag == 0:
        print band+'imgrecom.list has not been created'
    elif listflag == 1:
        print 'is used for ' + band + 'imgrecom.list'
    else:
        print 'are used for ' + band + 'imgrecom.list'



def mkimgrecom(dlist, band):
    
    #
    # a function for the recombine process in the pipeline 
    #
    
    fo = open(band+'imgrecom.list', 'w')
    fi = open('../'+dlist, 'r')
    for line in fi:
        v = line[:-1].split()
        print >> fo, '0.ffiles %s dummy %s %s 0' % (v[0], v[1], v[2])

    fi.close()
    fo.close()



def recomb(band, urms, lstarnum, reject):

    objname = 'recom'

    #
    # comarea
    #
    f = open(band+'recominfo')
    line = f.readline()
    values = line[:-1].split()
    comarea = '['+values[4]+':'+values[5]+','+values[6]+':'+values[7]+']'
    f.close()

    flist = band + 'tran.txt'
    fo = open('comlist', 'w')
    f = open(flist)
    i = 0
    for line in f:
        values = line[:-1].split()
        frame = values[0]
        xrms = float(values[12])
        yrms = float(values[13])
        starnum = float(values[14])
        if os.access(frame+'.fits', os.R_OK) and xrms <= urms and yrms <= urms and starnum >= lstarnum:
            print >> fo, "%s" % frame
            if i == 0:
                frame0 = frame
                i = i + 1
        else:
            errmsg = ''
            if xrms > urms:
                errmsg = errmsg + 'xrms = %.1f ' % xrms
            if yrms > urms:
                errmsg = errmsg + 'yrms = %.1f ' % yrms
            if starnum < lstarnum:
                errmsg = errmsg + 'starnum = %d ' % starnum

            if frame[:2] == '0.':
                print '%s was excluded because %s' % (frame[2:], errmsg)
            else:
                print '%s was excluded because %s' % (frame, errmsg)

    f.close()
    fo.close()

    iraf.images()
    iraf.imutil()
    iraf.imarith.unlearn()

    iraf.imstat.unlearn()
    iraf.imcombine.unlearn()
    iraf.ccdhedit.unlearn()
    iraf.wfits.unlearn()
    iraf.imdelete.unlearn()

    iraf.imstat.fields = 'midpt'
    iraf.imstat.lower = 'INDEF'
    iraf.imstat.upper = 'INDEF'
    iraf.imstat.nclip = 10
    iraf.imstat.lsigma = 3
    iraf.imstat.usigma = 3
    iraf.imstat.binwidth = 0.1
    iraf.imstat.format = 'yes'
    iraf.imstat.cache = 'yes'

    f = open('comlist')
    ncom = 0
    for frame in f:
       if ncom == 0:
           zframe = frame[:-1]
       else:
           sframe = frame[:-1]
           iraf.imarith(sframe+comarea, '-', zframe+comarea, 'difftmp')
           out_imstat = iraf.imstat('difftmp', Stdout=1)
           med = out_imstat[1]
           iraf.imarith(sframe, '-', med, sframe)
           iraf.imdel('difftmp')
       ncom = ncom + 1
    f.close()

    iraf.imstat.fields = 'midpt, stddev'
    iraf.imstat.lower = -1000
    out_imstat = iraf.imstat(frame0, Stdout=1)
    med, stdev = out_imstat[1].split()
    lowth = float(med) - 10 * float(stdev)
    
    iraf.imcombine( '@comlist', band+objname,
                    headers='', bpmasks='', rejmask='', nrejmasks='nrej',
                    expmasks='', sigma='', logfile='', combine='average',
                    reject=reject, project='no', outtype='real', outlimi='', offsets='none',
                    masktyp='none', maskval=0, blank=-9999, scale='none', zero='none', weight='none',
                    statsec='', expname='', lthresh=lowth, hthresh='INDEF',
                    nlow=0, nhigh=1, nkeep=1, mclip='yes', lsigma=3, hsigma=3,
                    rdnoise=15, gain=6, snoise='0.', sigscale=0.1,
                    pclip=-0.5, grow=0.)

    iraf.ccdhedit(band+objname, paramete="COMMETHOD", value="average", type="string")
    iraf.ccdhedit(band+objname, paramete="COMAREA", value=comarea, type="string")
    iraf.ccdhedit(band+objname, paramete="REJECT", value=reject, type="string")

    iraf.wfits('nrej.pl',band+'count.fits', Stdout=1)
    iraf.imarith(band+'count', '*', -1, band+'count')
    iraf.imarith(band+'count', '+', ncom, band+'count')

    fcount = pyfits.open(band+'count.fits')
    scidata = fcount[0].data
    pyfits.append(band+objname+'.fits', scidata)
    fcount.close()

    # the FITS file has an extension 2014-02-17
    fimg = pyfits.open(band+objname+'.fits', mode='update')
    prihdr = fimg[0].header
    prihdr['EXTEND'] = True
    fimg.close()

    
    iraf.imdelete('@comlist')
    os.remove('nrej.pl')
    os.remove(band+'recominfo')
    os.remove('comlist')
    os.remove(band+'count.fits')
    os.remove(band+'tran.txt')
    os.remove(band+'fwhm.txt')


def dorecom(band, ufwhm, uellip, urms, lstarnum, searchrad, mlim, reject, theta, ifwhm, minsep):

    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)
    
    itime = 0
    f = open(band+'imgrecom.list')
    for line in f:
        if re.search('\S', line[:-1]):
            val = line[:-1].split()
            if itime == 0:
                itime = float(val[5])
            elif float(val[5]) != itime:
                print 'multiple exposure time in %simgrecom.list' % band
                sys.exit()
    f.close()

    if os.path.exists('../ffiles') and not os.path.exists('../0.ffiles'):
        curdir = os.getcwd()
        darr = curdir.split('/')
        updir = curdir[:-1*len(darr[len(darr)-1])]
        os.symlink(updir + 'ffiles', updir + '0.ffiles') 
        # os.symlink('../ffiles', '../0.ffiles') does not work as the above does. 

    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)
    gflag = gmaptran.getgmap(band, ufwhm, uellip, searchrad, minsep, urms, lstarnum)
    # gmaptran.recominfo(band, urms, lstarnum)

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

        print 'combining'
        recomb(band, urms, lstarnum, reject)

    if os.path.exists('../0.ffiles'):
        os.remove('../0.ffiles')

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

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

    print 'done'



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

    #
    # display how to use
    #
    if argc < 2 or argvs[1] == '--help':
        print ''
        print 'usage: imgrecom.py init band[j|h|k]'
        print 'to make the frame list for imgrecom.py recom'
        print ''
        print 'usage: imgrecom.py recom band[j|h|k]'
        print 'to combine frames'
        print 'options: '
        print '-search=[float] : search radius for the sparsematch algorithm'
        print '-mlim=[float] : magnitude error upper limit'
        print '-ufwhm=[float] : upper limit of fwhm'
        print '-uellip=[float] : upper limit of ellipticity'
        print '-urms=[float] : upper limit of xrms and yrms in result of iraf.geomap. '
        print '-lstarnum=[int] : lower limit of the number of stars used in iraf.geotran . '
        print '-theta=[float] : the rotation angle of frames measured unti-clockwise in degrees'
        print '-ifwhm=[float] : an initial guess of fwhm'
        print '-minsep=[float] : minimum separation in xyxymatch'
        print '-reject=[sigclip|minmax] : the reject parameter of imcombine'
        print 'search=20, mlim=0.2, ufwhm=20, uellip=1, urms=1, lstarnum=1, theta=0, ifwhm=3.5, minsep=9 and reject=sigclip by default'
        print ''
        print 'usage: imgrecom.py --help'
        print 'to show this messege'
        print ''
        sys.exit()

    if argvs[1]!='init' and argvs[1]!='recom':
        print 'usage: imgrecom.py [init|recom] band[j|h|k] -option'
        sys.exit()
    if argvs[2]!='j' and argvs[2]!='h' and argvs[2]!='k':
        print 'usage: imgrecom.py [init|recom] band[j|h|k] -option'
        sys.exit()

    #
    # the default values of ufwhm and uellip are 20 and 1
    #
    searchrad = 20
    mlim = 0.2
    ufwhm = 20
    uellip = 1
    urms = 1
    lstarnum = 1
    reject = 'sigclip'
    theta = 0
    ifwhm = 3.5
    minsep = 9
    if argc > 3:
        for i in range(3, argc, 1):
            print argvs[i]
            if re.match('-ufwhm', argvs[i]):
                ufwhm = float(argvs[i][7:])
            elif re.match('-search', argvs[i]):
                searchrad = float(argvs[i][8:])
            elif re.match('-mlim', argvs[i]):
                mlim = float(argvs[i][6:])
            elif re.match('-uellip', argvs[i]):
                uellip = float(argvs[i][8:])
            elif re.match('-urms', argvs[i]):
                urms = float(argvs[i][6:])
            elif re.match('-lstarnum', argvs[i]):
                lstarnum = float(argvs[i][10:])
            elif re.match('-reject', argvs[i]):
                reject = argvs[i][8:]
            elif re.match('-theta', argvs[i]):
                theta = math.radians(float(argvs[i][7:]))
            elif re.match('-ifwhm', argvs[i]):
                ifwhm = float(argvs[i][7:])
            elif re.match('-minsep', argvs[i]):
                minsep = float(argvs[i][8:])
            else:
                print 'typo? -search=[float] -mlim=[float] -ufwhm=[float] -uellip=[float] -urms=[float] -lstarnum=[int] -theta=[float] -ifwhm=[float] -minsep=[float] -reject=[sigclip|minmax]' 


    if argvs[1]=='init':
        initrecom(argvs[2])
    if argvs[1]=='recom':
        if os.access(argvs[2]+'imgrecom.list', os.F_OK):
            dorecom(argvs[2], ufwhm, uellip, urms, lstarnum, searchrad, mlim, reject, theta, ifwhm, minsep)
        else:
            print argvs[2]+'imgrecom.list not found'
            print 'execute imgrecom.py init '+argvs[2]


                
