##^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^Lance Simms, Stanford University 2009
##ReturnLinearity.py
##
##PURPOSE: 
##  To examine the linearity of the pixels in an HxRG detector and return 
##  the linearity curves. 
##
##INPUTS: 
##      FitsFileName: string
##          The full path to the fits file that contains the ramps
##      SlopeFileName: string
##          The full path to the slopefitted file corresponding to 
##          FitsFileName
##
##OUTPUTS:
##  MODE = 2
##  MedADU: arr
##    The median ADU values for the x axis
##  DeltasMed: arr
##    The median point to point differences for y axis
##  DeltasMean: arr
##    The mean point to point differences for y axis
##  DeltasCumMed: arr
##    The median cumulative differences for y axis
##  DeltasCumMean: arr
##    The mean cumulative differences for y axis
##
##KEYWORDS:
##  XStart: int
##    X start coordinate for region of interest
##  XStop: int
##    X stop coordinate for region of interest
##  YStart: int
##    Y start coordinate for region of interest
##  YStop: int
##    Y stop coordinate for region of interest
##  UseRef: int	      
##    A Boolean; 0-Don't subtract reference 1-Do subtract
##  Mode: int      
##    0 - Examine ratios of delta approximations
##    1 - Examine deviation of read from straight line
##    2 - Examine median of pixels
##  DarkSub: int       
##    0 - Don't subtract a dark
##    1 - Subtract a median dark with the same NReads
##  TvFlag: int         
##    0 - Don't plot
##    1 - Plot
##    2 - Plot Ramps and Deltas
##  AllReads: int
##    0 - Collect data only for read specified by ReadNum
##    1 - Collect data for all reads in all exposures
##
##CALLING SEQUENCE:
##  LinArray = ExamineLinearity.py [FitsFileName] [SlopeFileName] 
##
##EXAMPLES:
##  MedADU, DeltasMed, DeltasMean, DeltasCumMed, DeltasCumMean = ReturnLinearity('/nfs/slac/g/ki/ki04/lances/H2RG-32-147/ASIC/07Dec19/Flat/Flat_G_H2RG_SIPIN_30_Reads_Dec19_2007_18_37_56.fits','/nfs/slac/g/ki/ki04/lances/H2RG-32-147/ASIC/Reduced/07Dec13/Flat/Flat_G_H2RG_SIPIN_15_Reads_Dec13_2007_21_35_50_G_Slope.fits', YStart=1700, YStop=1900, XStart=1000, XStop=1400, Mode=2, UseRef=1, CFac=1.45)
##
execfile('/afs/slac/u/ki/lances/python/python_HxRG/HxRG_Setup.py')
execfile('/afs/slac/u/ki/lances/python/SPIEPlotSettings.py')
PlotDir = '/nfs/slac/g/ki/ki04/lances/LSST/KPNO/Latex/Thesis/Linearity/Figures/'

from HxRG_Class import *
from ReturnCentroid import *
from ReturnRadialAvg import *
from Djs_Iterstat import *
from ReadNoise import *
import TableIO
import csv
import pdb
import asciidata

