#!/usr/bin/env python

import os
import sys
from pyraf import iraf
from iraf import obsutil
import pyfits
import numpy as np
import math
import re


def getbackstat(infits):

    img = pyfits.open(infits)
    comarea = img[0].header['comarea']
    ncomb = img[0].header['ncombine']
    itime = img[0].header['expos']
    data1 = eval('img[0].data'+comarea)
    img.close()

    zmed = np.median(data1)
    zstd = np.std(data1)

    while True:
        clip = np.where((np.fabs(data1 - zmed) < 3*zstd))
        med = np.median(data1[clip])
        std = np.std(data1[clip])
        if math.fabs(med - zmed) < 1 and math.fabs(zstd - std) < 1:
            break
        zmed = med
        zstd = std

    return med, std, ncomb, itime, comarea



def getfwhm(infits, mlim):
    
    iraf.psfmeasure.display = 'no'
    iraf.psfmeasure.imagecur = 'fwhmcoo.txt'
    iraf.psfmeasure.graphcur = 'q.txt'

    fo = open('q.txt', 'w')
    print >> fo, 'q'
    fo.close()
    
    imgfits = infits+'[0]'

    iraf.daofind.output = 'getfwhm.coo'
    iraf.daofind(imgfits)
    iraf.phot(imgfits, coords='getfwhm.coo', output='getfwhm.mag')
    # merr < mlim is just for excluding merr=INDEF 
    iraf.txdump('getfwhm.mag', fields='xc,yc,mag,merr,cier,sier,pier', expr='merr<'+str(mlim), Stdout='getfwhm.txt') 
    os.remove('getfwhm.mag')
    os.remove('getfwhm.coo')

    # maximum of 30 stars are enough to measure  
    # fwhm and ellipticity of a frame
    # the stars should be isolated, i.e. not contaminated by nearby stars. 
    num = prefwhm('getfwhm.txt')
    os.remove('getfwhm.txt')
    
    if num == 0:
        print 'isolated stars'
        return 999, 999, 0

    iraf.psfmeasure.radius = 2 * iraf.apphot.datapars.fwhmpsf            
    iraf.psfmeasure(imgfits, logfile='psfme.txt', Stderr=1, Stdout=1, StdoutG='/dev/null')

    fwhm, ellip = getpsf()

    os.remove('q.txt')
    os.remove('psfme.txt')
    os.remove('fwhmcoo.txt')

    return fwhm, ellip, num


    
def prefwhm(infile):

    # to remove stars have cier!=0 or sier!=0 or pier!=0  

    txcoo = []
    tycoo = []
    tmag = []
    tmerr = []
    f = open(infile)
    for line in f:
        values = line[:-1].split()
        # 4:cier 5:sier 6:pier 
        if values[4] == '0' and values[5] == '0' and values[6] == '0':
            txcoo.append(float(values[0]))
            tycoo.append(float(values[1]))
            tmag.append(float(values[2]))
            tmerr.append(float(values[3]))
    f.close()

    if len(txcoo) == 0:
        print 'no star with cier!=0 && sier!=0 && pier!=0'
        return 0


    # pick up isolated stars 

    xcoo = []
    ycoo = []
    mag = []
    merr = []
    for i in range(len(txcoo)):
        rmin = 1000
        for j in range(len(txcoo)):
            if i != j:
                rr = math.sqrt((txcoo[j] - txcoo[i]) ** 2 + (tycoo[j] - tycoo[i]) ** 2)
                if rr < rmin:
                    rmin = rr

        if rmin > 10:
            xcoo.append(txcoo[i])
            ycoo.append(tycoo[i])
            mag.append(tmag[i])
            merr.append(tmerr[i])
    
    if len(xcoo) == 0:
        print 'too crowded. fwhm and ellipticity may be unreliable.'
        for i in range(len(txcoo)):
            xcoo.append(txcoo[i])
            ycoo.append(tycoo[i])
            mag.append(tmag[i])
            merr.append(tmerr[i])

    #
    # use the 30 brightest stars for psfmeasure 
    #

    fo = open('fwhmcoo.txt', 'w')

    if len(xcoo) > 30:
        mdict = {}
        for i in range(len(merr)):
            mdict[i] = merr[i]

        psfnum = 0
        for k, v in sorted(mdict.items(), key=lambda x:x[1]):
            psfnum += 1
            if psfnum <= 30:
                print >> fo, xcoo[k], ycoo[k]
            else:
                break
        psfnum = 30
    else:
        psfnum = len(xcoo)
        for i in range(len(xcoo)):
            print >> fo, xcoo[i], ycoo[i]

    fo.close()

    return psfnum



