# Fits and subtracts an nth-order polynomial, then divides by the rms to return a S/N cube.
# This version uses the "robust" rms, which measures the median deviation of points from the median value, unlike
# the standard rms which uses the mean. This is more robust to outliers so should give more accurate results.

# An oddity occurred in the sigma-clipping phase of older versions of the code (which used the standard rms). If the
# deviant points were calculated with respect to the median of each spectrum, rather than from zero, then the
# resulting rms was higher and of higher variation from pixel to pixel. Renzograms showed significantly more noise
# and higher S/N values had to be used to properly reveal extended features.
# Using the robust noise fixes this. Alternatively one could use the deviation from zero instead of the mean, but in
# principle this should be less accurate. Still, I'm not sure exactly what's going on here.

import math
import numpy
import pyfits
import os

infile = 'VirgoOnly_h3_poly_2.fits'
clip = 2.0
niter =5
order =2

outfile = str(infile.split('.fits')[0])+str('_poly_')+str(order)+str('_SN.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)

FitsFile = pyfits.open(infile)

image = FitsFile[0].data
header =FitsFile[0].header

sizez = image.shape[0]
sizey = image.shape[1]
sizex = image.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(image[:,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)
				av  = numpy.median(tempspec)

				# Find the location of values to remove
				deviants = numpy.where(abs(tempspec - av) > 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)
			
			fitted = spectra - model
			
			# Get the robust rms, using the final masked spectrum
			medspec = numpy.median(fitted[tempchan])
			absdev = numpy.fabs(fitted[tempchan] - medspec)			# Fabs calculates an array of absolute values
			robrms = 1.4826 * numpy.median(absdev)							# See https://en.wikipedia.org/wiki/Median_absolute_deviation
			
			image[:,yp,xp] = fitted / robrms


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