import numpy as np

def usnob_apply_cuts(X):
	# USNO-B sources (not Tycho-2)
	I = (X.num_detections >= 2)
	# no diffraction spikes
	I = np.logical_and(I, np.logical_not(X.flags[:,0]))
	#X = X[I]
	#print '%i pass USNO-B diffraction spike cut' % len(X)

	if hasattr(X, 'an_diffraction_spike'):
		f = X.an_diffraction_spike
		print 'f shape', f.shape
		#u = unique(f)
		#print 'unique flags:', u
		for j in range(8):
			print 'flag', j, 'vals', np.unique(f[:,j])
		I *= (np.logical_not(f[:,7]) * np.logical_not(f[:,6]))
		#X = X[good]
		#print '%i pass AN diffraction spike cut' % len(X)
	return I

def usnob_compute_average_mags(X):
	# Compute average R and B mags.

	#epoch1 = (X.field_1 > 0)
	#epoch2 = (X.field_3 > 0)
	#nmag = np.where(epoch1, 1, 0) + np.where(epoch2, 1, 0)
	#summag = np.where(epoch1, X.magnitude_1, 0) + np.where(epoch2, X.magnitude_3, 0)
	#X.r_mag = np.where(nmag == 0, 0, summag / nmag)

	# rewritten to (hopefully) use less memory
	X.r_mag = np.zeros(len(X), np.float32)
	# sources that only have a first-epoch measurement
	I = np.flatnonzero((X.field_1 > 0) * (X.field_3 == 0))
	X.r_mag[I] = X.magnitude_1[I]
	# sources that only have a second-epoch measurement
	I = np.flatnonzero((X.field_1 == 0) * (X.field_3 > 0))
	X.r_mag[I] = X.magnitude_3[I]
	# sources that have both
	I = np.flatnonzero((X.field_1 > 0) * (X.field_3 > 0))
	X.r_mag[I] = (X.magnitude_1[I] + X.magnitude_3[I]) / 2.

	# B
	#epoch1 = (X.field_0 > 0)
	#epoch2 = (X.field_2 > 0)
	#nmag = np.where(epoch1, 1, 0) + np.where(epoch2, 1, 0)
	#summag = np.where(epoch1, X.magnitude_0, 0) + np.where(epoch2, X.magnitude_2, 0)
	#X.b_mag = np.where(nmag == 0, 0, summag / nmag)

	X.b_mag = np.zeros(len(X), np.float32)
	#I = np.logical_or(X.field_0 > 0, X.field_2 > 0)
	#X.b_mag[I] = (
	#	((X.field_0[I] > 0)*X.magnitude_0[I] + (X.field_2[I] > 0)*X.magnitude_2[I]) /
	#	((X.field_0[I] > 0)*1. + (X.field_2[I] > 0)*1.))

	# sources that only have a first-epoch measurement
	I = np.flatnonzero((X.field_0 > 0) * (X.field_2 == 0))
	X.b_mag[I] = X.magnitude_0[I]
	# sources that only have a second-epoch measurement
	I = np.flatnonzero((X.field_0 == 0) * (X.field_2 > 0))
	X.b_mag[I] = X.magnitude_2[I]
	# sources that have both
	I = np.flatnonzero((X.field_0 > 0) * (X.field_2 > 0))
	X.b_mag[I] = (X.magnitude_0[I] + X.magnitude_2[I]) / 2.
