# Adds flux with a Gaussian radial profile into a FITS file (e.g. fake sources)

import math
from math import *
import numpy
import pyfits
import os

infile = '7448Group_blank.fits'

flux = 10.0		# Total to inject
beam = 3.5		# Beam FWHM in pixels
width= 20		# Box size over which flux is injected

p = 25			# x-pixel center
q = 25			# y-pixel center

pi = 3.14159265359

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

#FWHM = 2sqrt(2.ln(2)) * sigma

sigma  = beam/2.354820045

rx = 0.0
ry = 0.0

for i in range(p-(width/2), p+(width/2)):
	for j in range(q-(width/2), q+(width/2)):
		#1D gaussian
		#r = sqrt( (i-p)*(i-p) + (j-q)*(j-q) )
		#f = (flux/(sigma*sqrt(2.0*pi))) * exp(-0.5 * (r/sigma)*(r/sigma))
		#image[20,j,i] = f
		
		#2D gaussian, volume = flux
		#rx = i-p
		#ry = j-q
				
		#A = flux/(2.0*pi*sigma*sigma)

		#f = A*exp(-( ((rx*rx)/(2.0*sigma*sigma)) + ((ry*ry)/(2.0*sigma*sigma))))
		
		#2D gaussian, height = flux
		
		rx = i-p
		ry = j-q
	       
		f = flux*exp(-( ((rx*rx)/(2.0*sigma*sigma)) + ((ry*ry)/(2.0*sigma*sigma))))
		
		image[20,j,i] = f
			
FitsFile = pyfits.writeto(outfile, image, header=header)