#!/usr/bin/env python
# A script to read in a csv file with positions and extents and use that to
# search for counterparts using tables from HEASARC
# A. Ptak, Mar. 2002

import saoimage_reg
import xa_util
import sys, os
import heasarc, CSV_ext, radec
from radec_id import radec_id
import radec
from math import sqrt

# columns here are HEASARC catalog name, short cat. name, cat. description,
# racol, deccol, idcol, extra data columns...
catdata_local = [["usnoa2", "usno", "USNO-A2 stars", "RA(ICRS)", "DE(ICRS)",
            "USNO-A2.0", "Rmag", "Bmag", "Mflag"],
           ["usnoa2", "usno", "USNO-A2 stars, old HEASARC format", "ra", "dec",
            "_unique_id", "rmag", "bmag", "magflag"],
           ["veron2001", "veron", "Veron 2001 QSOs", "ra", "dec",
            "none", "redshift", "app_mag", "object_type"],
           ["zcat", "zcat", "CfA Redshift Catalog", "ra", "dec",
            "none", "bmag", "radial_velocity", "radial_velocity_error",
            "redshift", "class"],
           ["rc3", "rc3", "RC3 Galaxy Catalog", "ra", "dec",
            "name", "bmag", "type", "_3rd_alt_name"],
           ["first", "first", "FIRST 20cm", "ra", "dec",
            "name", "flux_20_cm", "flux_20_cm_error", "int_flux_20_cm"]]
catdata_viz = [
           ["I/284", "usno-b", "USNO-B1 stars", "RAJ2000", "DEJ2000",
            "USNO-B1.0", "R1mag", "B1mag", "Imag", "R2mag", "B2mag"],
           ["I/271", "gsc2.2", "GSC2 stars", "RA(ICRS)", "DE(ICRS)",
            "none", "Rmag", "Bjmag", "Class", "a", "e"],
           ["B/2mass", "2mass", "2mass", "RAJ2000", "DECJ2000",
            "none", "Jmag", "Hmag", "Kmag", "extd_flg", "e_Jmag"]
           ]

class match_data:
    def __init__(self, init_data=None):
        self.pos1 = radec.radec()
        self.fldname = "none"
        self.telname = "none"
        self.detname = "none"
        self.flux = 0.
        self.sigmaj = 0.
        self.extflag = "F"
        self.id = -1
        self.pos2 = radec.radec()
        self.catalog = "none"
        self.catid = "none"
        #self.ds9_color = "green"
        self.extra_data = {} # additional data in the table
        self.id_name = "none"
        self.ra_name = "none"
        self.dec_name = "none"
        if init_data != None:
            self.setup(init_data)

    def __repr__(self):
        s = "%s %d: %s %.2g %.2f\" %s %s %.1f\" " % (self.fldname,
                                                               self.id,
                                                               self.pos1,
                                                               self.flux,
                                                               self.sigmaj,
                                                               self.catalog,
                                                               self.pos2,
                                                               d)
        k = self.extra_data.keys()
        k.sort()
        for key in k:
            # Works for number or strings
            s = s + " %s" % self.extra_data[key]
            
        return(s)

    def extract_from_hdb(self, data):
        if data.has_key(self.ra_name) and data.has_key(self.dec_name):
            self.pos2 = radec.radec(data[self.ra_name] + ' ' +
                                    data[self.dec_name])
        if self.id_name != "none":
            if not data.has_key(self.id_name):
                print "Error, " + self.id_name + " not found in data"
                return(1)
            else:
                self.catid = data[self.id_name]
        for key in self.extra_data.keys():
            if data.has_key(key):
                val = data[key]
                try:
                    fval = float(val)
                    val = fval
                except:
                    pass
                self.extra_data[key] = val
        return(0)
    
    def set_xray_data(self, ra, dec, flux, sigmaj, id, fldname, telname,
                      detname, extflag):
        self.pos1.ra = ra
        self.pos1.dec = dec
        self.flux = flux
        self.sigmaj = sigmaj
        self.id = id
        self.fldname = fldname
        self.detname = detname
        self.telname = telname
        self.extflag = extflag

    def setup(self, data):
        self.heasarc_name = data[0]
        self.catalog = data[1]
        self.desc = data[2]
        self.ra_name = data[3]
        self.dec_name = data[4]
        self.id_name = data[5]
        for val in data[6:]:
            self.extra_data[val] = 0.
        
