#!/usr/bin/env python

import numarray,pyfits
import os,sys,glob
import time
from PyCALib import *

# ************************************************************************
#
# This is the driver to compute the Asymmetry and Central Concentration 
# using the prescription of Abraham et al. (1996).
#
# It needs 4 inputs to run:
#
#  1) a catalog file with the IDs we want to use, on the first column
#  2) the detection catalog from Sextractor
#  3) the flux image where we want to measure A&C, this is usually the 
#     detection image from SExtractor
#  4) the Segmentation image from Sextractor
#
# The program creates <AC> directory where it puts the auxiliar images 
# and writes out an ac.log and and ac.info file with the information
# 
# Felipe Menanteau, JHU, Oct 2003
# Added command line options Jul 2004
#
# ************************************************************************

def main():
    
    # Get the options and arguments
    (opt,arg)  = cmdline()

    # Into local vars
    catids  = arg[0]
    catalog = opt.cat
    fluxima = opt.flux
    segmima = opt.segm
    bpz_cat = opt.bpz
    idcol   = opt.idcol

    # Read the detection catalog in
    #cat = detection_catalog(catalog,opt.verb)
    if opt.bpz:
        cat = apsis_catalogs(opt.dir,opt.verb)

    else:
        cat = detection_catalog(opt.cat)

        
    # Read the bpz information if required
    if opt.Z:
        bpz   = bpz_catalog(bpz_cat,opt.verb)
        (names,waves) = read_image_name(opt.file)

    # Define fitscut
    fitscut = check_exe('fitscut')
    outdir  = "./AC_files/"

    # Make the file repository directory if needed
    if not os.path.isdir(outdir):
        os.mkdir(outdir)
        print >> sys.stderr,' Made new directory '+ outdir
        
    aclog = open(outdir + "ac.log","w")
    acout = open(outdir + "ac.info","w")
    acout.write("# %5s %7s %7s %7s\n" % ("ID","C","A","Araw"))
    aclog.write("# Started on: %s\n" % (time.ctime()))
    
    # Open the catalog and proceed for each object in it
    for line in open(catids).readlines():
        fields = line.split()

        if fields[0][0] == "#":
            continue

        ID = fields[idcol]

        try:
            reblend = fields[3]
            cat.SIZE[ID] = 2.0*cat.SIZE[ID]
        except:
            reblend = None
        
        X  = cat.X_IMAGE[ID]
        Y  = cat.Y_IMAGE[ID]
        id = int(ID)
        key = None
        z   = None
        
        # Define filter according to photometric redshift
        if (opt.Z):

            z       = bpz.Z_B[ID]
            rank    = best_filter(names,waves,z)

            #print "z,rank",z,rank
            key     = rank[0]
            fluxima = names[key]

            # In case it's not detected in that filter
            i   = 0 
            while (cat.MAG[rank[i]][ID] == -99.):
                key = rank[i+1]
                fluxima = names[key]
                print >>sys.stderr, "# Changing %s from %s -> %s" % (ID,rank[i],rank[i+1])
                print >>sys.stderr, "# %s : %8.3f" % (key,cat.MAG[key][ID])
                i = i + 1
                
        # Size of the postage stamp
        if (cat.SIZE[ID] < 120):
            nx = 120
            ny = 120
        else:
            nx = cat.SIZE[ID]
            ny = cat.SIZE[ID]

        # Decide the geometry and the radius type
        if opt.AKron or opt.APetro:
            A_geometry = (cat.A_IMAGE[ID],cat.B_IMAGE[ID],cat.THETA[ID])
            if opt.AKron : A_Rtype = 'Kron'
            if opt.APetro: A_Rtype = 'Petro'

        else:
            A_geometry = None
            A_Rtype    = None

        if opt.CKron or opt.CPetro:
            C_geometry = (cat.A_IMAGE[ID],cat.B_IMAGE[ID],cat.THETA[ID])
            if opt.CKron : C_Rtype = 'Kron'
            if opt.CPetro: C_Rtype = 'Petro'
        else:
            C_geometry = None
            C_Rtype    = None

        # If object is reblended we don't want to use Kron's appertures
        if reblend == "R":
            A_geometry = None
            C_geometry = None
            msg_format = "# Reblended: ignoring geometry for %s "
            print >>sys.stderr, msg_format % ID
            aclog.write((msg_format+"\n") % ID)
            
        # Set the stamps names
        fluxstamp = outdir + ID + ".fits"
        segmstamp = outdir + ID + "_SEGM.fits"

        # Call fitscut on flux image and detection image
        cmd1 = "%s -f -x %s -y %s -r %d -c %d %s %s" % (fitscut,X,Y,nx,ny,fluxima,fluxstamp)
        cmd2 = "%s -f -x %s -y %s -r %d -c %d %s %s" % (fitscut,X,Y,nx,ny,segmima,segmstamp)
        os.system(cmd1)
        os.system(cmd2)
        
        # Run A/C on stamps
        # Concentration
        try:
            c = concentration(id, fluxstamp, segmstamp,
                              C_Rtype,C_geometry,write="yes",
                              plot=opt.plot)
            C = c.C
            if opt.Cds9: c.showds9(verb=0)
        
        except:
            print "# Error: Problem with calculation of concentration on obj: ",id
            aclog.write("# Error: Problem with calculation of concentration on obj: %s\n" % id)
            C = 999.999
            
        # Asymmetry
        try:
            a = asymmetry(id, fluxstamp, segmstamp,
                          A_Rtype, A_geometry,write="yes",
                          skyrms=cat.RMS[ID],smooth="yes",plot=opt.plot)
            A    = a.A
            Araw = a.Araw
            if opt.Ads9: a.showds9(verb=0)
            
        except:
            print "# Error: Problem with calculation of asymmetry on obj: ",ID
            aclog.write("# Error: Problem with calculation of asymmetry on obj: %s\n" % id)
            A    = 99.999
            Araw = 99.999
            
        format = "%7s %7.3f %7.3f %7.3f"
        outtup = (ID, C, A, Araw)

        acout.write( (format+"\n") % outtup)
        print  format  % outtup

        if a.msg: aclog.write(a.msg)
        if c.msg: aclog.write(c.msg)            

        # Clean up files if required
        if opt.clean:
            clean_up(ID,outdir)

    # Write to the log and bye
    aclog.write("# Ended on: %s\n" % (time.ctime()))





