import sys
import pyfits
from pylab import *
from astrometry.util.pyfits_utils import *
from astrometry.util.plotutils import *

if __name__ == '__main__':
	args = sys.argv[1:]
	if len(args) == 0:
		fn = 'echoPCII11.fits'
	else:
		fn = args[0]
	echo = table_fields(fn)

	#
	u = echo.dered_u
	g = echo.dered_g
	r = echo.dered_r
	p1s = 0.91*(u-g) + 0.415*(g-r) - 1.28
	
	# color cut
	CUT = (g < 20.3) * (u-g > 0.4) * (u-g < 1.4) * (g-r > 0.2) * (g-r < 0.7) * (p1s > -0.7) * (p1s < -0.25)
	CUT *= (r > 17)

	print '%i of %i pass color cut' % (sum(CUT), len(g))

	echo = table_fields(fn, CUT)

	c = 300000
	vr = echo.z * c

	vlo = -130
	vhi = -100

	pmlim = (-30,30)
	vrlim = (-300,300)
	subplot(2,1,1)
	plot(vr, echo.dered_r, 'r.')
	ylabel('(dereddened) r (mag)')
	xlim(vrlim)
	axvline(vlo)
	axvline(vhi)
	subplot(2,1,2)
	step = 15
	shift = 7.5
	bins = arange(-300+shift, 300+step+shift, step)
	hist(vr, bins=bins) #int(ceil((vr[I].max() - vr[I].min())/15)))
	xlim(vrlim)
	xlabel('radial velocity')
	savefig('rvsz.png')

	subplot(111)
	clf()
	J = (vr > vlo) * (vr < vhi)
	plot(echo.pmra, echo.pmdec, '.', color='0.5')
	plot(echo.pmra[J], echo.pmdec[J], 'r.')
	xlabel('pm ra')
	ylabel('pm dec')
	xlim(pmlim)
	ylim(pmlim)
	savefig('pm.png')

	clf()
	for i in range(len(echo.pmra)):
		if J[i]:
			fc = 'r'
		else:
			fc = 'k'
		ellipse(x=echo.pmra[i], y=echo.pmdec[i],
				width=2*echo.pmraerr[i], height=2*echo.pmdecerr[i],
				alpha=0.05, facecolor=fc, edgecolor='none')
	xlabel('pm ra')
	ylabel('pm dec')
	xlim(pmlim)
	ylim(pmlim)
	savefig('pm2.png')

	clf()
	subplot(2,1,1)
	plot(vr, echo.pmra, 'r.', alpha=0.5)
	xlabel('radial velocity')
	ylabel('pm ra')
	axvline(vlo)
	axvline(vhi)
	xlim(vrlim)
	ylim(pmlim)
	subplot(2,1,2)
	plot(vr, echo.pmdec, 'r.', alpha=0.5)
	xlabel('radial velocity')
	ylabel('pm dec')
	axvline(vlo)
	axvline(vhi)
	xlim(vrlim)
	ylim(pmlim)
	savefig('pm3.png')
	
	#clf()
	#plot(echo.pmra, echo.pmdec, '.', color='0.5')
	#plot(echo.pmra[CUT], echo.pmdec[CUT], 'r.')
	#xlabel('pm ra')
	#ylabel('pm dec')
	#xlim(pmlim)
	#ylim(pmlim)
	#savefig('pm4.png')