import numpy
import pyfits
from pyfits import open as pyopen
import os
import glob
import math
from math import *

# START OF THE BIT THE USER NEEDS TO READ
# Program to fit and remove known Gaussians from a data cube. Designed for the
# rare ideal case where all of the parameters of the sources are somehow already 
# known, i.e. from other observations.

# User inputs a file, "DeGaussSettings.txt", of the following format :
# Filename, xc, yc, a, b, pos
# Where xc, yc are the center coordinates in pixels, a and b are the major and
# minor FWHM in pixels, and pos is the position angle in degrees.
# "Filename" is a file containing the spectral profile of a source, in the following
# format : z, peak flux (Jy). Where z is the z pixel (channel) value.

# The input file can have multiple lines and so multiple input spectra. No fitting is
# done on the sources, Guassians exactly of the form specified are simply created and
# removed.
# Outputs a new FITS file, '*_degaussed.fits'

# Here set the FITS file to process :
FitsName = 'Blank.fits'
# Set the box over which the Gaussians should be removed. Use integers !
xmin = 0
ymin = 0
xmax = 50
ymax = 50
# END OF THE BIT THE USER NEEDS TO READ, UNLESS SOMETHING GOES HORRIBLY WRONG, WHICH
# IT DEFINITELY DEFINITELY WON'T.



# Change to directory where script is located
abspath = os.path.abspath(__file__)
dname = os.path.dirname(abspath)
os.chdir(dname)

# Open the parameters file, get the name of the FITS file

settings = open('DeGaussSettings.txt','r')

pi = 3.14159265359

outfile = FitsName+'_deguassed.fits'

FitsFile = pyopen(FitsName)

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

# Create blank FITS files of the same dimensions as this one
gaussimage = numpy.zeros_like(image)
degauimage = numpy.zeros_like(image)


for line in settings:
	# Get the basic parameters
	specfile, xc, yc, a, b, pos = line.split()

	spectra = str(specfile)
	xc = float(xc)
	yc = float(yc)
	a  = float(a)
	b  = float(b)
	pos= float(pos)
	
	pos = math.radians(float(pos))

	# Set Gaussian parameters	
	sigmax = a / (2.0*sqrt(2.0*log(2.0)))
	sigmay = b / (2.0*sqrt(2.0*log(2.0)))
	A = ((cos(pos)*cos(pos))/(2.0*sigmax*sigmax)) + ((sin(pos)*sin(pos))/(2.0*sigmay*sigmay))
	B = -(sin(2.0*pos)/(4.0*sigmax*sigmax)) + (sin(2.0*pos)/(4.0*sigmay*sigmay))
	C = ((sin(pos)*sin(pos))/(2.0*sigmax*sigmax)) + ((cos(pos)*cos(pos))/(2.0*sigmay*sigmay))

		
	spectrum = open(spectra,'r')
	
	# Get the peak flux value for each channel
	for line in spectrum:
		zp, flux = line.split()
		zp = int(zp)
		flux = float(flux)
		
		# Create the appropriate Gaussian over the required box
		for xp in range(xmin,xmax):
			for yp in range(ymin,ymax):
			
				rx = xp - (xc-xmin)
				ry = yp - (yc-ymin)
				
				f = flux*exp(-(A*(rx*rx) + (2.0*B*rx*ry) + (C*ry*ry)))

				gaussimage[zp,yp,xp] = gaussimage[zp,yp,xp] + f
			

degauimage = image - gaussimage

NewFitsFile = pyfits.core.writeto(outfile, degauimage, header=header, clobber='true')


print 'done !'