import numpy
import pyfits
from pyfits import open as pyopen
import os
import glob
import math
from math import *
import scipy
from scipy import optimize as opt
import time

global xg, yg

# Program to fit and remove non-circular Gaussians from a flat image.
# User specifies all the parameters of an initial guess, plus the range over
# which those parameters are allowed to vary. 

# In each channel, generates a Gaussian where the peak = peak value in that 
# channel. The center of the Gaussian is varied over the coordinates of the box.
# Residuals are calculated for each center, and the one with the lowest standard
# deviation used to determine the final center coordinates.

# Outputs a new FITS file, '*_degaussed.fits'


# 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')

# Make sure no line breaks in file name !
FitsName = str(settings.readline()).splitlines()[0]
guess = FitsName.split('.fits')[0]

pi = 3.14159265359

FWHM = 3.5		# Beam FWHM in pixels.

sigma = FWHM / (2.0*sqrt(2.0*log(2.0)))

outfile = str(guess)+'_NonCircGauss.fits'

FitsFile = pyopen(FitsName)

print 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)


InitialGuess = settings.readline()
Flux,x,y,a,b,pos,minx,maxx,miny,maxy = InitialGuess.split()
ParRange = settings.readline()
RF, Rx, Ry, Ra, Rb, Rpos = ParRange.split()


Flux = float(Flux)
x    = float(x)
y    = float(y)
a    = float(a)
b    = float(b)
pos  = math.radians(float(pos))
minx = int(minx)
maxx = int(maxx)
miny = int(miny)
maxy = int(maxy)
RF   = float(RF)
Rx   = float(Rx)
Ry   = float(Ry)
Ra   = float(Ra)
Rb   = float(Rb)
Rpos = math.radians(float(Rpos))

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))


x = x - 1.0
y = y - 1.0

pos = -pos

fluxrange = numpy.arange(Flux-RF, Flux+RF, 0.25)
xcrange   = numpy.arange(x-Rx, x+Rx, 0.1)
ycrange   = numpy.arange(y-Ry, y+Ry, 0.1)
aarange   = numpy.arange(a-Ra, a+Ra, 0.1)
barange   = numpy.arange(b-Rb, b+Rb, 0.1)
posrange  = numpy.arange(pos-Rpos, pos+Rpos, math.radians(1.0))

fluxrange = [Flux]
#xcrange = [x]
#ycrange = [y]
#barange = [a]
#aarange = [b]
posrange= [pos]

n = len(fluxrange)*len(xcrange)*len(ycrange)*len(aarange)*len(barange)*len(posrange)

print n, ' Gaussians to fit'
time.sleep(3)

print len(fluxrange)
print len(xcrange)
print len(ycrange)
print len(aarange)
print len(barange)
print len(posrange)

#time.sleep(5)

i = 0

minres = []

imslice = image[0,miny:maxy,minx:maxx]
gslice  = numpy.zeros_like(imslice)
rcube   = numpy.zeros((n,maxy-miny,maxx-minx))

gc = (2.0*sqrt(2.0*log(2.0)))


for aa in aarange:
	for ba in barange:
		for pang in posrange:

			sigmax = aa / gc
			sigmay = ba / gc

			A = ((cos(pang)*cos(pang))/(2.0*sigmax*sigmax)) + ((sin(pang)*sin(pang))/(2.0*sigmay*sigmay))
			B = -(sin(2.0*pang)/(4.0*sigmax*sigmax)) + (sin(2.0*pang)/(4.0*sigmay*sigmay))
			C = ((sin(pang)*sin(pang))/(2.0*sigmax*sigmax)) + ((cos(pang)*cos(pang))/(2.0*sigmay*sigmay))

			for flux in fluxrange: 
				for xc in xcrange:
					for yc in ycrange:				  

						i = i + 1

						print i


						for xp in range(0,maxx-minx):
							for yp in range(0,maxy-miny):

								rx = float(xp) - (xc-minx)
								ry = float(yp) - (yc-miny)

								f = Flux*exp(-(A*(rx*rx) + (2.0*B*rx*ry) + (C*ry*ry)))

								gslice[yp,xp] = gslice[yp,xp] + f
										  
						residual = imslice - gslice

						resdev = numpy.std(residual)

						minres.append([resdev, flux, xc, yc, aa, ba, pang, A, B, C])
						
						if i < n:
							rcube[i,:,:] = residual[:,:]
							#rcube[i,:,:] = gslice[:,:]
							
						gslice[:,:] = 0.0

# We have now iterated over all the parameters. Find the one which minimized the
# residual.

minres = numpy.asarray(minres)
minres = minres[minres[:,0].argsort()]
Flux = minres[0,1]
xc   = minres[0,2]
yc   = minres[0,3]
aa   = minres[0,4]
ba   = minres[0,5]
pang = minres[0,6]
A    = minres[0,7]
B    = minres[0,8]
C    = minres[0,9]

print 'The final parameters were : ', Flux,xc,yc,aa,ba,math.degrees(pang)

# Now use those parameters to construct and subtract the final, ULTIMATE GAUSSIAN !


for xp in range(minx,maxx+1):
	for yp in range(miny,maxy+1):
	
		rx = float(xp) - xc
		ry = float(yp) - yc
		
		f = Flux*exp(-(A*(rx*rx) + (2.0*B*rx*ry) + (C*ry*ry)))

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

degaussimage = image - gaussimage

NewFitsFile = pyfits.core.writeto(outfile, degaussimage, header=header, clobber='true')
NewFitsFile = pyfits.core.writeto('202BestFitGauass.fits', gaussimage, header=header, clobber='true')
NewFitsFile = pyfits.core.writeto('ResidualCube.fits', rcube, header=header, clobber='true')

print 'done !'