#PlotFitCoefficientsForM77.py
#
#PURPOSE:
#  To plot the spatial coefficients for M77 and try to come up
#  with an approximation for removing the persistence.
#
#INPUTS:
##  	ObjectName: str
##          The name of the object containing the stars and persistence
##  	ThisFilter: str
##    	   'I', 'G', or 'Y'
##      ExpFun1: int
##          Choice of the function to fit the radial profile for a given read
##          1 - A*Cos(B*r^C)*Exp(-r^2/D)
##          2 - A*Cos(B*r^C)*Exp(-r^2/D)+E
##          3 - A*Cos(r^C/(2*pi))*Exp(-r^2/D)+E
##          4 - A*Cos(B*r^C)*r**D
##
##EXAMPLES:
##  run PlotFitCoefficients.py 'M77' 'I' 3
##
#As a test plot: mplot(ReadNums+max(ReadNums)*FileNums, Fit_C0s)
from Return_Persistence_DbQuantities import *
from Djs_Iterstat import *
from HxRG_Class import *
from Linefit_func import *
from DecayJacobian import *

ObjectName = sys.argv[1]
ThisFilter = sys.argv[2]
ExpFunc1   = int(sys.argv[3])

##PARSE KEYWORDS
keywords    = ['MinFWHM=', 'MaxFWHM=']
MinFWHM     = 0
MaxFWHM     = 0

opts, extraparams = getopt.getopt(sys.argv[4:],'',keywords)
for o,p in opts:
  if o in ['--MinFWHM']:
    MinFWHM  = int(p)
  elif o in ['--MaxFWHM']:
    MaxFWHM  = int(p)

if ObjectName == 'NGC2683':
  Mode = 0
  QueryStr   = ' AND H2RG_32_147_ASIC_TemporalFits_Illum.FILE_PIXID REGEXP '+" 'NGC2683' "
  if ThisFilter == 'G':
    FileList='/nfs/slac/g/ki/ki03/lances/H2RG-32-147/ASIC/Reduced/07Dec12/NGC2683/NGC2683_G_RawFiles.lst'
  elif ThisFilter == 'I':
    FileList='/nfs/slac/g/ki/ki03/lances/H2RG-32-147/ASIC/Reduced/07Dec12/NGC2683/NGC2683_I_RawFiles.lst'
  elif ThisFilter == 'Y':
    FileList='/nfs/slac/g/ki/ki03/lances/H2RG-32-147/ASIC/Reduced/07Dec12/NGC2683/NGC2683_Y_RawFiles.lst'
 
elif ObjectName == 'M77':
  Mode = 0
  QueryStr   = ' AND H2RG_32_147_ASIC_TemporalFits_Illum.FILE_PIXID REGEXP '+"'M77'"
  if ThisFilter == 'G':
    FileList='/nfs/slac/g/ki/ki03/lances/H2RG-32-147/ASIC/Reduced/07Dec14/M77/M77_G_RawFiles.lst'
  elif ThisFilter == 'I':
    FileList='/nfs/slac/g/ki/ki03/lances/H2RG-32-147/ASIC/Reduced/07Dec14/M77/M77_I_RawFiles.lst'
  elif ThisFilter == 'Y':
    FileList='/nfs/slac/g/ki/ki03/lances/H2RG-32-147/ASIC/Reduced/07Dec14/M77/M77_Y_RawFiles.lst'

elif ObjectName == 'Dark':
  Mode = 1
  QueryStr   = 'NGC956_40Reads_Dec_19_2007_20_46_52_X_452_Y_1985'
  FileList   = '/nfs/slac/g/ki/ki03/lances/H2RG-32-147/NegPersistence/NGC2395_5_Reads_G_Band.lst'
  
if (MinFWHM != 0) & (MaxFWHM != 0):
  QueryStr = QueryStr+' and H2RG_32_147_ASIC_Slopes_Illum.FWHM < ' + str(MaxFWHM) +\
                      ' and H2RG_32_147_ASIC_Slopes_Illum.FWHM > ' + str(MinFWHM)

