#!/usr/bin/env python

# 2014-03-14 sexa2deg was separated from opm

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


def fitshead(infits):

    img = pyfits.open(infits)
    nax1 = img[0].header['NAXIS1']
    nax2 = img[0].header['NAXIS2']
    ra = img[0].header['RA']
    dec = img[0].header['DEC']
    img.close()

    size = str(math.floor(min(nax1, nax2) * 0.45))

    return size, ra, dec



def get2mass(vizier, size, ra, dec):

    if vizier == 'fr':
        vizurl = 'vizier.u-strasbg.fr'
    elif vizier == 'jp':
        vizurl = 'vizier.nao.ac.jp'
    elif vizier == 'uk':
        vizurl = 'vizier.ast.cam.ac.uk'
    elif vizier == 'cfa':
        vizurl = 'vizier.cfa.harvard.edu'
    elif vizier == 'hawaii':
        vizurl = 'www.ukirt.jach.hawaii.edu'
    elif vizier == 'canada':
        vizurl = 'vizier.hia.nrc.ca'
    elif vizier == 'china':
        vizurl = 'vizier.china-vo.org'
    elif vizier == 'india':
        vizurl = 'vizier.iucaa.ernet.in'
    elif vizier == 'russia':
        vizurl = 'vizier.inasan.ru'

    url = '"http://' + vizurl + '/viz-bin/asu-acl?-source=2MASS&-c.ra=' + ra + '&-c.dec=' + dec + '&-c.bs=' + size + '"'

    print 'Now retrieving and downloading 2mass catalog data'
    os.system('curl -# -o 2mass.out -O %s' % url)
    print ''

    f = open('2mass.out')
    lines = f.read()
    f.close()
    if lines.find('404 Not Found') != -1:
        print '%s is not alive' % vizurl
        sys.exit()



def edit2mass(band):

    if band == 'j':
        i = 0
    elif band == 'h':
        i = 1
    elif band == 'k':
        i = 2

    magmax = 0
    fout = open('2mass.tmp', 'w')
    fin = open('2mass.out')
    for line in fin:
        if line[:1] != '#' and line[:1] != '-' and line[:1] != ' ' and line[:5] != '2MASS' and len(line) > 10:
            v = line[:-1].split()
            qflag = line[82+i:83+i]
            #blflag = line[90+i:91+i]
            ccflag = line[94+i:95+i]
            mag = float(line[40+14*i:46+14*i])

            if qflag == 'A' and ccflag == '0':
                print >> fout, '%s %s %.3f' %  (v[1], v[2], mag)
                if magmax < mag:
                    magmax = mag
    fin.close()
    fout.close()

    exec("os.system('sort -n -k 3 2mass.tmp > good2mass.tmp')")
    os.remove('2mass.tmp')

    return magmax



def editsir(band, tmmax):

    magmin = 20
    magmax = 5

    xcoo = []
    ycoo = []
    mag = []
    merr = []
    fwhm = []
    ellip = []
    f = open(band+'sirphot.txt')
    for line in f:
        if line[:1] != '#':
            v = line[:-1].split()
            srmag = float(v[2])
            if srmag < tmmax and float(v[3]) < 0.1 and v[4] == '0' and v[5] == '0' and v[6] == '0' and v[7] != 'INDEF':
                if srmag > magmax:
                    magmax = srmag
                if srmag < magmin:
                    magmin = srmag

                xcoo.append(v[0])
                ycoo.append(v[1])
                mag.append(v[2])
                merr.append(v[3])
                fwhm.append(float(v[7]))
                ellip.append(v[8])
    f.close()

    fwhm = np.array(fwhm)
    med, std, clip = sirphot.getpsfstat(fwhm)

    fo = open('sirtmp.list', 'w')
    for i in clip[0]:
        print >> fo, xcoo[i], ycoo[i], mag[i]
    fo.close()

    exec("os.system('sort -n -k 3 sirtmp.list > inlist')")
    os.remove('sirtmp.list')

    return magmin, magmax



def trim2mass(magmin, magmax):

    fo = open('2mass.list', 'w')
    f = open('good2mass.tmp')
    for line in f:
        v = line[:-1].split()
        mag = float(v[2])
        if mag > magmin and mag < magmax:
            print >> fo, line[:-1]
    f.close()
    fo.close()
    os.remove('good2mass.tmp')



