#!/usr/bin/env python

import os
import sys
import re
import math
import shutil
import glob
import commands
import gmaptran
import sirrecom
import pyfits
import warnings
from pyraf import iraf

warnings.simplefilter('ignore', DeprecationWarning)


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

    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+'polrecom.list', os.F_OK):
        os.remove(band+'polrecom.list')

    listflag = 0

    for v in fdlist:

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

                polkeyflag = 1
                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'] 
                tpang = hdu[0].header['POL-AGL1'] 
                hdu.close()

                if polkeyflag == 0:
                    print 'FITS keyword POL-AGL1 was not found for %s/%s' % (v, file)
                else:
                    fnum = file[2:6]
                    objname[fnum] = tobjname
                    raoff[fnum] = traoff
                    decoff[fnum] = tdecoff
                    itime[fnum] = titime
                    polagl[fnum] = tpang
                    onum.append(fnum)

        onum.sort()
        if len(onum) == 0:
            print band+'band frames not found in ../'+v
            print 'or keyword POL-AGL1 is not found at the header of all the fits files in ../'+v
        else:
            if v == '0.ffiles':
                print 'ffiles'
            else:
                print v
            listflag = listflag+1
            fo = open(band+'polrecom.list','aw')
            for v1 in onum:
                print >> fo, v, v1, objname[v1], raoff[v1], decoff[v1], itime[v1], polagl[v1]
            fo.close()

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



def chkqui(tfnum, tpang):

    a00flag = 0
    a45flag = 0
    a22flag = 0
    a67flag = 0
    qflag = 0
    uflag = 0

    tqtxt = ""
    tutxt = ""
    titxt = ""
    qorder = ""
    uorder = ""

    f00 = open('a00comlist', 'aw')
    f45 = open('a45comlist', 'aw')
    f22 = open('a22comlist', 'aw')
    f67 = open('a67comlist', 'aw')
    fu = open('ulist', 'aw')
    fq = open('qlist', 'aw')
    fi = open('ilist', 'aw')
    fuc = open('ucomlist', 'aw')
    fqc = open('qcomlist', 'aw')
    fic = open('icomlist', 'aw')
    for i in range(len(tfnum)):

        if tpang[i] == '0.0':
            print >> f00, tfnum[i]
            tqtxt = tfnum[i] + " "
            qorder = "0"
            titxt = titxt + tfnum[i] + " "
            a00flag = 1
        if tpang[i] == '45.0':
            print >> f45, tfnum[i]
            tqtxt = tqtxt + tfnum[i] + " "
            qorder = qorder + "45"
            titxt = titxt + tfnum[i] + " "
            a45flag = 1
        if tpang[i] == '22.5':
            print >> f22, tfnum[i]
            tutxt = tfnum[i] + " "
            uorder = "22"
            titxt = titxt + tfnum[i] + " "
            a22flag = 1
        if tpang[i] == '67.5':
            print >> f67, tfnum[i]
            tutxt = tutxt + tfnum[i] + " "
            uorder = uorder + "67"
            titxt = titxt + tfnum[i] + " "
            a67flag = 1

    if a00flag == 1 and a45flag == 1 and qorder == '045':
        print >> fq, tqtxt
        qfname = (tqtxt.split())[0].replace('f', 'q')
        print >> fqc, qfname
        qflag = 1
    if a22flag == 1 and a67flag == 1 and uorder == '2267':
        print >> fu, tutxt
        ufname = (tutxt.split())[0].replace('f', 'u')
        print >> fuc, ufname
        uflag = 1
    if qflag == 1 and uflag == 1:
        print >> fi, titxt
        ifname = (titxt.split())[0].replace('f', 'i')
        print >> fic, ifname

    f00.close()
    f45.close()
    f22.close()
    f67.close()
    fq.close()
    fu.close()
    fi.close()
    fqc.close()
    fuc.close()
    fic.close()



def mkqui():

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

    print 'q ',
    sys.stdout.flush()
    cprog = 'spmkqux'
    f = open('qlist')
    for line in f:
        qv = line[:-1].split()
        qfname = qv[0].replace('f', 'q')
        os.system('%s/%s %s.fits %s.fits %s.fits' % (SCRIPTDIR, cprog, qv[0], qv[1], qfname))
    f.close()

    print 'u ',
    sys.stdout.flush()
    cprog = 'spmkqux'
    f = open('ulist')
    for line in f:
        uv = line[:-1].split()
        ufname = uv[0].replace('f', 'u')
        os.system('%s/%s %s.fits %s.fits %s.fits' % (SCRIPTDIR, cprog, uv[0], uv[1], ufname))
    f.close()

    print 'i ',
    sys.stdout.flush()
    cprog = 'spmkix'
    f = open('ilist')
    for line in f:
        iv = line[:-1].split()
        ifname = iv[0].replace('f', 'i')
        os.system('%s/%s %s.fits %s.fits %s.fits %s.fits %s.fits' % (SCRIPTDIR, cprog, iv[0], iv[1], iv[2], iv[3], ifname))
    f.close()



