# Spatial bandpass removal program that ignores NaNs
# Designed for use on a cube with bright sources already masked, as an alternative to MINMED or MEDMED
# Takes two input cubes : one with masking, one without. The masked cube is used to estimate the
# spatial bandpass. This is then subtracted from the unmasked cube.
# The advantage of this is that a much greater part of the scan can be used to
# estimate the bandpass, so it should be more accurate than MINMED or MEDMED.

import math
import numpy
import pyfits
import os

# Changes to the directory where the script is located
abspath = os.path.abspath(__file__)
dname = os.path.dirname(abspath)
os.chdir(dname)

maskedfile = 'AGES5523_6000-10200kms_masked.fits'
unmaskedfile = 'AGES5523_6000-10200kms.fits'

outfile = outfile = str(unmaskedfile.split('.fits')[0]+str('_UserMed.fits'))

MaskFitsFile = pyfits.open(maskedfile)

maskimage= MaskFitsFile[0].data
header = MaskFitsFile[0].header

sizez = maskimage.shape[0]
sizey = maskimage.shape[1]
sizex = maskimage.shape[2]

print sizex,sizey,sizez

UnmaskFitsFile = pyfits.open(unmaskedfile)

unmaskimage = UnmaskFitsFile[0].data


for yp in range(0,sizey):
	for zp in range(0,sizez):

		print yp

		bpass = numpy.array(maskimage[zp,yp,:])
		# Completely remove nans and indices
		bpass = bpass[numpy.logical_not(numpy.isnan(bpass))]
			
		med = numpy.median(bpass)
			
		unmaskimage[zp,yp,:] = unmaskimage[zp,yp,:] - med

FitsFile = pyfits.writeto(outfile,unmaskimage,header=header)