def join4ccmap():
    
    ra = []
    dec = []
    tmag = []
    f = open('2mass.list')
    for line in f:
        v = line[:-1].split()
        ra.append(v[0])
        dec.append(v[1])
        tmag.append(v[2])
    f.close()

    fo = open('ccmap.coo', 'w')
    f = open('opmout')
    for line in f:
        if not line.startswith('#'):
            v = line[:-1].split()
            id = int(v[6]) - 1
            if v[9] != tmag[id]:
                print 'matching is wrong in sirwcs.join4ccmap'
            else:
                print >> fo, v[1], v[2], ra[id], dec[id]
    f.close()
    fo.close()

    os.remove('2mass.list')



def ccmap(infits):

    band = infits[0]

    iraf.images()
    iraf.imcoords()
    iraf.ccmap.unlearn()
    iraf.ccsetwcs.unlearn()

    iraf.ccmap.lngunits = 'degrees'
    iraf.ccmap.latunits = 'degrees'
    iraf.ccmap.maxiter = 5
    iraf.ccmap.verbose = 'yes'
    iraf.ccmap.interactive = 'no'
    
    iraf.ccmap('ccmap.coo', band+'ccmap.db')
    # iraf.ccmap.update = 'yes' modifies band+'ccmap.db' and it makes a trouble when cctran
    # use ccsetwcs instead 

    iraf.ccsetwcs(infits+'[0]', band+'ccmap.db', 'ccmap.coo')

    os.remove('ccmap.coo')



def sexa2deg(ra, dec):

    hh, mm, ss = ra.split(':')
    degra = 15 * float(hh) + float(mm) / 4 + float(ss) / 240
    dd, dm, ds = dec.split(':')
    if dec[:1] == '-':
        degdec = float(dd) - float(dm) / 60 - float(ds) / 3600
    else:
        degdec = float(dd) + float(dm) / 60 + float(ds) / 3600

    return degra, degdec



def opm(infits, ra, dec, halfsize, srmin, ratio):

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

    degra, degdec = sexa2deg(ra, dec)

    os.system("%s/OPM/get_default_param param.tmp" % SCRIPTDIR)

    fo = open('param2', 'w')
    f = open('param.tmp')
    for line in f:
        v = line[:-1].split()
        if v[0] == 'MIRROR':
            print >> fo, '%-10s ON' % v[0]
        elif v[0] == 'MAGNIFY':
            print >> fo, '%-10s 0.45' % v[0]
        elif v[0] == 'HALFSIZE':
            print >> fo, '%-10s %.1f' % (v[0], halfsize)
        elif v[0] == 'MAG_CUT2':
            print >> fo, '%-10s %.2f' % (v[0], srmin)
        elif v[0] == 'NITER':
            print >> fo, '%-10s 10' % v[0]
        elif v[0] == 'MATCHING':
            print >> fo, '%-10s %f' % (v[0], ratio)
        else:
            print >> fo, line[:-1]
    f.close()
    fo.close()

    os.system('%s/OPM/local_coordinate_degree 2mass.list reflist %f %f %f' % (SCRIPTDIR, degra, degdec, halfsize))
    os.system('%s/OPM/opm inlist reflist opmout param2' % SCRIPTDIR)

    os.remove('param.tmp')
    os.remove('param2')
    os.remove('reflist')
    os.remove('inlist')