# ************************************
#    Parse the command line options
# ************************************
def cmdline():
    from optik import OptionParser

    usage = "usage: %prog <catids> [options] "
    parser = OptionParser(usage=usage)


    parser.add_option("-i","--idcol", type="int",
                      default=0,
                      dest="idcol", help="Column number with id")
    
    parser.add_option("-d","--dir", type="string",
                      default="",
                      dest="dir", help="base directory")

    parser.add_option("-c", "--cat",type="string",
                      dest="cat", help="input detection catalog")

    parser.add_option("-z", "--bpz",type="string",
                      dest="bpz", help="input detection catalog")

    parser.add_option("-f", "--flux",type="string",
                      dest="flux", help="input flux image")
    
    parser.add_option("-s", "--segm",type="string",
                      dest="segm", help="input segmentation image")
    
    parser.add_option("-Z", "--redshift",
                      action="store_true", dest="Z", default=0,
                      help="use the redshift")

    parser.add_option("-F", "--File",type="string",
                      dest="file", help="input flux file (images,wave)")

    parser.add_option("-P", "--Petro",
                      action="store_true", dest="Petro", default=0,
                      help="Use Petrosian apperture for both A and C")

    parser.add_option("--APetro",
                      action="store_true", dest="APetro", default=0,
                      help="Use Petrosian apperture for Asymmetry")

    parser.add_option("--CPetro",
                      action="store_true", dest="CPetro", default=0,
                      help="Use Petrosian apperture for Concentration")

    parser.add_option("-K", "--Kron",
                      action="store_true", dest="Kron", default=0,
                      help="Use Kron apperture for both A and C")

    parser.add_option("--AKron",
                      action="store_true", dest="AKron", default=0,
                      help="Use Kron apperture for Asymmetry")

    parser.add_option("--CKron",
                      action="store_true", dest="CKron", default=0,
                      help="Use Kron apperture for Concentration")
    
    parser.add_option("-v", "--verb",
                      action="store_true", dest="verb", default=1,
                      help="verbose")

    parser.add_option("-q", "--quiet",
                      action="store_false", dest="verb", default=1,
                      help="verbose quiet")

    parser.add_option("--ds9",
                      action="store_true", dest="ds9", default=0,
                      help="Show products in DS9 both A and C")

    parser.add_option("--Ads9",
                      action="store_true", dest="Ads9", default=0,
                      help="Show products in DS9 for Asymmetry")

    parser.add_option("--Cds9",
                      action="store_true", dest="Cds9", default=0,
                      help="Show products in DS9 for Concentration")

    parser.add_option("--plot",
                      action="store_true", dest="plot", default=0,
                      help="Plot eta(r) in pgplot")

    parser.add_option("--clean",
                      action="store_true", dest="clean", default=0,
                      help="Clean up aux files")

    
    (options, args) = parser.parse_args()

    if len(args) < 1:
        parser.error("Must supply at least one argument required")

    # Set the default values in case are not defined
    if not options.cat:
        options.cat = os.path.join(options.dir,"Catalogs/detectionImage.cat")

    if not options.bpz and options.Z:
        options.bpz = os.path.join(options.dir,"Catalogs/bpz.cat")

    if not options.segm:
        options.segm = os.path.join(options.dir,"Images/detectionImage_SEGM.fits")

    if not options.flux and not options.file :
        options.flux = os.path.join(options.dir,"Images/detectionImage.fits")

    if options.Kron :
        print "setting Kron's to 1"
        options.AKron=1
        options.CKron=1

    if options.Petro :
        print "setting Petrosian's to 1"
        options.APetro=1
        options.CPetro=1
        
    if options.ds9:
        options.Ads9=1
        options.Cds9=1

    print >>sys.stderr," Input files in use "
    print >>sys.stderr,  "\tDetection catalog   : %s" % options.cat
    print >>sys.stderr,  "\tPhoto-z   catalog   : %s" % options.bpz
    print >>sys.stderr,  "\tFlux image          : %s" % options.flux
    print >>sys.stderr,  "\tImages file         : %s" % options.file
    print >>sys.stderr,  "\tSegmentation image  : %s" % options.segm
    print >>sys.stderr," Options  "
    print >>sys.stderr,  "\tAsymmentry Kron     : %s" % options.AKron
    print >>sys.stderr,  "\tConcentration Kron  : %s" % options.CKron
    print >>sys.stderr,  "\tAsymmentry Petro    : %s" % options.APetro
    print >>sys.stderr,  "\tConcentration Petro : %s" % options.CPetro
    print >>sys.stderr,  "\tShow C in ds9       : %s" % options.Cds9
    print >>sys.stderr,  "\tShow A in ds9       : %s" % options.Ads9
    print >>sys.stderr,  "\tPlot eta(r) w/pgplot: %s" % options.plot
    print >>sys.stderr,  "\tClean up aux files  : %s" % options.clean
    print >>sys.stderr,  "\tverbose             : %s" % options.verb
    return options,args