def getpsfstat(data):

    zmed = np.median(data)
    zstd = np.sqrt(np.mean((data - zmed)**2))

    while True:
        clip = np.where((np.fabs(data - zmed) < 3 * zstd))
        med= np.median(data[clip])
        std = np.sqrt(np.mean((data[clip] - med)**2))
        if np.fabs(med - zmed) < 0.01 and np.fabs(std - zstd) < 0.01:
            break
        zmed = med
        zstd = std

    return med, std, clip



def readpsfmeasured(fname):

    flag = 0
    xcoo = []
    ycoo = []
    fwhm = []
    ellip = []
    f = open(fname)
    for line in f:
        if flag == 0:
            if re.match('\s+Image', line):
                flag = 1
        else:
            flag += 1
            if flag == 2:
                v = line[:-1].split()
                xcoo.append(float(v[1]))
                ycoo.append(float(v[2]))
                fwhm.append(v[4])
                ellip.append(v[5])
            else:
                if not re.search('\S', line[:-1]):
                    break
                else:
                    v = line[:-1].split()
                    xcoo.append(float(v[0]))
                    ycoo.append(float(v[1]))
                    fwhm.append(v[3])
                    ellip.append(v[4])
    f.close()

    return xcoo, ycoo, fwhm, ellip



def getpsf():

    dummy1, dummy2, ftmp, etmp = readpsfmeasured('psfme.txt')

    fwhm = []
    ellip = []
    for i in range(len(ftmp)):
        if ftmp[i] != 'INDEF' and etmp[i] != 'INDEF':
            fwhm.append(float(ftmp[i]))
            ellip.append(float(etmp[i]))

    fwhm = np.array(fwhm)
    ellip = np.array(ellip)

    fwhmmed, std, clip = getpsfstat(fwhm)
    ellipc = ellip[clip]
    ellipmed, std, clip = getpsfstat(ellipc)

    return fwhmmed, ellipmed



def fwhmmatch(x, y, xcoo, ycoo, rthresh):

    rmin = 1000
    imin = 0
    for i in range(len(xcoo)):
        r = math.sqrt( (x - xcoo[i])**2 + (y - ycoo[i])**2 )
        if r < rmin:
            rmin = r
            imin = i

    if rmin < rthresh:
        imatch = imin
    else:
        imatch = -1

    return imatch



def joinpsftxt(outtxt, afwhm):
    
    xcoo, ycoo, fwhm, ellip = readpsfmeasured('psfme.txt')
    
    rthresh = 0.5 * afwhm

    fo = open('joinpsf.tmp', 'w')
    f = open(outtxt)
    for line in f:
        v = line[:-1].split()
        x = float(v[0])
        y = float(v[1])
        imatch = fwhmmatch(x, y, xcoo, ycoo, rthresh)
        if imatch == -1:
            mfwhm = 'INDEF'
            mellip = 'INDEF'
        else:
            mfwhm = fwhm[imatch]
            mellip = ellip[imatch]

        print >> fo, '%.1f %.1f %s %s %s %s %s %s %s' % (x, y, v[2], v[3], v[4], v[5], v[6], mfwhm, mellip) 

    f.close()
    fo.close()
    
    os.rename('joinpsf.tmp', outtxt)



def onebyonefwhm(infits, outtxt, fwhm):
    
    imgfits = infits+'[0]'

    iraf.psfmeasure.unlearn()
    iraf.psfmeasure.display = 'no'
    iraf.psfmeasure.radius = 2 * fwhm            
    iraf.psfmeasure.imagecur = 'onepsf.coo'
    iraf.psfmeasure.graphcur = 'q.txt'
    iraf.psfmeasure.iterations = 2

    fo = open('q.txt', 'w')
    print >> fo, 'q'
    fo.close()

    f = open(outtxt)
    for line in f:
        v = line[:-1].split()
        print 'init', v[0], v[1]
        fo = open('onepsf.coo', 'w')
        print >> fo, v[0], v[1]
        fo.close()

        iraf.psfmeasure(imgfits, logfile='psfme.txt', Stderr=1, Stdout=1, StdoutG='/dev/null')

        fin = open('psfme.txt')
        flag = 0
        for line in fin:
            if flag == 0:
                if re.match('\s+Image', line):
                    flag = 1
            else:
                flag += 1
                if flag == 2:
                    v2 = line[:-1].split()
                    print v2[1], v2[2], v2[4], v2[5]
        fin.close()
        os.remove('psfme.txt')

    f.close()