def calib(band):
    
    dmag = []
    f = open('opmout')
    for line in f:
        if not line.startswith('#'):
            v = line[:-1].split()
            dmag.append(float(v[5]) - float(v[9]))
    f.close()

    dmag = np.array(dmag)
    med, std, dummy = sirphot.getpsfstat(dmag)
    
    print 'mag(sirius - 2mass) = %.2f (%.2f)' % (med, std)

    fo = open('sirphotwcs.tmp', 'w')
    f = open(band+'sirphot.txt')
    for line in f:
        v = line[:-1].split()
        mag = float(v[2]) - med
        print >> fo, '%s %s %s %s %.3f %s %s %s %s %s %s' % (v[0], v[1], v[0], v[1], mag, v[3], v[4], v[5], v[6], v[7], v[8])
    f.close()
    fo.close()

    # cctran 

    iraf.images()
    iraf.imcoords()
    iraf.cctran.unlearn()
    iraf.cctran.xcolumn = 3
    iraf.cctran.ycolumn = 4
    iraf.cctran.lngformat = '%0.2H'
    iraf.cctran.latformat = '%0.1h'
    
    iraf.cctran('sirphotwcs.tmp', band+'sirphotwcs.txt', band+'ccmap.db', 'ccmap.coo')

    os.remove('sirphotwcs.tmp')
    os.rename('opmout', band+'opmout')



def help():

    print ''
    print 'usage: sirwcs.py filename.fits'
    print 'e.g., sirwcs.py jM42.fits'
    print ''
    print 'options:'
    print '-vizier=[fr|jp|uk|cfa|hawaii|canada|china|india|russia] default : fr'
    print 'select a mirror site of vizier from which sirwcs get 2mass catalog.'
    print ''
    print '-noget2mass'
    print 'if you already have 2mass.out in the working directory, '
    print 'you can skip downloading 2mass catalog'
    print ''
    print '-jhk'
    print 'all the j, h, and k bands frames will be processed.'
    print 'e.g., when you have jM42.fits, hM42.fits and kM42.fits'
    print 'and if you want to process all the frames'
    print 'you can command either'
    print 'sirwcs.py jM42.fits -jhk, or sirwcs.py hM42.fits -jhk, or sirwcs.py kM42.fits -jhk'
    print ''
    print '-thresh=[float] default : 10'
    print 'daofind.sigma to detect stars in the SIRIUS frame'
    print ''
    print '-ifwhm=[float] default : 3'
    print 'an initial guess to estimate fwhm'
    print ''
    print '-aprad=[float] default : fwhm'
    print 'an aperture radius in pixels'
    print ''



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

    if argc < 2 or re.search('help', argvs[1]):
        help()
        sys.exit()

    if not os.path.exists(argvs[1]):
        print '%s not found' % argvs[1]
        sys.exit()
 
    infits = argvs[1]
    band = infits[0]
    if band != 'j' and band != 'h' and band != 'k':
        print 'input a correct fits file name which starts with j or h or k.'
        sys.exit()

    vizier = 'fr'
    ng2flag = 0
    thresh = 10
    ifwhm = 3.0
    ratio = 0.5

    # for sirphot
    mlim = 0.15
    pflag = 1
    aprad = 0

    if argc > 2:
        for item in argvs[2:]:

            if item.startswith('-vizier'):
                vizier = item[8:]
                print 'vizier = %s' % vizier
            elif item.startswith('-jhk'):
                band = 'all'
                print 'band = jhk'
            elif item.startswith('-noget2mass'):
                ng2flag = 1
                print 'noget2mass'
            elif 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:]
            else:
                print 'typo? -vizier=[mirror] -jhk -noget2mass -thresh=[float] -ifwhm=[float] -aprad=[thresh]'
                sys.exit()

    
    if ng2flag == 1 and not os.path.exists('2mass.out'):
        print 'No 2mass.out found'
        print '-noget2mass was ignored'
        ng2flag = 0

    if band == 'all':
        bandarr = ['j', 'h', 'k']
    else:
        bandarr = ['']
        bandarr[0] = band

    sizeflag = 0
    for band in bandarr:
        
        infits = band + argvs[1][1:]

        ifwhm = sirphot.sirphot(infits, ifwhm, thresh, mlim, pflag, aprad)
        # output : band+'sirphot.txt' 

        if sizeflag == 0:
            size, ra, dec = fitshead(infits)
            sizeflag = 1

        if ng2flag == 0:
            get2mass(vizier, size, ra, dec)
            ng2flag = 1    

        tmmax = edit2mass(band)

        srmin, srmax = editsir(band, tmmax)

        trim2mass(srmin, srmax)

        halfsize = 0.5 * float(size)

        opm(infits, ra, dec, halfsize, srmin, ratio)

        join4ccmap()

        ccmap(infits)
        
        calib(band)