## class match_data_usno(match_data):
##     def __init__(self):
##         match_data.__init__(self)
##         self.catalog = "usno"
##         self.catid = -1
##         self.extra_data["rmag"] = 0.
##         self.extra_data["bmag"] = 0.
##         #self.ds9_color = "cyan"

##     def __repr__(self):
##         s = match_data.__repr__(self) + " %.2f %.2f" % (\
##             self.extra_data["rmag"], self.extra_data["bmag"])
##         return(s)

##     def extract_from_hdb(self, data):
##         match_data.extract_from_hdb(self, data)        
##         self.catid = int(data[])

## class match_data_2mass(match_data):
##     def __init__(self):
##         match_data.__init__(self)
##         self.catalog = "2mass"
##         self.ds9_color = "magenta"
##         self.extra_data["Jmag"] = 0.
##         self.extra_data["Hmag"] = 0.
##         self.extra_data["Kmag"] = 0.

##     #def __repr__(self):
##     #    s = match_data.__repr__() + " %.2f %.2f" % (self.Rmag,
##     #                                                self.Bmag)
##     #    return(s)

##     def extract_from_hdb(self, data):
##         match_data.extract_from_hdb(self, data)
##         self.pos2.ra = float(data["RAJ2000"])
##         self.pos2.dec = float(data["DEJ2000"])
##         self.catid = radec_id("2mass", self.pos2)
        
def gen_match_data(table):
    """This function was a generator function to instantiate a match_data
    subclass, now there is only a single match_data class but this fn.
    instantiates it with the proper setup info for a given table.  Very
    inefficient at the moment"""
##     if table == "usno":
##         return(match_data_usno())
##     if table == "2mass":
##         return(match_data_2mass())
    for data in catdata:
        if data[1] == table:
            return(match_data(data))
    return(match_data())

argv = sys.argv
argc = len(argv)
if argc < 3:
    print "Usage: find_counterparts.py csvfile outfile (table=table regfile=regfile "
    print"            tableroot=tableroot make_src_tables= aspunc=0. maxtime=600. tableregionsize=2. ndatacols=5 query_vizier=0 verbose=1)"
    print "      or: find_counterparts.py tableroot ra=ra dec=dec radius=radius (table=table"
    print "           maxtime=600. tableregionsize=2. ndatacols=5 verbose=1)"
    print "Reads source list csvfile generated by XAssist and searches"
    print "  tables from HEASARC (listed below) for matches, saved in the form of a new "
    print "  comma-seperated values (csv) file and optionally also a ds9 region file"
    print "Alternate usage is to just download the catalog sources within a given radius "
    print "  (in arcmin.) of a position."
    print "table = name of HEASARC table to query.  Table names are:"
    print "  Tables local to HEASARC:"    
    for cat in catdata_local:
        print "        %20s: %10s" % (cat[2], cat[0])
    print "  Tables queried at Vizier through HEASARC:"    
    for cat in catdata_viz:
        print "        %20s: %10s" % (cat[2], cat[0])
    print "if no table is specified, all of the HEASARC \"local\" tables will be queried"
    print "  (i.e., the ones hosted at the HEASARC), and the Vizier ones will also be"
    print "  queried if query_vizier=1;  Note that the connection to HEASARC can hang"
    print "  when querying Vizier)."
    print "tableroot = root file name for the HEASARC table"
    print "If tableroot+catname.dat exists, it won't be re-created"
    print "If make_src_tables=yes, a separate csv output file will be created for each "
    print "  (X-ray) source, with filename basename(outfile) + _srcX.csv, basename(outfile)"
    print "  = outfile up to last '.'"
    print "tableregionsize = 2., if > 0 then ds9 region files are created with "
    print "  names tableroot+catname.reg, containing a circle region for each table"
    print "  object and radius=tableregionsize in arcsecs."
    print "Set ascunc to uncertainty for positions in csvfile (in arcsecs.)"
    print "maxtime = maximum time to wait for a response from HEASARC, in secs."
    print "ndatacols = number of columns in output csv file for data returned from catalog"
    print ""
    print "Example:"
    print "find_counterparts.py testy ra=\"10 01 57.8\" dec=\"+55 40 47\" radius=5. table=\"usnoa2\""
    sys.exit()

