from astro.main.SourceList import SourceList
from astro.main.AssociateList import AssociateList
import sys,os
import pylab, numpy
import random, math


(rr, ll, uu, bb) = (0.894, 0.900, 0.904, 0.910) #  PZ 3, SOA 1, LAS 3, POA 3

testarea = [ (ll,uu), (ll,bb), (rr,bb), (rr,uu) ]
testareaAL = [ ll,uu,ll,bb,rr,bb,rr,uu ]

attrlist = ['SLID','SID','RA','DEC']

s2s = {
	61804L: 'POA',
	155041L: 'PZ',
	244531L: 'LAS',
}

sourcelists = {
	'POA': (SourceList.SLID==61804)[0],
	'PZ': (SourceList.SLID==155041)[0],
	'LAS': (SourceList.SLID == 244531)[0], # DR3
}

colors = {
	'POA': 'r',
	'PZ':  'b',
	'LAS': 'y',
}


testdata = {}
for sl in sourcelists:
	#sourcelists[sl].debug = 10
	testdata[sl] = sourcelists[sl].get_data_dict_from_area_search_on_sources(attr_list=attrlist, Area=testarea)
	for k in testdata[sl]:
		testdata[sl][k] = numpy.array(testdata[sl][k])


datad = {'POA':{}, 'PZ':{}, 'SOA':{}, 'LAS':{}}
for sl in testdata:
	for si in range(len(testdata[sl]['SID'])):
		datad[sl][ testdata[sl]['SID'][si] ] = {'RA': testdata[sl]['RA'][si],'DEC': testdata[sl]['DEC'][si]}

# move the points around somewhat for enhanced visibility
datad['POA'][52207731]['RA'] += 0.00008
datad['PZ'][52209997]['RA'] += 0.00008
for sid in datad['PZ']:
	datad['PZ'][sid]['DEC'] -= 0.00008
for sid in datad['SOA']:
	datad['SOA'][sid]['DEC'] -= 0.00011


if(True):
	pylab.figure()
	for sl in sourcelists:
		pylab.plot( testdata[sl]['RA'], testdata[sl]['DEC'], colors[sl]+'o', label=sl)
	pylab.title('All Sources, Real Position\nRed=PhotoObjAll, Blue=Photoz, Yellow=lasSource')
	pylab.xlabel('RA (deg)')
	pylab.ylabel('DEC (deg)')
	pylab.savefig('allsourcesrealpos.png')
	
	pylab.figure()
	for sl in datad:
		for sid in datad[sl]:
			pylab.plot( [datad[sl][sid]['RA']], [datad[sl][sid]['DEC']], colors[sl]+'o')
	
	pylab.title('All Sources, Positions Shifted\nRed=PhotoObjAll, Blue=Photoz, Yellow=lasSource')
	pylab.savefig('allsources.png')


alids = [100651,100661] # POA+PZ+LAS,

als = [(AssociateList.ALID == alid)[0] for alid in alids]


for i in range(len(alids)):
	
	alid = alids[i]
	al = als[i]
	
	if False:
		adata = al.associates.get_data(['RA','DEC'])
		for aa in adata:
			ras = []
			decs = []
			for a in adata[aa]:
				ras.append(a[2])
				decs.append(a[3])
			pylab.plot(ras,decs)
	
	SD = al.process_params.SEARCH_DISTANCE / 3600.
	(SDp_ra, SDp_dec) = (0.8994,0.9080)
	
	pylab.figure()
	
	pylab.plot([ SDp_ra-(SD/2)/math.cos(SDp_dec *2*math.pi/360), SDp_ra+(SD/2)/math.cos(SDp_dec *2*math.pi/360)],
	           [ (SDp_dec), (SDp_dec)] , 'k-')
	pylab.plot([ (SDp_ra), (SDp_ra)],
	           [ (SDp_dec+SD/2), (SDp_dec-SD/2)] , 'k-')
	
	for sl in datad:
		for sid in datad[sl]:
			pylab.plot( [datad[sl][sid]['RA']], [datad[sl][sid]['DEC']], colors[sl]+'o')

	for aa in al.associates:
		ras = []
		decs = []
		for a in aa:
			ras.append( datad[s2s[a[0]]][a[1]]['RA'] + random.random()/60000 )
			decs.append( datad[s2s[a[0]]][a[1]]['DEC'] + random.random()/60000 )
		#pylab.plot( ras, decs, 'k-' )
		pylab.plot( ras, decs )
	
	pylab.title('AssociateList %i\nRed=PhotoObjAll, Blue=Photoz, Yellow=lasSource, Black=Search Distance' % alid)
	
	pylab.savefig('assoc'+str(alid)+'.png')


"""
alid = 100661
al = (AssociateList.ALID == alid)[0]
attrlist = ['RA','DEC']
adata = al.associates.get_data(attrlist)
for aa in adata:
    ras = []
    decs = []
    for a in adata[aa]:
        ras.append(a[2])
        decs.append(a[3])
    pylab.plot(ras,decs)

pylab.show()
"""