def ReturnLinearity(FitsFileName, SlopeFileName, UseRef=0, \
    XStart=-1, XStop=-1, YStart=-1, YStop=-1, \
    DarkSub=0, Mode=2, TvFlag=0, Verbose=0, CFac=-1): 
  
  ##INITIAL SETUP
  HxRG  = HxRG_C(FitsFileName=FitsFileName)
  SHxRG = HxRG_C(FitsFileName=SlopeFileName)
  
  if XStart == -1:
    XStart  = 4
  if XStop  == -1:
    XStop   = HxRG.FNAxis1-5
  if YStart == -1:
    YStart  = 4
  if YStop  == -1:
    YStop   = HxRG.FNAxis2-5
  
  ##FIND THE MATCHING DARK
  if DarkSub == 1:
     HxRG.Get_Median_Dark(HxRG.NReads)
     DarkHDU = pyfits.open(HxRG.DarkName)
  
  ##GO THROUGH THE FILES AND PLOT THE PERSISTENCE
  if TvFlag > 0:
    mplot.close('all')
    Fig = mplot.figure(0)
  
  ##GET THE RAW IMAGE..........................................
  ##Get the slope and the slope header
  FitsHDU      = pyfits.open(FitsFileName)
  HxRG.Get_Raw_Header(FitsFileName)
  
  #Get the Bias Read and allow for dark subtraction
  Im = FitsHDU[0].section[:,:,:]+float(HxRG.BZero)
  if UseRef == 1:
    if CFac == -1:
      HxRG.CFac = 3.1
      HxRG.CFac = 1.5
      HxRG.CFac = 1.45
      if HxRG.DetStr == 'H2RG-001':
        HxRG.CFac = 2.0
    else:
      HxRG.CFac = CFac
    Im   = ReturnRSCube(Im, HxRG.CFac, BadCols=HxRG.BadCols)
  if DarkSub == 1:
    Dark = DarkHDU[0].section[0:,:,:]
    Im   = Im -Dark
  
  ##GET THE SLOPEFITTED IMAGE..........................................
  ##Get the slope and the slope header
  SHxRG.Get_Raw_Header(SlopeFileName)
  FitsHDU     = pyfits.open(SlopeFileName)
  SlopeIm     = FitsHDU[0].section[0,:,:]
  SlopeHeader = FitsHDU[0].header
  
  ##Find the full range of pixels being considered
  XRange      = XStop - XStart
  YRange      = YStop - YStart
  
  ##Create Arrays to Hold ALl of the points
  ADUDepth   = 40000
  SlopeFracs = zeros(ADUDepth, dtype=float32)
  SlopeNums  = zeros(ADUDepth, dtype=float32)
  
  ###MODES 0 AND 1 -- LOOK AT INDIVIDUAL SLOPES#########################################
  if Mode == 0 or Mode == 1:
    for Row in arange(YRange)+YStart:
      print 'Row: ' + str(Row)
      for Col in arange(XRange)+XStart:
        ThisSlope  = SlopeIm[Row, Col]
        ThisRamp   = Im[:, Row, Col]
        ThisRampBS = Im[:, Row, Col] - Im[1, Row, Col]
        ThisRampEs = ThisRampBS
        ThisSlope  = SlopeIm[Row, Col]
    
        #Calculate the slopes for individual points
        if HxRG.ElecStr == 'LEACH':
          ReadTimes = (frange(HxRG.NAxis3-1)+1)*SHxRG.ITime/HxRG.NAxis3
        else:
          ReadTimes = (frange(HxRG.NAxis3-1)+1)*HxRG.FrameTime
    
        if HxRG.RawObjectKey.find('Median') == -1:
          ThisRamp   = Im[:, Row, Col]
          ThisRampWB = Im[1:,Row, Col]
          ThisRampBS = Im[:, Row, Col] - Im[0, Row, Col]
          ThisRampEs = ThisRampBS
          ThisRampFZ = ThisRampBS
          Deltas = (ThisRampBS[1:]-ThisRampBS[0])/ReadTimes[:-1]
        else:
          ThisRamp   = Im[1:, Row, Col]
          ThisRampWB = Im[1:,Row, Col]+Im[0,Row,Col]
          ThisRampBS = Im[1:, Row, Col] 
          ThisRampEs = ThisRampBS
          ThisRampFZ = append([0], ThisRamp)
          Deltas = ThisRamp[:]/(ReadTimes[1:]-ReadTimes[0])
          
        SlopeErrors = Deltas/ThisSlope
        if Mode == 0:
          for ThisSlopeError, i in zip(SlopeErrors, arange(SlopeErrors.size)):
            if HxRG.RawObjectKey.find('Median') == -1:
              SlopeFracs[ThisRampBS[i+1]] += ThisSlopeError
              SlopeNums[ThisRampBS[i+1]]  += 1
              m = ThisSlope
            else:
              SlopeFracs[ThisRampBS[i]] += ThisSlopeError
              SlopeNums[ThisRampBS[i]]  += 1
              m = ThisSlope
  
          if TvFlag == 2:
            mplot.figure(1)
            mplot.clf()
            mplot.plot(ThisRamp)
            mplot.figure(2)
            mplot.clf()
            mplot.plot(Deltas)
            mplot.figure(0)
            mplot.clf()
            if HxRG.RawObjectKey.find('Median') == -1:
              mplot.plot(ThisRampEs[1:], SlopeErrors, 'ro')
            else:
              mplot.plot(ThisRampEs, SlopeErrors, 'ro')
            ylim(.8, 1.2)
            pdb.set_trace()
    
        elif Mode == 1:
            UpperLinLim = 33000. 
            pmask = ones(ThisRampWB.size)
            pmask[where(ThisRampWB > UpperLinLim)] = 0
            m,b,Error,CosE = Linefit(arange(len(ThisRampFZ)),ThisRampFZ, \
                                     pmask=pmask, TvFlag=TvFlag)
            Errors = (m*arange(len(ThisRampFZ)))-ThisRampFZ
            if Verbose == 1:
              print 'Slope Fit from IDL :' + str(ThisSlope)
              print 'Slope Fit from Pyraf :' + str(m)
            if TvFlag > 0 :
              mplot.figure(0)
              mplot.clf()
              mplot.plot(arange(len(ThisRampFZ)),ThisRampFZ, 'b+')
              mplot.plot(arange(len(ThisRampFZ)),m*arange(len(ThisRampFZ)),'b')
              mplot.figure(1)
              mplot.clf()
              mplot.plot(Errors)
              ylim(-100,500)
              pdb.set_trace()
  
    TheseSlopeFracs = SlopeFracs/SlopeNums

    return TheseSlopeFracs

  ###MODE 2 -- LOOK AT COLLECTION OF PIXELS########################################
  elif Mode == 2:
  
    MedADU   = zeros(HxRG.NAxis3-1, dtype=float32)
    MedSig   = zeros(HxRG.NAxis3-1, dtype=float32)
    MeanADU  = zeros(HxRG.NAxis3-1, dtype=float32)
    MedDiff  = zeros(HxRG.NAxis3-1, dtype=float32)
    BiasRead = Im[0, YStart:YStop, XStart:XStop]
    SlopeImReg = SlopeIm[ YStart:YStop, XStart:XStop]
    MeanS, SigS, MedS, ModeS, MaskS = Djs_Iterstat(SlopeIm, RejVal=-10000., SigRej=1.5)
  
    for ReadNum in arange(HxRG.NAxis3-1)+1:
      if HxRG.RawObjectKey.find('Median') == -1:
        ThisRead = Im[ReadNum, YStart:YStop, XStart:XStop] - BiasRead
        ThisDiff = Im[ReadNum,   YStart:YStop, XStart:XStop] - \
                   Im[ReadNum-1, YStart:YStop, XStart:XStop]
      else:
        ThisRead = Im[ReadNum, YStart:YStop, XStart:XStop]
        if ReadNum == 1:
          ThisDiff = Im[ReadNum, YStart:YStop, XStart:XStop]
        if ReadNum == 1:
          ThisDiff = Im[ReadNum, YStart:YStop, XStart:XStop]-\
                     Im[ReadNum-1, YStart:YStop, XStart:XStop]
  
      Mean, Sig, Med, Mode, Mask = Djs_Iterstat(ThisRead, RejVal=-10000., SigRej=1.5)
      MeanD, SigD, MedD, ModeD, MaskD = Djs_Iterstat(ThisDiff, RejVal=-10000., SigRej=1.5)
      MedSig[ReadNum-1]  = Sig
      MedADU[ReadNum-1]  = Med
      MeanADU[ReadNum-1] = Mean
      MedDiff[ReadNum-1] = MedD

  #Create Arrays for point to point difference   
  DeltasMed      = zeros(HxRG.NAxis3-1,dtype=float32)
  DeltasMean     = zeros(HxRG.NAxis3-1,dtype=float32)
  DeltasMed[0]   = MedADU[0]
  DeltasMean[0]  = MeanADU[0]
  DeltasMed[1:]  = MedADU[1:] - MedADU[0:-1]
  DeltasMean[1:] = MeanADU[1:] - MeanADU[0:-1]

  if TvFlag > 0:
    mplot.figure(0)
    mplot.clf()
    mplot.plot(MedADU[0:-1],  DeltasMed[:-1], 'b')
    mplot.plot(MedADU[0:-1],  DeltasMed[:-1], 'bo')
    mplot.plot(MeanADU[0:-1], DeltasMean[:-1])
    mplot.plot(MedADU[0:-1],  MedDiff[:-1], 'ro')
    mplot.figure(1)
    mplot.clf()

  #Create Arrays for cummulative differences
  if HxRG.RawObjectKey.find('Median') == -1:
    DeltasCumMed  = (MedADU[:-1])
    DeltasCumMean = (MeanADU[:-1])
  else:
    DeltasCumMed  =  MedADU[1:] 
    DeltasCumMean =  MeanADU[1:]  
  
  DeltasCumMed  = DeltasCumMed/((HxRG.FrameTime)*(arange(DeltasCumMed.size)+1))
  DeltasCumMean = DeltasCumMean/((HxRG.FrameTime)*(arange(DeltasCumMean.size)+1))
  
  if TvFlag > 0:
    mplot.figure(1)
    mplot.clf()
    mplot.plot(MedADU[0:-1], DeltasCumMed/(DeltasCumMed[0:3].mean()), 'b')
    mplot.plot(MedADU[0:-1], DeltasCumMed/(DeltasCumMed[0:3].mean()), 'bo')
    mplot.plot(MedADU[0:-1], DeltasCumMean/(DeltasCumMean[0:3].mean()))
    ylim(.8,1.1)
    mplot.xlabel(r'\textbf{Mean Signal ($e^{-}$ in well)}')
    mplot.ylabel(r'\textbf{Fractional Countrate}')
  
  return MedADU, DeltasMed, DeltasMean, DeltasCumMed, DeltasCumMean