def sirphot(infits, ifwhm, thresh, mlim, pflag, aprad):

    med, stddev, ncomb, itime, comarea = getbackstat(infits)

    iraf.digiphot(_doprint = 0)
    iraf.apphot(_doprint = 0)
    
    iraf.digiphot.unlearn()
    iraf.apphot.unlearn()
    iraf.daofind.unlearn()
    iraf.ptools(_doprint = 0)
    iraf.ptools.unlearn()
    iraf.txdump.unlearn()
    iraf.obsutil()
    iraf.obsutil.unlearn()
    iraf.psfmeasure.unlearn()

    iraf.apphot.datapars.datamax = 20000
    iraf.apphot.datapars.datamin = med - 5 * stddev
    iraf.apphot.datapars.sigma = stddev
    iraf.apphot.datapars.readnoise = 30 * math.sqrt(ncomb)
    iraf.apphot.datapars.epadu = 5 * ncomb
    iraf.apphot.datapars.itime = itime
    iraf.apphot.datapars.fwhmpsf = ifwhm
    iraf.apphot.findpars.threshold = thresh
    iraf.apphot.findpars.sharphi = 0.8
    iraf.apphot.daofind.interac = 'no'
    iraf.apphot.daofind.verify = 'no'
    
    iraf.apphot.phot.interactive = 'no'
    iraf.apphot.phot.verify = 'no'
    iraf.apphot.phot.verbose = 'no'
    if aprad == 0:
        iraf.apphot.photpars.apertures = ifwhm
    else:
        iraf.apphot.photpars.apertures = aprad

    if infits[0] == 'j':
        iraf.apphot.photpars.zmag = 21
    elif infits[0] == 'h':
        iraf.apphot.photpars.zmag = 21.2
    elif infits[0] == 'k':
        iraf.apphot.photpars.zmag = 20.3
    else:
        iraf.apphot.photpars.zmag = 21

    fwhm, ellip, num = getfwhm(infits, mlim)
    print 'fwhm = %.2f ellip = %.2f num = %d' % (fwhm, ellip, num)


    iraf.apphot.datapars.fwhmpsf = fwhm
    if fwhm < 5:
        iraf.apphot.centerpars.cbox = 5
    else:
        iraf.apphot.centerpars.cbox = fwhm
                
    iraf.apphot.fitskypars.annulus = 3 * fwhm
    iraf.apphot.fitskypars.dannulus = 2 * fwhm
    if aprad == 0:
        iraf.apphot.photpars.apertures = fwhm
    else:
        iraf.apphot.photpars.apertures = aprad
    
    # daofind 
    mlim = 0.2
    band = infits[:1]
    imgfits = infits+'[0]'
    outtxt = band+'sirphot.txt'
    iraf.daofind.output = 'sirphot.coo'
    iraf.daofind(imgfits)
    iraf.phot(imgfits, coords='sirphot.coo', output='sirphot.mag')

    if os.access('sirphot.mag', os.F_OK):
        iraf.txdump('sirphot.mag', fields='xc,yc,mag,merr,cier,sier,pier', expr='merr<'+str(mlim), Stdout=outtxt)
        os.remove('sirphot.mag')
    else:
        print 'No stars found'

    os.remove('sirphot.coo')

    if pflag == 1:
        iraf.psfmeasure.radius = 2 * fwhm            
        iraf.psfmeasure.imagecur = outtxt
        iraf.psfmeasure.graphcur = 'q.txt'
        iraf.psfmeasure.iterations = 3

        fo = open('q.txt', 'w')
        print >> fo, 'q'
        fo.close()

        iraf.psfmeasure(imgfits, logfile='psfme.txt', Stderr=1, Stdout=1, StdoutG='/dev/null')
        joinpsftxt(outtxt, fwhm)

        os.remove('q.txt')
        os.remove('psfme.txt')


    return fwhm
    


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

    if argc < 2:
        print 'usage : sirphot.py img.fits'
        print 'option '
        print '-ifwhm=[float]'
        print '-thresh=[float]'
        print '-aprad=[float] aperture radius in pixels'
        print '-nopsf'
        sys.exit()

    if not os.path.exists(argvs[1]):
        print '%s not found' % argvs[1]
        sys.exit()

    ifwhm = 3.0
    thresh = 10
    pflag = 1
    aprad = 0
    mlim = 0.2

    for item in argvs:
        if item.startswith('-ifwhm'):
            ifwhm = float(item[7:])
            print 'ifwhm = %s' % item[7:]
        elif item.startswith('-thresh'):
            thresh = float(item[8:])
            print 'thresh = %s' % item[8:]
        elif item.startswith('-aprad'):
            aprad = float(item[7:])
            print 'aprad = %s' % item[7:]
        elif item.startswith('-nopsf'):
            pflag = 0
            print 'fwhm and ellipticity will not be measured'

    sirphot(argvs[1], ifwhm, thresh, mlim, pflag, aprad)



