# SPIRE Line Fitting script example based on using the SpectrumExplorer 
# SpectrumFitterGUI
#
# This script was created by running a single line fit by hand in the GUI,
# then using the option in the GUI to save the procedure as a script. It
# was then extended to include other lines read from a catalogue.
#
# OT1 DP Workshop at ESAC 16-18 March 2011
#

# Read in the point source spectrum product
# Either from a FITS file (the end result of the User Reprocessing Script)
# pointSourceSds = fitsReader('1342189124_spectrum_point_unapod.fits')
# Or read from the observation context pool
obs = getObservation(1342189124, poolName="spireSpec1342189124_NGC7027")
# Using the syntax from the Product Viewer:
pointSourceSds = obs.refs["level2"].product.refs["HR_unapodized_spectrum"].product
# (or the simplified syntax!)
# pointSourceSds = obs.level2.getProduct("HR_unapodized_spectrum")

# Define some constants
from herschel.share.unit import Constant
k = Constant.K_BOLTZMANN.value
c = Constant.SPEED_OF_LIGHT.value
sourceVelocity = 20.0*1e3

# Get the satellite radial velocity from the metadata
radialVelocity = pointSourceSds.meta['radialVelocity'].value

# Open a file for the line fit results
fitResults = open('C:\\Users\\etp87\\Desktop\\DPWorkshop\\fit_results.dat','w')

# Read the line catalogue
name      = String1d([])
lineFreqs = Double1d([])
file  = open('C:\\Users\\etp87\\Desktop\\DPWorkshop\\spirelines.dat','r')
for line in file.readlines():
    if line[0] != "#":
        name.append(line.split()[0])
        lineFreqs.append(float(line.split()[1]))
file.close()

# Correct the catalogue frequencies for the source velocity (for initial guess)
lineFreqsCorr = lineFreqs*(1.-sourceVelocity/c)    

pl=PlotXY()
# loop over the two centre detectors
for detector in ["SLWC3", "SSWD4"]:
    spireSpectrum1d = pointSourceSds['0000'][detector].copy()
    # convert the wavescale to GHz
    gHzSpectrum = convertWavescale(ds=spireSpectrum1d, to="GHz")
    # correct to LSR using satellite radial velocity
    gHzSpectrum['frequency'].data = gHzSpectrum['frequency'].data*(1.-radialVelocity/c)
    # Plot the spectrum    
    pl.addLayer(LayerXY(gHzSpectrum['frequency'].data, gHzSpectrum['flux'].data, color=java.awt.Color.BLUE))
    #
    # Initialise the Fitter
    sf = SpectrumFitter(gHzSpectrum, 0, 1, False)
    # Specify fit engine (1 = LevenbergMarquardt,2 = Amoeba, 3 = Linear, 4 = MP, 5 = Conjugated Gradient).
    sf.useFitter(1)
    # Add the models. First, polynomial for continuum (+initial guess)
    Mc1 = sf.addModel('poly', [3.0], [514.0577074148017,-1.3697323948473303,0.0012781730567420267,-2.8923270579794556E-7])
    #
    # Select only lines within the band and add a model for each
    sel = lineFreqsCorr.where(lineFreqsCorr.gt(MIN(gHzSpectrum['frequency'].data)).and(lineFreqsCorr.lt(MAX(gHzSpectrum['frequency'].data))))    
    for idx in sel.toInt1d():
        line = lineFreqsCorr[idx]
        exec("M%i = sf.addModel('sinc', [10.0, line, (100*c*0.04)/(Math.PI*1e9)])"%idx)
        # fix the width of the line to the resolution
        exec("M%i.setFixed(Int1d([2]))"%idx)
        # Plot the rest frequencies (not shifted by source velocity)
        pl.addLayer(LayerXY(Double1d([lineFreqs[idx],lineFreqs[idx]]),Double1d([0,200])))
        pl.addAnnotation(Annotation(lineFreqs[idx], -50, name[idx], angle=90))
    #
    # Do the fit, calculate the residual
    sf.doGlobalFit()
    sf.residual()
    tm = sf.getTotalModel()
    #
    # Plot the total model, and residual
    pl.addLayer(LayerXY(gHzSpectrum['frequency'].data, tm, color=java.awt.Color.RED))
    pl.addLayer(LayerXY(gHzSpectrum['frequency'].data, gHzSpectrum['flux'].data-tm, color=java.awt.Color.GREEN))
    #
    # calculate line fluxes from the fit parameters
    for idx in sel.toInt1d():
        # line area in W/m2
        exec("area =  1e-26*M%i.getFittedParameters()[0]*Math.PI*1e9*M%i.getFittedParameters()[2]"%(idx,idx))
        exec("print name[idx], M%i.getFittedParameters(), area"%idx)
        exec("fitResults.write('%%s %%s %%s %%s \\n'%%(name[idx], lineFreqs[idx], M%i.getFittedParameters()[1], area))"%idx)

pl.xtitle="Frequency (GHz)"
pl.ytitle="Flux Density (Jy)"

fitResults.close()