# Masks werid-shaped volumes in a FitsFile from a list of coordinates
# Takes as input a FITS file and an ASCII mask file
# Input mask format : x, y, z, flux
# Creates a new image of the same shape as the original FITS file image 
# but filled with NaN, then repopulates this with only the points in 
# the text file

import math
import numpy
from astropy.io import fits as pyfits
import os

infile   = 'GALFAM33_LikeAGES_h3.fits'
textfile = 'AGESOnlyM33.txt'
outfile  = str(infile.split('.fits')[0])+str('_Masked.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]

print(sizez,sizey,sizex)

# Create an array of the same original shape but fill with Nan
newimage = numpy.full_like(image, numpy.nan)

coords = open(textfile,'r')

for line in coords:
	# Skip header and other commented lines
	if '#' in line:
		continue
		
	x, y, z, v = line.split()
	xp = int(x)
	yp = int(y)
	zp = int(z)
	
	v = float(v)
	
	newimage[zp,yp,xp] = v
	
FitsFile = pyfits.writeto(outfile, newimage, header=header, overwrite=True)