import sys

import matplotlib
matplotlib.use('Agg')

from math import pi
from pylab import *
from numpy import *
from numpy.random import *
import pyfits

# Given an image and an xylist (including estimated image sigma),
# look at a cutout around each source position, add noise, and recompute
# dcen3x3 to find the noisy peak.

def dcen3a(f0, f1, f2):
	s = 0.5 * (f2 - f0)
	d = 2. * f1 - (f0 + f2)
	if (d <= 1.e-10*f0):
		return None
	aa = f1 + 0.5 * s * s / d
	sod = s / d
	return sod * (1. + (4./3.) * (0.25 * d / aa) * (1. - 4. * sod * sod)) + 1.


def dcen3b(f0, f1, f2):
	a = 0.5 * (f2 - 2*f1 + f0)
	b = f1 - a - f0
	xc = -0.5 * b / a
	if not (0.0 < xc < 2.0):
		return None
	return xc

def dcen3(f0, f1, f2):
	return dcen3a(f0,f1,f2)

def dcen3x3(image):
	my0 = dcen3(image[0,0], image[1,0], image[2,0])
	my1 = dcen3(image[0,1], image[1,1], image[2,1])
	my2 = dcen3(image[0,2], image[1,2], image[2,2])

	mx0 = dcen3(image[0,0], image[0,1], image[0,2])
	mx1 = dcen3(image[1,0], image[1,1], image[1,2])
	mx2 = dcen3(image[2,0], image[2,1], image[2,2])

	if None in [ mx0, mx1, mx2, my0, my1, my2 ]:
		return None

	# x = (y-1) mx + bx
	bx = (mx0 + mx1 + mx2) / 3.
	mx = (mx2 - mx0) / 2.

	# y = (x-1) my + by
	by = (my0 + my1 + my2) / 3.
	my = (my2 - my0) / 2.;

	# find intersection
	xc = (mx * (by - my - 1.) + bx) / (1. + mx * my)
	yc = (xc - 1.) * my + by

	# check that we are in the box
	if not ((0.0 < xc < 2.0) and (0.0 < yc < 2.0)):
		return None

	return (xc,yc)