tableroot = "temp"
keywds = xa_util.get_keywords(argv[1:])
csvfile = None
if keywds.has_key("ra"):
    tableroot = argv[1]
else:
    csvfile = argv[1]
    outfile = argv[2]
#print keywds
#sys.exit()
        
regfile = None
aspunc = 0.
hdbcom = "browse_extract.pl"
maxtime = 1800.
verb=1
if keywds.has_key("verbose"):
    verb = int(keywds["verbose"])
    if verb > 1:
        print "verbose set to ", verb
if keywds.has_key("regfile"):
    regfile = keywds["regfile"]
    if verb > 1: print "Set regfile to " + regfile
if keywds.has_key("tableroot"):
    tableroot = keywds["tableroot"]
if keywds.has_key("hdbcom"):
    hdbcom = keywds["hdbcom"]
if keywds.has_key("maxtime"):
    maxtime = float(keywds["maxtime"])
table = None
if keywds.has_key("table"):
    table = keywds["table"]
tableregionsize = 2.
if keywds.has_key("tableregionsize"):
    tableregionsize = float(keywds["tableregionsize"])
ndatacols = 5
if keywds.has_key("ndatacols"):
    ndatacols = int(keywds["ndatacols"])

query_vizier = 0
if keywds.has_key("query_vizier"):
    query_vizier = int(keywds["query_vizier"])

# Following only relevant if not reading sources
if not csvfile:
    ra = keywds["ra"]
    dec = keywds["dec"]
    rad = float(keywds["radius"])
    str = "%s %s" % (ra, dec)
    print "radec = " + str
    try:
        posn = radec.radec(str)
    except:
        print str + " does not appear to be a valid position"
        sys.exit()
    catdata = catdata_local
    if query_vizier:
        catdata += catdata_viz
    for data in catdata:
        if table:
            if data[0] != table:
                continue
        if verb >=1: print "Querying %s..." % data[2]
        tab = heasarc.heasarc(data[0], "%f, %f" % (posn.ra, posn.dec),
                              radius=rad,
                              outfile=tableroot+"_%s.dat" % data[1],
                              hdbcom=hdbcom, maxtime=maxtime,
                              verbose=verb)
        nmat = len(tab.recs)
        if verb >=1: print "Found %d records" % nmat
        if tableregionsize > 0. and len(tab.recs) > 0:
            tabregfile = tableroot + "_%s.reg" % data[1]
            reg = saoimage_reg.ds9_reg()
            for rec in tab.recs:
                tabrec = gen_match_data(data[1])
                if tabrec.extract_from_hdb(rec):
                    break
                reg.add_object("circle", tabrec.pos2.ra,
                               tabrec.pos2.dec, [tableregionsize],
                               radec=1, text="%s" % (data[1] + "-" + \
                                                     tabrec.catid),
                               color="blue")
            reg.Save(tabregfile)
        sys.stdout.flush()
    sys.exit()
    
sepsrcs = 0
if keywds.has_key("make_src_tables") and keywds["make_src_tables"] == "yes":
    sepsrcs = 1
if keywds.has_key("aspunc"):
    aspunc = float(keywds["aspunc"])
if verb >=1: print "Reading source parameters..."
csv = CSV_ext.CSV_ext(csvfile)
if verb >=1: print "Read %d source parameter lines" % len(csv.csv)
racol = csv.FindColNum("ra")
deccol = csv.FindColNum("dec")
fluxcol = csv.FindColNum("flux")
sigmajcol = csv.FindColNum("sigmaj")
fldnamecol = csv.FindColNum("fldname")
telnamecol = csv.FindColNum("telesc")
detnamecol = csv.FindColNum("instr")
idcol = csv.FindColNum("idnum")
extcol = csv.FindColNum("ext")
if verb > 1:
    print "sigmajcol = ", sigmajcol

# Make list of unique flds
if verb >=2: print "Making list of unique fields..."
class field_data:
    def __init__(self, ra0, dec0, ra1, dec1):
        self.ra0 = ra0
        self.dec0 = dec0
        self.ra1 = ra1
        self.dec1 = dec1
        self.tables = {}