FitsFiles=ReadFileList(FileList)
StartTimes = Persistence_Get_Dark_Times(FileList)
HxRG = HxRG_C(FitsFileName=FitsFiles[0])
HxRG.Get_Raw_Header(FitsFiles[0])
AvgExpTime=mean(StartTimes[1:]-StartTimes[0:-1])

FilePixIDs, ReadNums, FileNums, Fit_C0s, Fit_C1s, Fit_C2s, Fit_C3s, Fit_C4s, Filters=\
  Return_PersistenceFit_DbQuantities('H2RG-32-147', 'ASIC', ObjectName, QueryStr, Mode, ThisFilter=ThisFilter, \
                                     ExpFunc1=ExpFunc1)

NumFiles = Fit_C1s.size/(max(FileNums)*max(ReadNums))
NumReads = max(ReadNums)
NumDiths = max(FileNums)
FileNums.shape = (NumFiles, NumDiths, NumReads)
ReadNums.shape = (NumFiles, NumDiths, NumReads)
Fit_C0s.shape  = (NumFiles, NumDiths, NumReads)
Fit_C1s.shape  = (NumFiles, NumDiths, NumReads)
Fit_C2s.shape  = (NumFiles, NumDiths, NumReads)
Fit_C3s.shape  = (NumFiles, NumDiths, NumReads)
Fit_C4s.shape  = (NumFiles, NumDiths, NumReads)

NumPoints = NumReads*NumDiths
ReadTimes      = zeros(NumPoints, dtype =float64)
FitC0AvgArr    = zeros(NumPoints, dtype =float64)
FitC1AvgArr    = zeros(NumPoints, dtype =float64)
FitC2AvgArr    = zeros(NumPoints, dtype =float64)
FitC3AvgArr    = zeros(NumPoints, dtype =float64)
FitC4AvgArr    = zeros(NumPoints, dtype =float64)

if Mode == 0:

  for FileNum in arange(NumDiths):
    for ReadNum in arange(NumReads):
      FMean, FSig, FMedian, FMask = Djs_Iterstat(Fit_C0s[:,FileNum,ReadNum],RejVal=-1, SigRej=3.0)
      FitC0AvgArr[ReadNum+FileNum*NumReads] = FMean
      FMean, FSig, FMedian, FMask = Djs_Iterstat(Fit_C1s[:,FileNum,ReadNum],RejVal=-1, SigRej=4.0, Min=0.7, Max=1.8)
      FitC1AvgArr[ReadNum+FileNum*NumReads] = FMean
      FMean, FSig, FMedian, FMask = Djs_Iterstat(Fit_C2s[:,FileNum,ReadNum],RejVal=-1, SigRej=2.0, Min=0, Max=600)
      FitC2AvgArr[ReadNum+FileNum*NumReads] = FMean
      FMean, FSig, FMedian, FMask = Djs_Iterstat(Fit_C3s[:,FileNum,ReadNum],RejVal=-1, SigRej=2.0, Min=0, Max=600)
      FitC3AvgArr[ReadNum+FileNum*NumReads] = FMean
      FMean, FSig, FMedian, FMask = Djs_Iterstat(Fit_C4s[:,FileNum,ReadNum],RejVal=-1, SigRej=3.0)
      FitC4AvgArr[ReadNum+FileNum*NumReads] = FMean
      ReadTimes[ReadNum+FileNum*NumReads] = FileNum*AvgExpTime+ReadNum*HxRG.FrameTime

else:
  for FileNum in arange(NumDiths):
    for ReadNum in arange(NumReads):
      FitC0AvgArr[ReadNum+FileNum*NumReads] = Fit_C0s[:,FileNum, ReadNum]
      FitC1AvgArr[ReadNum+FileNum*NumReads] = Fit_C1s[:,FileNum, ReadNum]
      FitC2AvgArr[ReadNum+FileNum*NumReads] = Fit_C2s[:,FileNum, ReadNum]
      FitC3AvgArr[ReadNum+FileNum*NumReads] = Fit_C3s[:,FileNum, ReadNum]
      FitC4AvgArr[ReadNum+FileNum*NumReads] = Fit_C4s[:,FileNum, ReadNum]
      ReadTimes[ReadNum+FileNum*NumReads]   = FileNum*AvgExpTime+ReadNum*HxRG.FrameTime

