#hifidemo_scripting1_apr2010.py: demo of how to write scripts to
#    manipulate HIFI data in HIPE. Baseline subtraction using 
#    polynomial fitting is taken as
#    an example, since HIPE does not provide a convenient
#    way to do this, presently.
#
#
# v0.0 06/Apr/2010 ACAB: created for April 2010 HIFI DP Workshop
# v0.1 07/Apr/2010 ACAB: use getObservation
#
#

#Basics of Jython scripting is described in the "Scripting and 
#Data Mining Manual", accessible from the Hipe Help menu

#load band 4B SScan LChop observation 
#from ~/.hcss/lstore/
obs=getObservation("1342180554",poolName="pool_1342180554")

#extract WBS-H-USB product from ObservationContext
htp=obs.refs["level2"].product.refs["WBS-H-USB"].product

#This htp contains a collection of level 2 "spectrumdatasets".
#A single level 2 spectrumdataset, say number 60, can be selected as follows.
#(use the context viewer to find out which spectrumdatasets are available)
sds=htp.refs["60"].product["dataset"]

#A spectrumdataset can consist of multiple scans (read-outs), each of
#which can contain multiple 'segments' (WBS sub-bands or HRS modules).

#    -number of scans
scans=sds.rowCount
print scans
#    -array of segment indices, e.g. [1,2,3,4] for WBS bands 1-5.
segments=sds.getSegmentIndices()
print segments

#REST OF SCRIPT: Subtract polynomial fit of flux over entire IF from
#                individual sub-bands.

#The example spectrumdataset is a WBS spectrum with 4 sub-bands.
#If we want to work on the overall spectrum, 'stitch' them first.
#(If you don't how to do this, use GUI under Applicable Tasks and copy
# equivalent command from console) 
sds_s = stitch(ds=sds,variant="crossoverPoints",edgeTolerance=0.01,stepsize=0.0)

#Get most basic spectrum to work on: a single segment
#Since we 'stitched' the sub-bands above, this consists
#of one spectrum covering the entire IF.
segments_s=sds_s.getSegmentIndices()
print segments_s
spectrum_s = sds_s.getPointSpectrum(0).getSegment(segments_s[0])

#currently, the fitting tools cannot deal with NaNs,
#so remove those first
flux=spectrum_s.getFlux()
index=flux.where(IS_NAN)
flux[index]=MEDIAN(flux)
spectrum_s.setFlux(flux)

#initialize fitting task for polynomial  baseline fitting
sf = SpectrumFitter(spectrum_s,False)
                         
#Select which fitting method to use.
#The LevenbergMarquardt method is good for non-linear fits
#even if the initial guess is far away from the solution,
#i.e. for 'global' fits.
#Other options are 'linear' and 'amoeba'. The latter
#is better in finding the *absolute* minimum in chi^2.
sf.useFitter("LevenbergMarquardt")

# Initialize the polynomial model
polynomialOrder=2
initialGuesses = [0., 0., 0.]
model = sf.addModel("poly",[polynomialOrder],initialGuesses)

# To not include emission lines in the fit,
# set their weight to 0 and the rest to 1.
model.setWeight(0., 2000.0, 1.0)
model.setWeight(1115.5, 1116.0, 0.0)

#do the fitting:
sf.doGlobalFit()
sf.residual()

# Get the fit parameters and their standard deviations 
fitParams = sf.getResult()["Parameters"].data
print fitParams
fitErrors = sf.getResult()["StdDev"].data
print fitErrors

#plot original spectrum, baseline, residual, and masked areas
Wave=spectrum_s.getWave()
originalFlux = spectrum_s.getFlux()
residualFlux = sf.getResidual()
baselineFlux = model.getSynthetic()
channelMask = model.getMask()
p=PlotXY()
p.batch=1
ll=[]
l  = LayerXY(Wave, originalFlux)
ll.append(l)
l  = LayerXY(Wave, residualFlux)
ll.append(l)
l  = LayerXY(Wave, baselineFlux)
ll.append(l)
l  = LayerXY(Wave, channelMask)
ll.append(l)
p.layers=ll
p.batch=0

#So far we worked on a stitched spectrum. Now subtract baseline
#from original, non-stitched spectrum.

segments=sds.getSegmentIndices()
for p in segments:
 wavep=sds.getPointSpectrum(0).getSegment(p).getWave()
 basep=Polynomial(fitParams)(wavep)
 fluxp=sds.getPointSpectrum(0).getSegment(p).getFlux()-basep
 sds.getPointSpectrum(0).getSegment(p).setFlux(fluxp)

#put modified dataset back in htp
htp.setProduct("60", DatasetWrapper(sds))

#Since 'htp' is bound to 'obs', 'obs' has been changed as well.
#If 'htp' were generated with a deep copy, e.g. htp1=htp.copy(),
#one would put it back in 'obs' as follows.
#obs.level["level2"].setProduct("WBS-H-USB",htp)