fldlist = {}
for data in csv.csv:
    ra = float(data[racol])
    dec = float(data[deccol])
    fldname = data[fldnamecol]
    sigmaj = data[sigmajcol]
    if sigmaj == "":
        sigmaj = 0.1
    else:
        sigmaj = float(sigmaj)/3600.
    if not fldlist.has_key(fldname):
        r = field_data(ra-sigmaj, dec-sigmaj, ra+sigmaj, dec+sigmaj)
        fldlist[fldname] = r
    else:
        if fldlist[fldname].ra0 > ra - sigmaj:
            fldlist[fldname].ra0 = ra - sigmaj
        if fldlist[fldname].dec0 > dec - sigmaj:
            fldlist[fldname].dec0 = dec - sigmaj
        if fldlist[fldname].ra1 < ra + sigmaj:
            fldlist[fldname].ra1 = ra + sigmaj
        if fldlist[fldname].dec1 < dec + sigmaj:
            fldlist[fldname].dec1 = dec + sigmaj

if verb >=2: print "Found %d unique fields" % len(fldlist)

# Now get usno stars for each field
for fldname in fldlist.keys():
    r = fldlist[fldname]
    p1 = radec.radec()
    p1.ra = r.ra0
    p1.dec = r.dec0
    p2 = radec.radec()
    p2.ra = r.ra1
    p2.dec = r.dec1
    if verb >=1: print "RA/DEC range for %s = %s - %s" % (fldname, p1, p2)
    ra = 0.5*(r.ra0 + r.ra1)
    dec = 0.5*(r.dec0 + r.dec1)
    rad = p1.Dist(p2)*60.

    catalogs = []
    old = 0
    if old == 1:
        if verb >=1: print "Querying USNO stars..."
        usno = heasarc.heasarc("usnoa2", "%f, %f" % (ra, dec),
                               radius=rad, outfile=tableroot+"_usno.dat",
                               hdbcom=hdbcom, maxtime=maxtime,verbose=verb)
        nmat = len(usno.recs)
        if verb >=1: print "Found %d records" % nmat
        fldlist[fldname].tables["usno"] = usno
        if verb >=1: print "Querying 2mass..."
        k = heasarc.heasarc("B/2mass", "%f, %f" % (ra, dec),
                            radius=rad, outfile=tableroot+"_2mass.dat",
                            hdbcom=hdbcom, maxtime=maxtime,verbose=verb)
        nmat = len(k.recs)
        if verb >=1: print "Found %d records" % nmat
        fldlist[fldname].tables["2mass"] = k

        if verb >=1: print "Querying Veron2001..."
        v = heasarc.heasarc("veron2001", "%f, %f" % (ra, dec),
                            radius=rad, outfile=tableroot+"_veron.dat",
                            hdbcom=hdbcom, maxtime=maxtime,verbose=verb)
        nmat = len(v.recs)
        if verb >=1: print "Found %d records" % nmat
        fldlist[fldname].tables["veron"] = v

        if verb >=1: print "Querying CfA Redshift Catalog..."
        z = heasarc.heasarc("zcat", "%f, %f" % (ra, dec),
                            radius=rad, outfile=tableroot+"_zcat.dat",
                            hdbcom=hdbcom, maxtime=maxtime,verbose=verb)
        nmat = len(z.recs)
        if verb >=1: print "Found %d records" % nmat
        fldlist[fldname].tables["zcat"] = z

        if verb >=1: print "Querying RC3..."
        rc3 = heasarc.heasarc("rc3", "%f, %f" % (ra, dec),
                              radius=rad, outfile=tableroot+"_rc3.dat",
                              hdbcom=hdbcom, maxtime=maxtime,verbose=verb)
        nmat = len(rc3.recs)
        if verb >=1: print "Found %d records" % nmat
        fldlist[fldname].tables["rc3"] = rc3
    else:
        catdata = catdata_local
        if query_vizier:
            catdata += catdata_viz
        for data in catdata:
            if table:
                if data[0] != table:
                    continue
            if verb >=1: print "Querying %s..." % data[2]
            tab = heasarc.heasarc(data[0], "%f, %f" % (ra, dec),
                                  radius=rad,
                                  outfile=tableroot+"_%s.dat" % data[1],
                                  hdbcom=hdbcom, maxtime=maxtime,verbose=verb)
            nmat = len(tab.recs)
            if verb >=1: print "Found %d records" % nmat
            sys.stdout.flush()
            fldlist[fldname].tables[data[1]] = tab
            if tableregionsize > 0. and len(tab.recs) > 0:
                tabregfile = tableroot + "_%s.reg" % data[1]
                reg = saoimage_reg.ds9_reg()
                for rec in tab.recs:
                    tabrec = gen_match_data(data[1])
                    if tabrec.extract_from_hdb(rec):
                        break
                    reg.add_object("circle", tabrec.pos2.ra,
                                   tabrec.pos2.dec, [tableregionsize],
                                   radec=1, text="%s" % (data[1] + "-" + \
                                                         tabrec.catid),
                                   color="blue")
                reg.Save(tabregfile)

