##^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^Lance Simms, Stanford University 2009
#Return_PSF_2_1m_TrackingErrors.py
#
#PURPOSE:
#  To numerically integrate the function believed to represent the 
#  telescope tracking errors and produce a PSF for the system.
#
#INPUTS:
#  m_dec: float    
#    Slope of tracking error in DEC (arcsec/sec)
#  m_ra: float
#    Slope of tracking error in RA  (arcsec/sec)
#  eps_see: float 
#    Seeing disc in pixels          (arcsec)
#  f_osc: float
#    Frequency of oscillation of RA (1/sec)
#  A_osc: float
#    Amplitude of oscillation of RA (arcsec)
#  t_exp: float 
#    Exposure time in seconds       (seconds)
#  I_o: float
#    Maximum intensity of star      (ADU)
#KEYWORDS:
#  XSize, YSize: ints
#    The x and y size of the box used to contain the stars
#  XOff, YOff, ints
#    The offset from 0 for each box
#  BoxSize: int
#    The size of the box to use for the radial profile
#  Num_ts: int
#    The number of times over the total exposure time, t_exp, 
#    for which the pointing was corrected.  This term should 
#    be used to account for adjustment to the telescope
#  RA_DEC_Ratio: int
#    The ratio of the number of times that the DEC adjustment is 
#    mad per unit to that of the RA adjustment
#  Mode: int
#    The mode to use for the radial average
#  TvFlag: int
#    0 - Don't plot the star
#    1 - Plot the star
#
from scipy.integrate import quad
from ReturnRadialAvg import *
from RadialFits import *
execfile('/afs/slac/u/ki/lances/python/python_HxRG/HxRG_Setup.py')
FWToStd = 2.35482

def Return_PSF_2_1m_TrackingErrors(m_dec, m_ra, eps_see, \
    f_osc, A_osc, t_exp, I_o, XSize=31, YSize=31, XOff=15, YOff=15, \
    Num_ts = 1, BoxSize = 10, Mode = 0, TvFlag=0, RA_DEC_Ratio = 4):
 
  PSFFunc = lambda t, x, y, I_o, m_dec, m_ra, dec_t_o, ra_t_o, f_osc, A_osc, eps_see : \
            I_o*exp(-(x-m_ra*(t-ra_t_o)-\
                      A_osc*(sin(2*pi*f_osc*t)-sin(2*pi*f_osc*ra_t_o)))**2/\
                     (2*(eps_see/FWToStd)**2))\
               *exp(-(y-m_dec*(t-dec_t_o))**2/\
                     (2*(eps_see/FWToStd)**2))

  #Arrays
  if Num_ts == 1 :
    t_inits        = [0] 
    t_finals       = [t_exp]  
    RA_Move_Times  = [0]
    DEC_Move_Times = [0]
  else:
    t_interval = t_exp/Num_ts
    t_inits    = arange(Num_ts)     * t_interval
    t_finals   = (arange(Num_ts)+1) * t_interval
    
    #Number of DEC moves per RA move 
    RA_TimeInt     = t_interval
    RA_Move_Times  = arange(int(Num_ts))*RA_TimeInt
    DEC_TimeInt    = (t_exp/Num_ts)*RA_DEC_Ratio
    DEC_Move_Times = (arange(int(Num_ts))/RA_DEC_Ratio)*DEC_TimeInt

  XVals = frange(XSize)-XOff
  YVals = frange(YSize)-YOff
  Image = zeros([YSize,XSize], dtype=float)
  XCoo  = zeros([YSize,XSize], dtype=int16)
  YCoo  = zeros([YSize,XSize], dtype=int16)

  for t_init, t_final, t_ind in zip(t_inits, t_finals, arange(Num_ts)):
    DEC_t_o = DEC_Move_Times[int(t_ind)]
    RA_t_o  = RA_Move_Times[int(t_ind)]
    for i, Xc in zip(arange(XSize), XVals):
      for j, Yc in zip(arange(YSize), YVals):
        Int = quad(PSFFunc, t_init, t_final, args=(Xc,Yc, I_o, \
                   m_dec, m_ra, DEC_t_o, RA_t_o, f_osc, A_osc, eps_see))
        Image[j,i] = Image[j,i]+Int[0]
        #Only update the spatial grid on the first time loop
        if t_ind == 0:
          XCoo[j,i]  = i
          YCoo[j,i]  = j
  MaxIm         = Image.max()
  MaxImLocX     = where(Image == MaxIm)[0] 
  MaxImLocY     = where(Image == MaxIm)[1]
  BoxSize       = BoxSize
  XCStart       = MaxImLocX[0]-BoxSize-1
  XCStop        = MaxImLocX[0]+BoxSize
  YCStart       = MaxImLocY[0]-BoxSize-1
  YCStop        = MaxImLocY[0]+BoxSize
  #Find the new Centroid
  Centroid      = ReturnCentroid(Image, XCStart, XCStop, YCStart, YCStop)  
  XStart        = Centroid[0]
  XStop         = Centroid[1]
  YStart        = Centroid[2]
  YStop         = Centroid[3]
  IntArr3       = ReturnRadialAvg(Image[YStart:YStop, XStart:XStop], Mode=0, \
                  DoCentroid=0)
  Radii, IntArr = ReturnRadialAvg(Image[YStart:YStop, XStart:XStop], Mode=1, \
                  DoCentroid=0)
  #Do some radial fits
  Height, MuX, MuY, FWHMX, FWHMY, FWHM, E, EA = \
           moments(IntArr, SubSize=1)
  Amp_G,   FWHM_G,   Sky_G,   Fit_G = RadialGaussFits(Radii, IntArr,\
       IncFitData=1)
  Amp_G_M, FWHM_G_M, Beta_G_M, Sky_G_M, Fit_G_M = RadialMoffatFits(\
      Radii, IntArr, IncFitData=1)
  
  if TvFlag == 1:  
    #Show the image
    mplot.figure(0)
    mplot.clf()
    mplot.matshow(Image, extent=[YVals.min(), YVals.max(), XVals.min(), XVals.max()], \
                  fignum=0)
    mplot.colorbar()
    
    #Show the radial profile
    mplot.figure(1)
    mplot.clf()
    mplot.plot(Radii, IntArr/t_exp, 'ro')
    mplot.plot(Radii, Fit_G/t_exp, 'k-') 
    mplot.plot(Radii, Fit_G_M/t_exp, 'b-')
    xlim(0,14)
    print Image.min()
    print Image.max()
    FWHM_Gauss = 2*sqrt(2*log(2))*FWHM_G
    FWHM_Moff  = 2*sqrt(2**(-1/Beta_G_M)-1)*FWHM_G_M
    print 'Gauss FWHM  : ' + str(FWHM_Gauss)
    print 'Moffat FWHM : ' + str(FWHM_Moff)
    print 'Ellipticity : ' + str(E)

  return Height, MuX, MuY, FWHMX, FWHMY, FWHM, E, EA, Amp_G, FWHM_G, Fit_G, Amp_G_M, FWHM_G_M, Beta_G_M, Fit_G_M