def read_image_name(file):

    print >>sys.stderr, "\tOpening file names :",file

    names = {}
    waves = {}

    for line in open(file).readlines():
        
        fields = line.split()
        if fields[0][0] == "#":
            continue

        filter        = fields[0]
        names[filter] = fields[1]
        waves[filter] = float(fields[2])
        
    return names,waves

def best_filter(names,waves,z,lo=4450.):

    # Get the close band to B-band rest-frame
    lw  = lo*(1. + z)
    # and we find the closest
    wo   = waves.values()
    keys = waves.keys()
    dist = abs (lw - numarray.asarray(wo))

    rank = []
    for i in range(len(wo)):
        dmin = numarray.where( dist == numarray.sort(dist)[i])
        idx  = dmin[0][0]
        key  = keys[idx]
        rank.append(key)
        
    return rank

def check_exe(exe):

    path = os.environ['PATH'].split(':')
    for p in path:

        f = os.path.join(p,exe)
        if os.path.isfile(f):
            print >>sys.stderr,"# Found %s in %s" % (exe,f)
            return f
    print sys.stderr,"# ERROR: Couldn't find %s" % exe
    return 

def clean_up(ID,outdir):
    ''' Cleans up the aux files'''

    files = glob.glob(os.path.join(outdir,ID)+"*.fits")

    print >>sys.stderr,"Cleaning up files"
    for file in files:
        os.remove(file)

    return

main()