import numpy, re, sqlcl

def query(sql, url = sqlcl.dr4_url, fmt='csv'):
    """Like sqlcl.query (which it uses), but also perform the readlines()
    and throw errors as appropriate (or at least sometimes)"""
    
    sql = re.sub('^[ \t\n]*select[ \t]', 'select ', sql) # work-around CAS parser bug
    result = sqlcl.query(sql, url, fmt).readlines()

    if re.match('^ERROR', result[0]):
        msg = result[3][0:-1]
    elif re.match('^error_message', result[0]):
        msg = result[1][0:-1]
    elif result[-1].find('</font') >= 0:
        for i in range(-1, -20, -1):
            if result[i].find('Active Server Pages') >= 0:
                msg = "".join(result[i:])

                msg = re.sub('\n\n', '\n', re.sub('</font>|<p>|<font [^>]*>', '', msg))
                break
    else:
        msg = None

    if msg:            
        msg = re.sub('[.:][ \t\n]*$', '', msg)
        raise SyntaxError, msg + ":\n" + sql

    return result

def extract(lines, names):
    """
    Given a set of lines returned by a skyserver query (or
    a sqlcl query with the .readlines() method applied), extract
    the columns specified by "names" and return numpy arrays.

    See also demo() for plotting these values using sm
    """
    fields = lines[0][0:-1].split(',')

    inames = []
    for n in names:
        iname = -1
        for i in range(0, len(fields)):
            if fields[i] == n:
                iname = i
                break

        if iname < 0:
            raise NameError, ("Field %s not found" % n)

        inames += [iname]

    vecs = []
    for iname in inames:
        vecs.append(numpy.arange(len(lines) - 1, dtype = 'd'))

    j = 0
    for line in lines[1:]:
        vals = line[:-1].split(',')

        i = 0
        for iname in inames:
            vecs[i][j] = float(vals[iname]); i += 1

        j += 1

    if len(vecs) == 1:
        return vecs[0]
    else:
        return vecs

def getPhotoFlags(key = None):
    """Initialise a dictionary of flag names; if key is specified, only
    set that entry (and return it's value)

    This routine is usually called by photoFlags()
"""

    global photoFlagsArr

    try:
        type(photoFlagsArr)
    except NameError:
        photoFlagsArr = {}

    if key:
        flag_names = [key]
    else:
        flag_names = query("select dbo.fPhotoFlagsN(~0)")[1:][0][0:-1].split(' ')

    try:
        return photoFlagsArr[key]
    except KeyError:
        for n in flag_names:
            photoFlagsArr[n] = query("select dbo.fPhotoFlags('" + n + "')")[1:][0][0:-1]
            if photoFlagsArr[n]:
                photoFlagsArr[n] = int(photoFlagsArr[n])

    if key:
        return photoFlagsArr[key]

def photoFlags(name):
    """Return the value of a photoFlag (e.g. photoFlags('INTERP')); queries the CAS as needed"""
    
    try:
        return photoFlagsArr[name]
    except:
        getPhotoFlags(name)

        return photoFlagsArr[name]

def demo(dev = 'x11'):
    """Demonstrate using this file in conjunction with sm"""
    import sm

    dmag = 0.10
    badFlags = photoFlags('BRIGHT') | photoFlags('BLENDED') | \
               photoFlags('INTERP_CENTER') | photoFlags('NOTCHECKED')

    sql = """
select top 10000
   psfMag_g as g,     psfMag_r as r,     psfMag_i as i
from
   star O
where
   g < 20
   and psfMagErr_g < %g
   and psfMagErr_r < %g
   and psfMagErr_i < %g
   and 0 = (flags & 0x%x)""" % (dmag, dmag, dmag, badFlags)

    (g, r, i) = extract(query(sql), ('g', 'r', 'i'))

    if dev:
        sm.device(dev)
    sm.erase(); sm.ctype()

    gr_lim = (-0.5, 2); ri_lim = (-0.5, 2)
    sm.limits(gr_lim, ri_lim)
    sm.box()
    sm.ptype([11])
    sm.points(g-r, r-i)
    sm.ptype()
    
    sm.xlabel("g - r")
    sm.ylabel("r - i")

if __name__=='__main__':
    import sm
    demo('x11')
    print "Hit any key to exit"
    sm.cursor()