def mkcomlist(band, urms, lstarnum):

    raoff = {}
    decoff = {}
    polagl = {}
    f = open(band+'polrecom.list')
    for line in f:
        values = line[:-1].split()
        yymmdd = (values[0].split('.'))[0]
        fnum = yymmdd+'.'+band+'f'+values[1]
        raoff[fnum] = values[3]
        decoff[fnum] = values[4]
        polagl[fnum] = values[6]
    f.close()

    onum = []
    xrms = {}
    yrms = {}
    starnum = {}
    f = open(band+'tran.txt')
    for line in f:
        values = line[:-1].split()
        fnum = values[0]
        if int(values[14]) >= lstarnum and float(values[12]) < urms and float(values[13]) < urms:
            xrms[fnum] = float(values[12])
            yrms[fnum] = float(values[13])
            starnum[fnum] = int(values[14])
            onum.append(fnum)
    f.close()

    onum.append('-1')
    raoff['-1'] = '-9999'
    decoff['-1'] = '-9999'
    polagl['-1'] = '-99'

    for palist in ['a00comlist', 'a45comlist', 'a22comlist', 'a67comlist', 'ulist', 'qlist', 'ilist', 'ucomlist', 'qcomlist', 'icomlist']:
        if os.access(palist, os.F_OK):
            os.remove(palist)

    oraoff = '9999'
    odecoff = '9999'
    tfnum = []
    tpang = []
    for v in onum:
        if oraoff != raoff[v] or odecoff != decoff[v]:
            if oraoff != '9999':
                #
                # special treatment is needed for an observation without dithering.                                                                                # software yet to be developed                                                                                                                     #    

                if len(tfnum) > 4:
                    print 'more than 4 frames observed for a cycle of'
                    print tfnum
                    print 'this cycle was skipped'
                else:
                    # a00list, ... , ulist, qlist, ilist are created in chkqui 
                    chkqui(tfnum, tpang)
                tfnum = []
                tpang = []
            oraoff = raoff[v]
            odecoff = decoff[v]

        tfnum.append(v)
        tpang.append(polagl[v])


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

    fi.close()
    fo.close()



def polrecom(band, reject):

    SCRIPTDIR = os.path.realpath(os.path.dirname(sys.argv[0]))
    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()

    #
    # read fits header of the 1st frame
    #
    f = open(band+'tran.txt')
    line = f.readline()
    v1 = line[:-1].split()
    v2 = v1[0].split('.')
    file = '../' + v2[0] + '.ffiles/' + v2[1] + '.fits'
    f.close()
    
    f1 = pyfits.open(file)
    ohdr = f1[0].header
    f1.close()

    comarr = ['i', 'q', 'u', 'a00', 'a45', 'a22', 'a67']
    for v in comarr:
        print v,
        sys.stdout.flush()
        flist = v + 'comlist'

        if reject == 'sigclip':
            ireject = 0
        else:
            if v == 'q' or v == 'u':
                ireject = 2
            else:
                ireject = 1

        cprog = 'recompolx3'
        lnum = commands.getoutput('%s/%s %s %s %s %s %d' % (SCRIPTDIR, cprog, band, objname, v, flist, ireject))
        fcount = pyfits.open('count.fits')
        scidata = fcount[0].data
        pyfits.append(band+v+objname+'.fits', scidata)
        fcount.close()
        os.remove('count.fits')
        fim = pyfits.open(band+v+objname+'.fits', mode='update')
        imhdr = fim[0].header
        imhdr.update('OBJECT', ohdr['OBJECT'])
        imhdr.update('FILTER', ohdr['FILTER'])
        imhdr.update('EXPOS', ohdr['EXPOS'])
        imhdr.update('DATE_UTC', ohdr['DATE_UTC'])
        imhdr.update('TIME_UTC', ohdr['TIME_UTC'])
        imhdr.update('DATE_LT', ohdr['DATE_LT'])
        imhdr.update('TIME_LT', ohdr['TIME_LT'])
        imhdr.update('RA', ohdr['RA'])
        imhdr.update('DEC', ohdr['DEC'])
        imhdr.update('AIRMASS', ohdr['AIRMASS'])
        imhdr.update('COMAREA', comarea)
        imhdr.update('NCOMBINE', int(lnum))
        imhdr.update('COMMETHO', 'average')
        fim.flush()
        fim.close()


    aiquclean()

    os.remove(band+'recominfo')
    print ''


def aiquclean():

    if not os.access('iqu', os.F_OK):
        os.mkdir('iqu')

    for vl in ['a00', 'a45', 'a22', 'a67']:
        if os.access(vl+'comlist', os.F_OK):
            f = open(vl+'comlist')
            for fnum in f:
                os.remove(fnum[:-1]+'.fits')
            f.close()

    for vl in ['icomlist', 'ucomlist', 'qcomlist']:
        if os.access(vl, os.F_OK):
            f = open(vl)
            for fnum in f:
                if re.match('0.', fnum):
                    shutil.move(fnum[:-1]+'.fits', 'iqu/'+fnum[2:-1]+'.fits')
                else:
                    shutil.move(fnum[:-1]+'.fits', 'iqu')
            f.close()


def clean():

    mfkeep = 0
    iqukeep = 0

    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')
            
    if iqukeep == 0:
        shutil.rmtree('iqu')

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

    file=['ilist', 'ulist', 'qlist']
    for v in file:
        if os.access(v, os.F_OK):
            os.remove(v)


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+'polrecom.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 %spolrecom.list' % band
                sys.exit()

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

        mkcomlist(band, urms, lstarnum)

        print 'making i,q,u, images for each cycle ',
        mkqui()
        print 'done'

        print 'combining'
        polrecom(band, reject)
        print 'done'
        print ''

    os.remove(band+'fwhm.txt')
    clean()

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

    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: polrecom.py init band[j|h|k]'
        print 'to make the frame list for polrecom.py recom'
        print ''
        print 'usage: polrecom.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: polrecom.py --help'
        print 'to show this messege'
        print ''
        sys.exit()

    if argvs[1] != 'init' and argvs[1] != 'recom':
        print 'usage: polrecom.py [init|recom] band[j|h|k] -option'
        sys.exit()
    if argvs[2] != 'j' and argvs[2] != 'h' and argvs[2] != 'k':
        print 'usage: polrecom.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] -reject=[minmax] -theta=[float] -ifwhm=[float] -minsep=[float]' 


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


                