if ExpFunc1 == 2:
  #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  #FIRST COEFFICIENT - Diffusion Term
  mplot.figure(0)
  mplot.clf()
  GP = ( FitC3AvgArr != -1 )
  mplot.plot(ReadTimes[GP], (FitC3AvgArr[GP])**(1./8), 'ro')
  BadPoints = ones(NumPoints)
  BadPoints[FitC3AvgArr == -1]=0
  m,b,error,cos_hit = Linefit(ReadTimes[0:120], FitC3AvgArr[0:120]**(1./8), \
                            pmask=BadPoints[0:120], RejectCos=0)
  mplot.plot(ReadTimes, m*ReadTimes+b,'k-')
  mplot.title('Diffusion Term')
  mStr = "%0.4f" % m
  bStr = "%0.4f" % b
  mplot.figtext(0.2,0.85, r'\textbf{m = '+ mStr +'}', fontsize=14)
  mplot.figtext(0.2,0.80, r'\textbf{b = '+ bStr +'}', fontsize=14)
  mplot.twinx()
  mplot.plot(ReadTimes[GP], FitC3AvgArr[GP], 'bo')
  mplot.plot(ReadTimes, (ReadTimes*m+b)**8,'g')
  
  #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  #SECOND COEFFICIENT - Stretching Term
  mplot.figure(1)
  mplot.clf()
  GP = ( FitC2AvgArr != -1 )
  mplot.plot(ReadTimes[GP], (FitC2AvgArr[GP])**(1./5), 'ro')
  BadPoints = ones(NumPoints)
  BadPoints[FitC2AvgArr == -1]=0
  m,b,error,cos_hit = Linefit(ReadTimes[0:120], (FitC2AvgArr[0:120])**(1./5), \
                              pmask=BadPoints[0:120], RejectCos=0)
  
  mplot.plot(ReadTimes, m*ReadTimes+b,'k-')
  mplot.title('Stretch Term')
  mStr = "%0.4f" % m
  bStr = "%0.4f" % b
  mplot.figtext(0.2,0.25, r'\textbf{m = '+ mStr +'}', fontsize=14)
  mplot.figtext(0.2,0.2, r'\textbf{b = '+ bStr +'}', fontsize=14)
  mplot.twinx()
  mplot.plot(ReadTimes[GP], FitC2AvgArr[GP], 'bo')
  mplot.plot(ReadTimes, (ReadTimes*m+b)**5,'g')
  
  #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  mplot.figure(2)
  mplot.clf()
  GP = ( FitC1AvgArr != -1 )
  mplot.plot(ReadTimes[GP], FitC1AvgArr[GP]**.3, 'ro')
  BadPoints = ones(NumPoints)
  BadPoints[FitC1AvgArr == -1]=0
  m,b,error,cos_hit = Linefit(ReadTimes[0:120], FitC1AvgArr[0:120]**.3, \
                              pmask=BadPoints[0:120], RejectCos=0)
  mplot.plot(ReadTimes, m*ReadTimes+b,'k-')
  mplot.title('Zeros Term')
  mStr = "%0.4f" % m
  bStr = "%0.4f" % b
  mplot.figtext(0.2,0.25, r'\textbf{m = '+ mStr +'}', fontsize=14)
  mplot.figtext(0.2,0.2, r'\textbf{b = '+ bStr +'}', fontsize=14)
  
  mplot.twinx()
  mplot.plot(ReadTimes[GP], FitC1AvgArr[GP], 'bo')
  mplot.plot(ReadTimes, (ReadTimes*m+b)**(10./3),'g')
 
