import pyfits

import matplotlib
matplotlib.use('Agg')
from pylab import *

from matplotlib.patches import Rectangle

from numpy import *

def density_plot(x, y, nbins=100):
    ax = [x.min(), x.max(), y.min(), y.max()]
    counts,xedges,yedges = histogram2d(x, y, bins=nbins)
    width = xedges[1] - xedges[0]
    height = yedges[1] - yedges[0]
    axes=gca()
    for (xe,countrow) in zip(xedges,counts):
        for (ye,c) in zip(yedges,countrow):
            culur='%g'%(c/counts.max())
            r=Rectangle([xe,ye], width, height,
                        axes=axes, clip_box=axes.bbox,
                        edgecolor=culur,facecolor=culur)
            axes.add_artist(r)
    axis(ax)


if __name__ == '__main__':

    p = pyfits.open('grillmair.fits')[1].data
    ra = p.field('ra')
    dec = p.field('dec')
    umag = p.field('psfmag_u') - p.field('extinction_u')
    gmag = p.field('psfmag_g') - p.field('extinction_g')
    rmag = p.field('psfmag_r') - p.field('extinction_r')
    imag = p.field('psfmag_i') - p.field('extinction_i')
    zmag = p.field('psfmag_z') - p.field('extinction_z')

    goodi = (gmag > -9999) * (rmag > -9999) * (imag > -9999)

    ra = ra[goodi]
    dec = dec[goodi]
    umag = umag[goodi]
    gmag = gmag[goodi]
    rmag = rmag[goodi]
    imag = imag[goodi]
    zmag = zmag[goodi]
    
    #clf()
    #plot(ra, dec, 'r.')
    #savefig('radec.png')

    clf()
    density_plot(ra, dec)
    xlabel('RA (deg)')
    ylabel('Dec (deg)')
    savefig('radec.png')
    
    clf()
    #plot(gmag-imag, rmag, 'r.')
    density_plot(gmag-imag, rmag)
    xlabel('g-i (mag)')
    ylabel('r (mag)')
    savefig('colormag.png')