import matplotlib
if __name__ == '__main__':
	matplotlib.use('Agg')
import unittest
import pyfits
from astrometry.util.fits import *
from astrometry.blind.plotstuff import *
import numpy as np
import pylab as plt
from math import pi,sqrt

class TestPlotstuff2(unittest.TestCase):

	def testSubpixelShift(self):
		W,H = 100,100
		cx,cy = 34.2, 56.7
		X,Y = np.meshgrid(np.arange(W), np.arange(H))
		S = 2.
		G = 1./(2.*pi*S**2) * np.exp(-((X-cx)**2+(Y-cy)**2)/(2.*S**2))
		#G /= G.max()
		print G.sum()
		print np.sum(G*X), np.sum(G*Y)
		imfn = 'test-plotstuff-2.fits'
		pyfits.writeto(imfn, G, clobber=True)

		wcs = anwcs_create_box(33.5, 10.4, 1., W, H)
		anwcs_rotate_wcs(wcs, 25.)
		wcsfn = 'test-plotstuff-2.wcs'
		anwcs_write(wcs, wcsfn)

		plot = Plotstuff()
		plot.outformat = PLOTSTUFF_FORMAT_PNG
		plot.size = (W,H)
		plot.wcs_file = wcsfn
		plot.color = 'black'
		plot.plot('fill')

		im = plot.image
		im.image_low = 0
		im.image_high = G.max()
		plot_image_set_wcs(im, wcsfn, 0)
		plot_image_set_filename(im, imfn)
		plot.plot('image')

		plotfn = 'test-plotstuff-2.png'
		plot.write(plotfn)

		I = plt.imread(plotfn)
		I = I[:,:,0]
		sx,sy = (I*X).sum()/I.sum(), (I*Y).sum()/I.sum()
		print sx,sy
		ex,ey = cx,cy
		self.assertTrue(abs(sx-ex) < 0.1)
		self.assertTrue(abs(sy-ey)< 0.1)

		# Shift the plot's WCS CRPIX.
		dx = 0.25
		dy = 0.3
		sip = anwcs_get_sip(wcs)
		tan = sip.wcstan
		plotstuff_set_wcs(plot.pargs, wcs)

		xy = plot.xy
		plot_xy_set_wcs_filename(xy, wcsfn)
		plot.marker = 'crosshair'

		# Move the plot WCS origin
		for step in range(16):
			if step < 8:
				tan.set_crpix(tan.crpix[0]+dx, tan.crpix[1])
				ex += dx
			else:
				tan.set_crpix(tan.crpix[0], tan.crpix[1]+dy)
				ey += dy
			#anwcs_print_stdout(wcs)
			plot.color = 'black'
			plot.plot('fill')
			plot.plot('image')
			plotfn = 'test-plotstuff-2b%i.png' % step
			plot.write(plotfn)
			#anwcs_write(plot.pargs.wcs, 'test-plotstuff-2b.wcs')
			I = plt.imread(plotfn)
			I = I[:,:,0]
			sx,sy = (I*X).sum()/I.sum(), (I*Y).sum()/I.sum()
			print sx,sy
			print ex,ey
			self.assertTrue(abs(sx-ex) < 0.1)
			self.assertTrue(abs(sy-ey) < 0.1)

			plot.color = 'red'
			plot_xy_clear_list(xy)
			# don't plot at sx,sy / ex,ey -- original image coords
			# are unchanged.
			plot_xy_vals(xy, cx+1, cy+1)
			plot.plot('xy')
			plotfn = 'test-plotstuff-2c%i.png' % step
			plot.write(plotfn)
			# visual check that plot xy symbols match source centroid -- yes
			

		# Scan image WCS in RA,Dec; check that recovered source position
		# through plot WCS matches.
		# reset...
		plot.wcs_file = wcsfn
		plotwcs = anwcs_open(wcsfn, 0)

		wcs = anwcs_open(wcsfn, 0)
		im.wcs = wcs
		sip = anwcs_get_sip(wcs)
		tan = sip.wcstan
		ddec = 1.2 * 1./W
		ok,era,edec = anwcs_pixelxy2radec(wcs, cx, cy)
		print era,edec
		for step in range(16):
			tan.set_crval(tan.crval[0], tan.crval[1] + ddec)
			edec += ddec

			plot.color = 'black'
			plot.plot('fill')
			plot.plot('image')
			plotfn = 'test-plotstuff-2d%i.png' % step
			plot.write(plotfn)

			I = plt.imread(plotfn)
			I = I[:,:,0]
			sx,sy = (I*X).sum()/I.sum(), (I*Y).sum()/I.sum()
			#print sx,sy
			ok,ra,dec = anwcs_pixelxy2radec(plotwcs, sx, sy)

			#print era,edec
			#print ra,dec
			print 'dRA,dDec', ra-era, dec-edec
			self.assertTrue(abs(ra-era) < 1e-4)
			self.assertTrue(abs(dec-edec) < 1e-4)

if __name__ == '__main__':
	unittest.main()
