# Program designed to get really accurate baselines. User first masks the galaxies using FRELLED/miriad, and maybe also runs
# UserMed. The user should then also mask any visible RFI and the Milky Way.
# The program takes the unmasked (UserMed) cube and the masked cube as input, fits the baselines to the masked cube
# (to avoid including real sources in the baseline fit) and subtracts them from the UserMed cube.


import math
import numpy
import pyfits
import os

maskedfile = 'AGES7448_100-5000_poly2_masked.fits'
unmaskedfile='AGES7448_100-5000_poly2_unmasked.fits'
clip = 2.0
niter =5
order =2
outfile = str(unmaskedfile.split('.fits')[0])+str('_poly_')+str(order)+str('_goodbaselines.fits')

os.system('cls')

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

# Change to the directory where the .blend is located
#os.chdir(Blender.sys.expandpath('//'))

MaskFitsFile = pyfits.open(maskedfile)
UnMaskFitsFile=pyfits.open(unmaskedfile)

# Masked image from which the baseline will be estimated
mimage = MaskFitsFile[0].data
header = MaskFitsFile[0].header

# Cleaned image from which the baseline will be subtracted
cimage = UnMaskFitsFile[0].data

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

# Create an array to hold the channel numbers
channel = numpy.array([i for i in range(0,sizez)])


for xp in range(0,sizex):
	print xp
	for yp in range(0,sizey):

		spectra = numpy.array(mimage[:,yp,xp])
		
		nanlist=[]
		for i in range(0,len(spectra)):
			if numpy.isnan(spectra[i]):
				nanlist.append(i)
		
		tempspec = numpy.delete(spectra,nanlist)
		tempchan = numpy.delete(channel,nanlist)
		
		# Create duplicate arrays to fit sigma-clipped polynomials
		#tempspec = numpy.array(spectra)
		
		if len(tempspec)>order :
			for i in range(0,niter):
				length = len(tempchan)
					
				# Fit a polynomial
				temppoly = numpy.polyfit(tempchan, tempspec,order)
				tempmodl = numpy.array([0.0 for k in range(0,length)])
				
				for f in range(0,order+1):
					fn = float(-f + order)
					tempmodl = tempmodl + temppoly[f]*pow(tempchan,fn)
					
				tempspec = tempspec - tempmodl
				
				rms = numpy.std(tempspec)
				
				# Find the location of values to remove
				deviants = numpy.where(tempspec > clip*rms)
				
				if (len(deviants)>0 and len(deviants) < length):
					tempspec = numpy.delete(tempspec,deviants)
					tempchan = numpy.delete(tempchan,deviants)
				
			# Fit a polynomial to the original spectrum, masking all deviant values
			poly = numpy.polyfit(tempchan,spectra[tempchan],order)
			model = numpy.array( [0.0 for k in range(0,len(spectra))])
			for f in range(0,order+1):
				fn = float(-f + order)
				model = model + poly[f]*pow(channel,fn)
			
			cspectra = numpy.array(cimage[:,yp,xp])
			
			fitted = cspectra - model
			
			cimage[:,yp,xp] = fitted


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