""" 
wisecc.py reads in a table of WISE catalog sources and plots a color-color diagram

To generate the table, go to the web site
http://wise2.ipac.caltech.edu/docs/release/allsky/
then follow the link "Access the Catalog and Image data via IRSA".
You might want to check out some of the other links on that page
in order to get a better understanding of WISE and the data products.

On the page
http://irsa.ipac.caltech.edu/Missions/wise.html
click on the link "WISE Catalog Search"
select "WISE All-Sky Source Catalog" and click on the red Select button.

For this example, do a "Single Object Search".  (The search will
actually return multiple objects, but examines only one part of the sky.)

Type in coordinates or an oject name, specify a cone search with a radius
of 10 arcminutes, and hit "Run Query".  You'll be brought to a page
with the search results.  You may need to wait a few seconds for the
search to complete.

Click on "Download Table" and save the file in a name of your choice.
Edit the line in this file (wisedatafile = ...) to use your chosen file name.  

In ipython do "run wisecc"
"""


def float_safe(s):
# convert string to float, 
# catches ValueError exception and returns zero
# this handles the presence of 'null' for numeric values
  try:
    x = float(s)
  except ValueError:
    x = 0.0
  return x


wisedatafile = 'wise_allsky.wise_allsky_4band_p3as_psd14626.tbl'

import matplotlib.pyplot as plt
from numpy import array, where


f = open(wisedatafile, 'r')
name = []
w1, w1err, w1snr = [], [], []
w2, w2err, w2snr = [], [], []
w3, w3err, w3snr = [], [], []
s = f.readline()
while len(s) > 0:
  if (s[0] == ' ') and (len(s) > 1): # source lines start with space
    t = s.split()
    # print 't=', t
    name.append(t[0])
    w1.append(float(t[6]))
    w1err.append(float_safe(t[7]))
    w1snr.append(float(t[8]))
    w2.append(float(t[10]))
    w2err.append(float_safe(t[11]))
    w2snr.append(float(t[12]))
    w3.append(float(t[14]))
    w3err.append(float_safe(t[15]))
    w3snr.append(float(t[16]))
  s = f.readline()
# convert numerical data to numpy arrays
w1, w1err, w1snr = array(w1), array(w1err), array(w1snr)
w2, w2err, w2snr = array(w2), array(w2err), array(w2snr)
w3, w3err, w3snr = array(w3), array(w3err), array(w3snr)

# keep only sources with snr >= 4 in all three bands
good = where((w1snr >= 4) & (w2snr >= 4) & (w3snr >= 4))
w1, w1err = w1[good], w1err[good]
w2, w2err = w2[good], w2err[good]
w3, w3err = w3[good], w3err[good]

# plot color-color diagram
plt.clf()
#plt.plot(w2-w3, w1-w2, '+b')
plt.errorbar(w2-w3, w1-w2, w2err-w3err, w1err-w2err, '+k')
plt.xlabel('W2-W3 (mag)')
plt.ylabel('W1-W2 (mag)')
plt.show()