# maskTimelines.py
#
# Scritp to mask bright sources before runing the basene subtraction
# starting from the Level1 product as available in the archive
#
# v. 1.0 - Luca Conversi & Anthony Smith - 2011/05/25


########################################################################
# Loading a small map observation of Neptune
########################################################################
obs     = getObservation(0x50004662, useHsa=1)
level1  = obs.level1
channel = 'PSW'



########################################################################
# Definition of the function used to (un)mask the selected data points
########################################################################
def maskTimeline(input, nSamples=None, ra=None, dec=None, radiusArcmin=None, unmask=False):
	"""Mask all except the first and last nSamples in each scan,
	and all within radiusArcmin of (ra, dec).
	"""
	try: # Treat it as a context
		for i in range(input.count):
			input.getRefs().set(i, ProductRef(maskTimeline(input.getProduct(i), nSamples, ra, dec, radiusArcmin, unmask)))
	except: # pointed product
		for channel in input.getChannelKeySet():
			mask = input.getMask(channel)
			if nSamples is not None:
				# Mask all except the first and last nSamples in each scan
				if unmask:
					mask[nSamples:-nSamples] = SpireMask.MASTER.unset(mask[nSamples:-nSamples])
				else:
					mask[nSamples:-nSamples] = SpireMask.MASTER.set(mask[nSamples:-nSamples])
			if ra is not None and dec is not None and radiusArcmin is not None:
				# Mask all within radiusArcmin of (ra, dec)
				raTimeline = input.getRa(channel)
				decTimeline = input.getDec(channel)
				d2 = ((raTimeline - ra) * COS(dec * Math.PI / 180.)) ** 2 + (decTimeline - dec) ** 2
				maskRegion = d2 < (radiusArcmin / 60.) ** 2
				if unmask:
					mask[mask.where(maskRegion)] = SpireMask.MASTER.unset(mask[mask.where(maskRegion)])
				else:
					mask[mask.where(maskRegion)] = SpireMask.MASTER.set(mask[mask.where(maskRegion)])
			input.setMask(channel, mask)
		return input



########################################################################
# Original map - No masking and median baseline subtraction
########################################################################
map_orig = obs.level2.refs[channel].product



########################################################################
# Mask all points except the first and last nSamples in each scan
# and uses a median baseline subtraction on each scan
########################################################################

# Selecting first & last 50 data points, masking the others
sample = 50
maskTimeline(level1, nSamples=sample)

# Run median baseline subtraction
scans = baselineRemovalMedian(input=level1, wholeTimeline=False)

# Remove the mask set above
maskTimeline(level1, nSamples=sample, unmask=True)
maskTimeline(scans, nSamples=sample, unmask=True)

# Run map-maker and compute difference
map_samples = naiveScanMapper(scans, array=channel, wcs=map_orig.getWcs())
difference1 = imageSubtract(image1=map_samples, ref=1, image2=map_orig)



########################################################################
# Mask all points within radiusArcmin of (RA, Dec)
# and uses a median baseline subtraction on each scan
########################################################################

# Selecting a circle of 2 arcmin radius centred on the map, 
# i.e. selecting the WCS reference pixel (which is roughly in the centre)
ra  = map_orig.getWcs().getCrval1()
dec = map_orig.getWcs().getCrval2()
rad = 2
maskTimeline(level1, ra=ra, dec=dec, radiusArcmin=rad)

# Run median baseline subtraction
scans = baselineRemovalMedian(input=level1, wholeTimeline=False)

# Remove the mask set above
maskTimeline(level1, ra=ra, dec=dec, radiusArcmin=rad, unmask=True) 
maskTimeline(scans, ra=ra, dec=dec, radiusArcmin=rad, unmask=True) 

# Run map-maker and compute difference
map_radius  = naiveScanMapper(input=scans, array=channel, wcs=map_orig.getWcs())
difference2 = imageSubtract(image1=map_radius, ref=1, image2=map_orig)



########################################################################
# Mask all points except the first and last nSamples in each scan
# and uses a 1st order polynominal baseline subtraction on each scan
########################################################################

# Selecting first & last 50 data points, masking the others
sample = 50
maskTimeline(level1, nSamples=sample)

# Run polynomial baseline subtraction
baselineRemovalPolynomial(input=level1, polyDegree=1, wholeTimeline=False)

# Remove the mask set above
maskTimeline(level1, nSamples=sample, unmask=True)
maskTimeline(scans, nSamples=sample, unmask=True)

# Run map-maker and compute difference
mapPoly_samples = naiveScanMapper(scans, array=channel, wcs=map_orig.getWcs())
difference3     = imageSubtract(image1=mapPoly_samples, ref=1, image2=map_orig)




########################################################################
# Mask all points within radiusArcmin of (RA, Dec)
# and uses a 1st order polynominal baseline subtraction on each scan
########################################################################

# Selecting a circle of 2 arcmin radius centred on the map, 
# i.e. selecting the WCS reference pixel (which is roughly in the centre)
ra  = map_orig.getWcs().getCrval1()
dec = map_orig.getWcs().getCrval2()
rad = 2
maskTimeline(level1, ra=ra, dec=dec, radiusArcmin=rad)

# Run median polynomial subtraction
scans = baselineRemovalPolynomial(input=level1, polyDegree=1, wholeTimeline=False)

# Remove the mask set above
maskTimeline(level1, ra=ra, dec=dec, radiusArcmin=rad, unmask=True) 
maskTimeline(scans, ra=ra, dec=dec, radiusArcmin=rad, unmask=True) 

# Run map-maker and compute difference
mapPoly_radius = naiveScanMapper(input=scans, array=channel, wcs=map_orig.getWcs())
difference4    = imageSubtract(image1=mapPoly_radius, ref=1, image2=map_orig)