if ExpFunc1 == 3:

  #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  #SECOND COEFFICIENT - Stretching Term
  mplot.figure(1)
  mplot.clf()
  GP = ( FitC2AvgArr != -1 )
  mplot.plot(ReadTimes[GP], (FitC2AvgArr[GP])**(1./2), 'ro')
  BadPoints = ones(NumPoints)
  BadPoints[FitC2AvgArr == -1]=0
  m,b,error,cos_hit = Linefit(ReadTimes[0:120], (FitC2AvgArr[0:120])**(1./2), \
                              pmask=BadPoints[0:120], RejectCos=0)

  mplot.plot(ReadTimes, m*ReadTimes+b,'k-')
  mplot.title('Diffusion Term')
  mStr = "%0.4f" % m
  bStr = "%0.4f" % b
  mplot.figtext(0.2,0.25, r'\textbf{m = '+ mStr +'}', fontsize=14)
  mplot.figtext(0.2,0.2, r'\textbf{b = '+ bStr +'}', fontsize=14)
  mplot.twinx()
  mplot.plot(ReadTimes[GP], FitC2AvgArr[GP], 'bo')
  mplot.plot(ReadTimes, (ReadTimes*m+b)**2,'g')

  #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  mplot.figure(2)
  mplot.clf()
  GP = ( FitC1AvgArr != -1 ) & (FitC1AvgArr > 0.4)
  GP = GP[0:199]
  mplot.plot(ReadTimes[GP], FitC1AvgArr[GP], 'ro')
  BadPoints = ones(NumPoints)
  BadPoints[FitC1AvgArr == -1]=0
  mplot.title('Zeros Term')
  mStr = "%0.4f" % m
  bStr = "%0.4f" % b
  mplot.figtext(0.2,0.25, r'\textbf{m = '+ mStr +'}', fontsize=14)
  mplot.figtext(0.2,0.2, r'\textbf{b = '+ bStr +'}', fontsize=14)
  
  #Try a fitting function for this curve:
  FitFuncOpt = 2
  if FitFuncOpt == 1:
    FitFunc = lambda p, x: p[0]*exp(-(x+p[1])/p[2])+p[3]
    ErrFunc = lambda p, x, y: y-FitFunc(p,x)
    p0 = [.3, -90, 90, 1.]
    p1, Success = optimize.leastsq(ErrFunc, p0[:], \
                    args = (ReadTimes[GP], FitC1AvgArr[GP]), full_output=0,\
                    Dfun = SpatialSpreadJacobian4, maxfev=100000000, col_deriv=0,\
                    xtol = 1.e-14, gtol=1e-15, ftol=1e-15)
    mplot.plot(ReadTimes, p1[0]*exp(-(ReadTimes+p1[1])/p1[2])+p1[3], 'b')
  elif FitFuncOpt == 2:
    FitFunc = lambda p, x: p[0]*exp(-(x)/p[1])+p[2]
    ErrFunc = lambda p, x, y: y-FitFunc(p,x)
    p0 = [.3, 120, 1.]
    p1, Success = optimize.leastsq(ErrFunc, p0[:], \
                    args = (ReadTimes[GP], FitC1AvgArr[GP]), full_output=0,\
                    Dfun = SpatialSpreadJacobian3, maxfev=100000000, col_deriv=0,\
                    xtol = 1.e-14, gtol=1e-15, ftol=1e-15)
    mplot.plot(ReadTimes, p1[0]*exp(-(ReadTimes+10)/p1[1])+p1[2], 'b')

  elif FitFuncOpt == 3:
    FitFunc = lambda p, x: p[0]*exp(-(x)/p[1])+p[2]
    ErrFunc = lambda p, x, y: y-FitFunc(p,x)
    p0 = [.3, 90, 1.]
    p1, Success = optimize.leastsq(ErrFunc, p0[:], \
                    args = (ReadTimes[GP], FitC1AvgArr[GP]), full_output=0,\
                    xtol = 1.e-14, gtol=1e-15, ftol=1e-15)
    mplot.plot(ReadTimes, p1[0]*exp(-(ReadTimes)/p1[1])+p1[2], 'b')
 
  