if __name__ == '__main__':
	imgfn = sys.argv[1]
	xyfn = sys.argv[2]
	
	print 'FITS Image', imgfn
	print 'xylist', xyfn

	p = pyfits.open(imgfn)
	I = p[0].data
	print 'Image is', I.shape

	p = pyfits.open(xyfn)
	xy = p[1].data
	x = xy.field('X')
	y = xy.field('Y')
	flux = xy.field('FLUX')
	print 'Sources:', len(x)

	hdr = p[1].header
	sigma = hdr['ESTSIGMA']
	print 'Estimated sigma', sigma

	N = 100

	dx = []
	dy = []
	dx2 = []
	dy2 = []
	fluxes = []
	fluxes2 = []
	Nout = 0
	Nout2 = 0
	for i in range(len(x)):
		#print 'peak',i
		ix = round(x[i]) - 1
		iy = round(y[i]) - 1
		cutout = I[iy-1:iy+2, ix-1:ix+2]
		if cutout.shape != (3,3):
			print 'cutout is', cutout.shape
			continue
		#cutout5 = I[range(iy-2, iy+3, 2), range(ix-2, ix+3, 2)]
		cutout5 = array([ I[range(iy-2, iy+3, 2), ix-2],
						  I[range(iy-2, iy+3, 2), ix  ],
						  I[range(iy-2, iy+3, 2), ix+2], ])
						  
		if cutout5.shape != (3,3):
			print 'cutout5 has shape', cutout5.shape
			print 'yrange', range(iy-2, iy+3, 2)
			print 'xrange', range(ix-2, ix+3, 2)
			print 'cutout5', cutout5
			continue

		for j in xrange(N):
			noise = normal(0, sigma, size=(3,3))
			img = cutout + noise
			cen = dcen3x3(img)
			# original center:
			xx = x[i] - ix
			yy = y[i] - iy
			if cen is not None:
				(cx,cy) = cen
				dx.append(cx - xx)
				dy.append(cy - yy)
				fluxes.append(flux[i])
			else:
				Nout += 1
				noise5 = normal(0, sigma, size=(3,3))
				cen2 = dcen3x3(cutout5 + noise5)
				if cen2 is None:
					Nout2 += 1
					continue
				(cx2, cy2) = cen2
				cx = 1. + (cx2-1.) * 2.
				cy = 1. + (cy2-1.) * 2.
				dx2.append(cx - xx)
				dy2.append(cy - yy)
				fluxes2.append(flux[i])


	dx = array(dx)
	dy = array(dy)
	dx2 = array(dx2)
	dy2 = array(dy2)

	print 'A total of', Nout,  'of', (N*len(x)), '(%i %%)' % int(round(Nout*100./float(N*len(x)))),  'peaks moved outside the 3x3 box.'
	print 'A total of', Nout2, 'of', (N*len(x)), '(%i %%)' % int(round(Nout2*100./float(N*len(x)))), 'peaks moved outside the 5x5 box.'
	figure()
	clf()
	subplot(1,2,1)
	plot(dx, dy, 'r.')
	axis('equal')
	ylim(-3,3)
	a=axis()
	axhline(0)
	axvline(0)
	axis(a)
	title('3x3 centroid error')
	xlabel('distance (pixels)')
		
	subplot(1,2,2)
	plot(dx2, dy2, 'b.')
	axis('equal')
	ylim(-3,3)
	axhline(0)
	axvline(0)
	axis(a)
	title('5x5 centroid error')
	xlabel('distance (pixels)')
	savefig('dxdy.png', dpi=75)
	subplot(111)

	dist = sqrt(dx**2 + dy**2)
	dist2 = sqrt(dx2**2 + dy2**2)
	clf()
	subplot(2,1,1)
	(nb,bins,patches) = hist(dist, 40)
	xlim(0,bins[-1])
	title('3x3 centroid error distance')
	#xlabel('pixels')
	subplot(2,1,2)
	hist(dist2, bins=bins)
	xlim(0,bins[-1])
	title('5x5 centroid error distance')
	xlabel('pixels')
	savefig('dists.png', dpi=75)

	clf()

	dboth = hstack((dx,dy))
	sigd = std(dboth)
	md = mean(dboth)
	xx = arange(dboth.min(), dboth.max(), 0.01)
	yy = (len(dboth)*(xx.max()-xx.min())/40) / (sigd * sqrt(2. * pi)) * exp(-((xx-md)**2)/(2*sigd**2))

	subplot(2,1,1)
	(n,bins,patches) = hist(dboth, 40)
	plot(xx, yy, 'r-')
	title('3x3 centroid coordinate errors (std=%g)' % sigd)
	#xlabel('pixels')

	dboth2 = hstack((dx2,dy2))
	sigd2 = std(dboth2)
	md2 = mean(dboth2)
	if len(dboth2):
		xx2 = arange(dboth2.min(), dboth2.max(), 0.01)
		yy2 = (len(dboth2)*(xx2.max()-xx2.min())/40) / (sigd2 * sqrt(2. * pi)) * exp(-((xx2-md2)**2)/(2*sigd2**2))
	else:
		xx2 = []
		yy2 = []
	subplot(2,1,2)
	hist(dboth2, bins=bins)
	plot(xx2, yy2, 'r-')
	title('5x5 centroid coordinate errors (std=%g)' % sigd2)
	xlabel('pixels')
	savefig('dboth.png', dpi=75)

	clf()
	subplot(2,1,1)
	plot(fluxes, dist, 'r.')
	ylabel('pixel distance')
	#xlabel('flux')
	title('3x3 centroid')
	subplot(2,1,2)
	plot(fluxes2, dist2, 'b.')
	ylabel('pixel distance')
	xlabel('flux')
	title('5x5 centroid')
	savefig('fluxdist.png', dpi=75)
	
	

