# Simple script to replace empty spectra in a cube with NaNs.

import math as maths
import numpy
from astropy.io import fits as pyfits
import os
import shutil
import sys
import warnings
import sys

# Suppress some common pointless warning messages
warnings.filterwarnings("ignore")

# *** USER PARAMETERS ***
# Specify the name of the FITS file (must be in current directory and include a '.fits' extension)
infilename = 'InputCube.fits'
# *** END OF USER PARAMETERS ***


basename = infilename.split('.fits')[0]

FitsFile = pyfits.open(infilename)

image  = FitsFile[0].data
header = FitsFile[0].header
	
FitsFile.close()
	
sizez = image.shape[0]
sizey = image.shape[1]
sizex = image.shape[2]


for xp in range(sizex):
	sys.stdout.write('\r')
	sys.stdout.write('Checking pixel '+str(xp+1)+' of '+str(sizex)+'...')
	sys.stdout.flush()	
	
	for yp in range(sizey):
		spectrum = numpy.array(image[:,yp,xp])
		
		# Must sum to EXACTLY zero
		if numpy.sum(spectrum) == 0.0:
			image[:,yp,xp] = numpy.nan
			

NewFitsFile = pyfits.writeto(basename+'_NaNs.fits', image, header=header)	

print('Done !')