# This jython script is not meant to be run end to end.
# In the hands on demo, we will use flagging routines to clean some 
# spectral scan data, then run the decon routine on it.


#####################################
#
#  DEFINE SOME UTILITY FUNCTIONS

def clearAllFlags(htp):
  # This routine clears all flags in all datasets in a HTP.
  # We use this to simplify some training, since the SPG managed
  # to flag spurs in our training datasets already.
  #
  SPUR=HifiMask.SPUR_CANDIDATE
  # inspect summary table to get the index of all data, then loop through them
  for htpindex in htp["summary"]["dataset"].data :
    spectrum= htp.get(htpindex)  # fetch the spectrum from the HTP
    spectrumType=spectrum.meta["sds_type"].value # determine type (COMB, HC, SCIENCE, TUNE)
    nrows=spectrum.getLength()  # get number of spectrum in this dataset
    rflags=spectrum["rowflag"].data
    rflags=rflags*0
    spectrum["rowflag"].data=rflags
    for dsidx in range(nrows) :  # loop through them
      for sb in (spectrum.getSegmentIndices()): # loop through the subbands
        segment = spectrum.getPointSpectrum(dsidx).getSegment(sb)
        flag=segment.getFlag() # fetch the flags
        flag=SPUR.unset(flag)  # clear the spur flag
        segment.setFlag(flag)  # reinsert into the spectrum
    htp.set(htpindex,spectrum,spectrumType) # reinsert modified spectra back into the HTP
  return htp




def flagData(htp,htpindex,freq1,freq2) :
   # This routine takes an htp, and flags a range of frequencies
   # as a spur.
   # The user needs to enter the index of the dataset, and the lower
   # and upper frequency range.
   #
   SPUR=HifiMask.SPUR_CANDIDATE
   #
   spectrum=htp.get(htpindex)
   spectrumType=spectrum.meta["sds_type"].value # determine type (COMB, HC, SCIENCE, TUNE)
   #
   nrows=spectrum.getLength()  # get number of spectrum in this dataset
   for dsidx in range(nrows) :  # loop through them
     for sb in (spectrum.getSegmentIndices()): # loop through the subbands
       segment = spectrum.getPointSpectrum(dsidx).getSegment(sb)
       flag=segment.getFlag() # fetch the flags
       freq=segment.getWave() # fetch the frequencies
       idx=freq.where((freq>=freq1).and(freq<=freq2))
       if idx.length()>0 :
         flag[idx]=SPUR.set(flag[idx])  # set the flag
         segment.setFlag(flag)  # reinsert into the spectrum
   htp.set(htpindex,spectrum,spectrumType) # reinsert modified spectra back into the HTP
   return htp

#  END OF FUNCTION DEFINITIONS
# 
#####################################


########################
#
#  START OF SCRIPT
#

obs=getObservation("1342181161",poolName="1342181161_sscan_dbs_1b_r4")

# extract the level 2tree
lev2=obs.getProduct("level2")
# from that, extract the HTP for the WBS-H backend
htpwbsh=lev2.getProduct("WBS-H-USB")
htpwbshlsb=lev2.getProduct("WBS-H-LSB")
# now clear all the flags, since the SPG already caught them.
htpwbsh=clearAllFlags(htpwbsh)
htpwbshlsb=clearAllFlags(htpwbshlsb)

# there is a slight bug in decon.
# normally we need only have to worry about flagging the USB track at level 2
# However the rowflags are being accidentally read from the LSB, so we
# need to clear the LSB rowflags to trick decon into thinking the data is good.
#
################ IMPORTANT   ###########
# At this stage, you have a htp that needs some cleaning.
# Inspect the data and identify those datasets affected by spurs
# Then using either the SpectrumExplorer or the clean function
# included here, flag the spurs, then execute the final few lines
# of the script to re-insert the data to the obs context, then run decon

# be sure to record the dataset index, and frequency range affected
htpwbsh=flagData(htpwbsh,32,589.5,590.0)
htpwbsh=flagData(htpwbsh,33,590.5,591.0)
htpwbsh=flagData(htpwbsh,34,591.5,592.0)
htpwbsh=flagData(htpwbsh,35,592.0,592.5)
htpwbsh=flagData(htpwbsh,36,593.0,593.5)
#########################################


# OK now with everything flagged, re-insert data to the context and run decon
# re-insert the HTP back into the observation context
lev2.setProduct("WBS-H-LSB",htpwbshlsb)
lev2.setProduct("WBS-H-USB",htpwbsh)
obs.setProduct("level2",lev2)
# run deconvolution

# We set the max_iterations to prevent the task from running too long
# in case the spurs are not flagged properly.

decon_result = doDeconvolution(obs=obs,max_iterations=20)