# Simulates the behaviour of MinMed on a reduced data cube
# User specifies only number of boxes and cube to operate on
# At every channel, splits spatial bandpass into n boxes and finds the median in each. Subtracts
# the minimum of the medians from the entire bandpass.

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)

infile = 'AGES7448_group.fits'
outfile = str(infile.split('.fits')[0]+str('_SimMinMed.fits'))

# Specify number of boxes
n = 5


FitsFile = pyfits.open(infile)

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

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

print sizex,sizey,sizez

for yp in range(0,sizey):
	for zp in range(0,sizez):
		
		print yp
		
		vlist = []
		medlist = []
		
		bpass = numpy.array(image[zp,yp,:])
			
		# No. boxes examined so far
		nb = 0
		# No. values in each box	
		lb = int(float(len(bpass))/float(n))
			
		for i in range(0,n):
			for j in range(nb*lb,(nb*lb)+lb):	
				vlist.append(bpass[j])
				
			# Remove nans from vlist. Do this now, rather than removing from the bandpass, because
			# this makes it easier to iterate over the (fixed) length of the bandpass (rather than
			# what would be the variable length of vlist).
			# Have to convert vlist into numpy array first
			vlist = numpy.array(vlist)
			vlist = vlist[numpy.logical_not(numpy.isnan(vlist))]
					
			med =  numpy.median(vlist)
			medlist.append(med)
				
			vlist = []
				
			nb = nb + 1
			
		minmed = numpy.min(medlist)
			
		bpass = bpass - minmed
			
		image[zp,yp,:] = bpass
			
FitsFile = pyfits.writeto(outfile,image,header=header)