##^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^Lance Simms, Stanford University 2009
#RADIALFITS.py
#
#PURPOSE:
#  To do a gaussian, lorentzian, or moffat fit to a
#  stellar radial profile.
#INPUTS:
#  Intensity: float, int, or double array 
#   The y values of intensity.  This can be a 1d or 2d array
#   If 2-d, the second moments are computed for input to the 
#   fit.
#  Radii: float, int, or double array
#    THe x values of radius
#
#KEYWORDS:
#  IncFitData: int
#    0 - Don't return the fit
#    1 - Return the fit
#
#RETURNS: 
#    p1: float
#      The Amplitude
#    p2: float
#      The standard deviation (multiply by 2.35482 to get FWHM)
#AUTHOR:
#  Lance Simms, Stanford University 9/08
#
#

import optimize
import pylab
import matplotlib.pylab as mplot
from numpy import *
from FitGaussian import *
import pdb

##RADIAL GAUSSIAN
def RadialGaussFits(Radii, Intensity, IncFitData=0):

 #GAUSSIAN: 
 # p[0] = Amplitude
 # p[1] = Sigma
 # p[2] = Background
 GaussFunc     = lambda p, x: p[0]*exp(-.5*x**2/(p[1]**2))+p[2]
 GaussErrFunc  = lambda p, x, y: y-GaussFunc(p,x)

 if size(Radii.shape) == 1:
   #1d array for Radii
   FWHMX = 2
 else:
   #2d array for Radii
   Height, MuX, MuY, FWHMX, FWHMY, FWHM, E, EA = moments(Intensity)

 #Calculate the Second moment and print it
 print 'FWHM from Moments:   ' + str(FWHM)
 print 'FWHM X from Moments: ' + str(FWHMX)
 print 'FWHM Y from Moments: ' + str(FWHMY)   

 BackgroundEst     = Intensity.min()
 PeakEst           = Intensity.max()-BackgroundEst
 Radii1d           = Radii.copy()
 Radii1d.shape     = Radii.size
 Intensity1d       = Intensity.copy()
 Intensity1d.shape = Intensity1d.size
 p0=[PeakEst, sqrt(FWHM), Intensity.min()]

 p1, Success = optimize.leastsq(GaussErrFunc, p0[:], col_deriv=1,\
                    args = (Radii1d, Intensity1d), full_output=0,\
                    maxfev=10000000., xtol = 1.e-15, ftol=1.e-15,  \
                    Dfun = GaussJacobian)

 if IncFitData == 0:
   return p1[0], p1[1], p1[2]
 else: 
   FittedData = p1[0]*exp(-.5*Radii**2/(p1[1]**2))+p1[2]
   return p1[0], p1[1], p1[2], FittedData 

def GaussJacobian(p, x, y, Col=1):

  A, Sig, Sky = p
  if Col == 1:
    J = zeros([3, len(x)], dtype=float32 )

    ##dI/dA
    J[0,:] = -exp(-.5*x**2/Sig**2)

    ##dI/dSig
    J[1,:] = -(A*x**2/(Sig**3))*exp(-.5*x**2/Sig**2)

    ##dI/dSky
    J[2,:] = -1

    return J

  elif Col == 0:
    J = zeros([len(x),3], dtype=float32 )

    ##dI/dA
    J[:,0] = -exp(-.5*x**2/Sig**2)

    ##dI/dSig
    J[:,1] = -(A*x**2/(Sig**3))*exp(-.5*x**2/Sig**2)

    ##dI/dSky
    J[:,2] = -1

    return J

##MOFFAT FIT
def RadialMoffatFits(Radii, Intensity, IncFitData=0):

 #GAUSSIAN: 
 # p[0] = Amplitude
 # p[1] = Sigma
 # p[2] = Background
 MoffatFunc     = lambda p, x: p[0]*(1+(x/p[1])**2)**p[2]+p[3]
 MoffatErrFunc  = lambda p, x, y: y-MoffatFunc(p,x)

 if size(Radii.shape) == 1:
   #1d array for Radii
   FWHMX = 2
 else:
   #2d array for Radii
   Height, MuX, MuY, FWHMX, FWHMY, FWHM, E, EA = moments(Intensity)

 #Calculate the second moment and print it out 
 FWHM_SecMom = 2*FWHM*sqrt(2*log(2))
 print 'FWHM :'+ str(FWHM_SecMom)

 FWHM_SecMom          
 BackgroundEst     = Intensity.min()
 PeakEst           = Intensity.max()-BackgroundEst
 Radii1d           = Radii.copy()
 Radii1d.shape     = Radii.size
 Intensity1d       = Intensity.copy()
 Intensity1d.shape = Intensity1d.size
 p0=[PeakEst, sqrt(FWHM), -2., Intensity.min()]
 p1, Success = optimize.leastsq(MoffatErrFunc, p0[:], col_deriv=1,\
                    args = (Radii1d, Intensity1d), full_output=0,\
                    maxfev=10000000., xtol = 1.e-15,\
                    Dfun = MoffatJacobian)

 if IncFitData == 0:
   return p1[0], p1[1], p1[2], p1[3]
 else: 
   FittedData = p1[0]*(1+(Radii/p1[1])**2)**p1[2]+p1[3]
   return p1[0], p1[1], p1[2], p1[3], FittedData 

def MoffatJacobian(p, x, y, Col=1):

  A, Sig, Beta, Sky = p
  if Col == 1:
    J = zeros([4, len(x)], dtype=float32 )

    ##dI/dA
    J[0,:] = -(1+(x/Sig)**2)**Beta

    ##dI/dSig
    J[1,:] = 2*A*Beta*(x**2/Sig**3)*(1+(x/Sig)**2)**(Beta-1)

    ##dI/dBeta
    J[2,:] = -A*log(1+(x/Sig)**2)*(1+(x/Sig)**2)**Beta

    ##dI/dSky
    J[3,:] = -1

    return J
  elif Col == 0:
    J = zeros([len(x),4], dtype=float32 )

    ##dI/dA
    J[:,0] = -(1+(x/Sig)**2)**Beta

    ##dI/dSig
    J[:,1] = 2*A*Beta*(x**2/Sig**3)*(1+(x/Sig)**2)**(Beta-1)

    ##dI/dBeta
    J[:,2] = -A*log(1+(x/Sig)**2)*(1+(x/Sig)**2)**Beta

    ##dI/dSky 
    J[:,3] = -1 

    return J