matches = []
if verb >=1:
    print "Checking for catalog matches for %d source in memory" % len(csv.csv)
j = 0
nsrcs = len(csv.csv)
for data in csv.csv:
    j = j+1
    id = int(data[idcol])
    fldname = data[fldnamecol]
    detname = data[detnamecol]
    telname = data[telnamecol]
    ext = data[extcol]
    #if ext == "T":
    #    if verb >=1: print "Source %d in %s is extended, continuing"
    #    continue
    ra = float(data[racol])
    dec = float(data[deccol])
    try:
        flux = float(data[fluxcol])
    except:
        flux = 0.
    try:
        sigmaj = float(data[sigmajcol])
    except:
        sigmaj = 0.1
    match_rad = sqrt((sigmaj*3.)**2 + aspunc**2)
    ad = radec.radec()
    ad.ra = ra
    ad.dec = dec
    if verb >=2: print "Chandra source %s %.2g %.2f" % (ad, flux, sigmaj)
    if verb >=2: print "%.1f%% complete" % (j*100./nsrcs)
    i = 1
    for table in fldlist[fldname].tables.keys():
        if verb >=2: print "Searching for matches in %s..." % table
        for rec in fldlist[fldname].tables[table].recs:
            #if verb >=1: print "Offset %d = %.2f\"" % (i, float(rec["OFFSET"])*60.)
            tabrec = gen_match_data(table)
            if tabrec.extract_from_hdb(rec):
                break
            #pos = radec.radec(rec["ra"] + ' ' + rec["dec"])
            d = tabrec.pos2.Dist(ad)*3600
            if d <= match_rad:
                tabrec.set_xray_data(ra, dec, flux, sigmaj, id, fldname,
                                     telname, detname, ext)
                #tabrec.pos2 = pos
                #tabrec.Rmag = float(rec["rmag"])
                #tabrec.Bmag = float(rec["bmag"])
                matches.append(tabrec)
                i = i + 1

if verb >=1: print "Summary:"
reg = saoimage_reg.ds9_reg()
outfil = open(outfile, 'w')
outfil.write("fldname,telname,detname,id,xradec,xra,xdec,flux,semimaj_axis,catalogname,catid,catradec,catra,catdec,offset")
for i in range(1,ndatacols+1):
    outfil.write(",colname%d,coldata%d" % (i,i))
outfil.write("\n")

if sepsrcs:
    srcfiles = {}
    srcroot = xa_util.rootname(outfile) + "_src"

for match in matches:
    d = match.pos1.Dist(match.pos2)*3600
    xid = "%s_%s" % (match.detname, match.id)
    reg.add_object("circle", match.pos2.ra, match.pos2.dec, [2.], radec=1,
                   text="%s-%s/%s %.1f\"" % (match.catalog,
                                             match.catid,
                                             xid, d), color="cyan")
    if verb >=1: print match
    line = "%s,%s,%s,%d,%s,%f,%f,%.2g,%.2f,%s,%s,%s,%f,%f,%f" % (\
        match.fldname,
        match.telname,
        match.detname,
        match.id,
        match.pos1,
        match.pos1.ra,
        match.pos1.dec,
        match.flux,
        match.sigmaj,
        match.catalog,
        match.catid,
        match.pos2,
        match.pos2.ra,
        match.pos2.dec,
        d)
    keys = match.extra_data.keys()
    keys.sort()
    for i in range(ndatacols):
        if len(keys) > i:
            line = line + ",%s,%s" % (keys[i], match.extra_data[keys[i]])
        else:
            line = line + ",,"
    outfil.write(line + "\n")
    if sepsrcs:
        srcfile = srcroot + "%d.csv"  % match.id 
        if not srcfiles.has_key(srcfile):
            srcfiles[srcfile] = open(srcfile, 'w')
        srcfiles[srcfile].write(line + "\n")
        
if regfile is not None:
    if verb >= 1:
        print "Saving " + regfile
    reg.Save(regfile)