##^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^Lance Simms, Stanford University 2009
##FitDonut
##
##PURPOSE:
##  To attempt to fit the radial profile of core-halo persistence 
##INPUTS
##  x - The radii points
##  y - The intensity in ADU
##      p[0] = Max intensity - A
##      p[1] = Frequency normalization - f
##      p[2] = Radial power in cosine  - p
##      p[3] = Diffusion term - sig
##      p[4] = Dark current term
##OPTIONAL INPUTS:
##  ReturnFit: int
##        0 - only return the coefficients
##        1 - return the fit in addition to the coefficients
##  Col - 0 - Do the col_deriv=0
##        1 - Do the col deriv=1, apparently this works alot better
##
execfile('/afs/slac/u/ki/lances/python/python_HxRG/HxRG_Setup.py')
import pdb

def FitDonut(p, x, y, Col=0, ReturnFit=0):

  #Define the fit function
  FitFunc  = lambda p, x: p[0]*cos(p[1]*x**p[2])*exp(-x**2/p[3])+p[4]
  ErrFunc1 = lambda p, x, y: y-FitFunc(p,x)

  #Try to fit
  p0=[max(Sig), .5, .75, 50, 0]
  p1, success = optimize.leastsq(ErrFunc1, p0[:], \
                               args = (r, Sig),full_output=0,\
                               maxfev=100000.,\
                               Dfun = DonutJacobian, col_deriv=1,\
                               xtol = 1.e-12, gtol=1e-15, ftol=1e-15)

  return p1

def DonutJacobian(p, x, y, Col=1):
  #'The Jacobian of the errfunc is -Jacobian of the func'
  if p.size == 5:
    A, f, pow, sig, dc = p
    if Col == 1:
      J = zeros([5, len(x)], dtype=float32 )

      ##dI/dA
      J[0,:] = -cos(f*x**pow)*exp(-x**2/sig)

      ##dS/df
      J[1,:] = A*(x**pow)*sin(f*x**pow)*exp(-x**2/sig)

      ##dS/dp
      J[2,:] = A*log(x)*(f*x**pow)*sin(f*x**pow)*exp(-x**2/sig)

      ##dS/dSig
      J[3,:] = -A*cos(f*x**pow)*(x**2/sig**2)*exp(-x**2/sig)

      ##dS/dSk 
      J[4,:] = -1
      return J

    else:

      J = zeros([len(x),5], dtype=float32 )
      ##dI/dA
      J[:,0] = -cos(f*x**pow)*exp(-x**2/sig)

      ##dS/df
      J[:,1] = A*(x**pow)*sin(f*x**pow)*exp(-x**2/sig)

      ##dS/dp
      J[:,2] = A*log(x)*(f*x**pow)*sin(f*x**pow)*exp(-x**2/sig)

      ##dS/dTau2
      J[:,3] = -A*cos(f*x**pow)*(1/sig**2)*exp(-x**2/sig)

      ##dS/dSk 
      J[:,4] = -1

      return J

  if p.size == 4:
    A, pow, sig, dc = p
    if Col == 1:
      J = zeros([4, len(x)], dtype=float32 )

      ##dI/dA
      J[0,:] = -cos(x**pow/(2*pi))*exp(-x**2/sig)

      ##dS/dp
      J[1,:] = A*log(x)*(x**pow/(2*pi))*sin(x**pow/(2*pi))*exp(-x**2/sig)

      ##dS/dSig
      J[2,:] = -A*cos(x**pow/(2*pi))*(x**2/sig**2)*exp(-x**2/sig)

      ##dS/dSk 
      J[3,:] = -1
      return J

    else:

      J = zeros([len(x),4], dtype=float32 )
      ##dI/dA
      J[:,0] = -cos(x**pow/(2*pi))*exp(-x**2/sig)

      ##dS/dp
      J[:,1] = A*log(x)*(x**pow/(2*pi))*sin(x**pow/(2*pi))*exp(-x**2/sig)

      ##dS/dTau2
      J[:,2] = -A*cos(x**pow/(2*pi))*(1/sig**2)*exp(-x**2/sig)

      ##dS/dSk 
      J[:,3] = -1

      return J