#Demonstration on how to apply the semi-extended source correction, in addition to point-source calibration,
#to a PACS SED spectrum of a semi-extended source (spatial extent < ~25") with specExtendToPointCorrection task
#
#Spectrometer Workshop: PACS
#23 April 2014
#Created by E. Puga


import java
import herschel
from herschel.pacs.spg.spec.beams import ExtendedSource
from math import *
from herschel.ia.gui.plot.renderer.PlotLegendEngine import Position

###################################################################
#Classes and functions
###################################################################

#Gaussian Model definition (Circular Gaussian with offset wrt centre of central spaxel)
class MySource(ExtendedSource):
    sigma = 0.0
    offRa = 0.0
    offDec = 0.0
    def __init__(self,diameter, offRa, offDec):
        self.sigma = diameter/2.3548
        self.offRa = offRa
        self.offDec = offDec
       
    def getFlux(self,relRa,relDec):
        flux = exp(-(0.5*((relRa-self.offRa)/self.sigma)**2+0.5*((relDec-self.offDec)/self.sigma)**2))
        return flux

    def renderGauss(self, size=109, dpix=0.5):
        self.image=Double2d(size, size)
        for y in range( size ):
           for x in range( size ):
             xim = (-dpix)*(x-(size/2+0.5))
             yim = (dpix)*(y-(size/2+0.5))
             self.image.set(y,x, source.getFlux(xim, yim))
        return self.image

    def getSourceModelImage(self, size=109, dpix=0.5):
        return self.renderGauss(size, dpix)

def flagSet(band,gd,flag,wve):
    flg=flag.copy()
    if (gd==1):
        if (band == "B3A"):
            sey=[55.0,70.0]
        elif (band == "B2A"):
            sey=[55.0,73.0]
        elif (band == "B2B"):
            sey=[70.0,95.0]
        else:
            sey=[102.,190.]
        idx=wve.where(wve < sey[0])
        flg[idx]= 1.0
        idx=wve.where(wve > sey[1])
        flg[idx]= 1.0
    else:
        sey=ABS(len(gd)/2.) - len(gd)/2
        if ((len(gd) > 0) &(sey != 0.0)):
            print "Error: you have an odd number of ""good"" entries"
            return
        flg[:]= 1.0
        for i in range(0,len(gd),2):
            idx=wve.where((wve >= gd[i])&(wve <= gd[i+1]))
            if (idx.length() == 0):
                print "Error: your limits",gd[i],gd[i+1],"are outside the range of the spectrum; start again"
                flg=None
            else: 
                flg[idx]= 0.0
    return(flg)

def colorShade(col=java.awt.Color(128,128,128), npts=9):
    r=col.getRed()
    g=col.getGreen()
    b=col.getBlue()
    nr=(npts-1)/2
    steps=Double1d.range(nr)*(1./(nr+1))+(1./(nr+1))
    listShades= [java.awt.Color(int(r*step), int(g*step), int(b*step)) for step in steps]
    listTints= [java.awt.Color(int(r+(step*(255-r))), int(g+(step*(255-g))), int(b+(step*(255-b)))) for step in steps]
    listAll=[listShades, [col], listTints]
    listFinal=[num for elem in listAll for num in elem]
    return listFinal

#A utility
def ls(directory="./",pattern="",case=1,verbose=1):
	"""ls(directory="./",pattern="",case=1)
Returns the files in "directory" (relative path allowed) with filename containing "pattern" (incl. wildcards)
case	: 0/1 for case (in)sensitive listing (prefer case = 1 for pattern filtering)
"""
	import sys,os,fnmatch
	if case:
		allfiles = SORT(String1d(os.listdir(directory)))
	else:
		allfiles = String1d([file.lower() for file in SORT(String1d(os.listdir(directory)))])
	files = allfiles
	if pattern:
		files = [file for file in allfiles if fnmatch.fnmatch(file, pattern)]
	if verbose: 
		for file in files: print file
	return files


###################################################################
#Data, source name and Model Parameters                 
###################################################################
path= '/home/epuga/PACS_ICC/PACS_TelBackNorm/calset64/'
sourceName='NGC7027'
model='Gaussian' 
diam, offRa, offDec=12.0, 0., 0. #in arcsecs

#Model Visualization
source = MySource(diam, offRa, offDec)
Display(SimpleImage(image=source.getSourceModelImage(109, 0.5)))


#obsidList=[1342186968,1342186969]#NGC7027 According to Pierre, this is the real SED, but it was so early in time (OD182), that there is no orderSel
# Chose the obsids
obsids = {}
obsids[1342186968] = "SEDA" 
obsids[1342186969] = "SEDB"
#Each pool has been separated in red and blue

###################################################################

calTree=getCalTree()


#Load FinalCubes and concatenate them
cubes=[]
for obsid in obsids.keys():
   for camera in ['blue','red']:
      pattern='*'+ str(obsid) +'*'+ camera+'*'+ '_FinalCubes'
      files=ls(directory=path,pattern = pattern, case=1, verbose=0)
      for file in files:
         print file
         cubes.append(readSliced(file, poolLocation=path, verbose=0))
cubes=concatenateSliced(cubes)

#Plot of spectrum in central spaxel
x,y=2,2
pinit = plotCubes(cubes,[],x=x,y=y)

###################################################################
#Point Source Correction & E2P Correction
###################################################################

fullSpec = Spectrum1d()
segments = Int1d()
flags = Int1d()
fullSpec9 = Spectrum1d()
fullSpecCorr = Spectrum1d()
fullSpec9Corr = Spectrum1d()
fullCorrection1 = Spectrum1d()
fullCorrection9 = Spectrum1d()

verbose=0 # for extractCentralSpectrum
smoothing = 'wavelet' # or: smoothing = 'filter'
nLowFreq = 4

for slice in range(len(cubes.refs)):
  rcube = cubes.get(slice)
  band = rcube.meta['band'].value
  c1,c9,c129 = extractCentralSpectrum(rcube, slice=slice, smoothing=smoothing,
                nLowFreq=nLowFreq, calTree=calTree, verbose=verbose)
  
  if (model == 'Gaussian'):
     # in the case of a Gaussian, it is the FWHM
     source = MySource(diam, offRa, offDec)
     central3x3=False
     e2p1 = specExtendedToPointCorrection(c1,calTree=calTree,userSource=source,central3x3=central3x3, extra=True)   
     central3x3=True
     e2p9 = specExtendedToPointCorrection(c9,calTree=calTree,userSource=source,central3x3=central3x3, extra=True)   
  else:  
     #specExtendedToPointCorrection this assumes a disc top hat
     central3x3=False
     e2p1 = specExtendedToPointCorrection(c1,calTree=calTree,srcRadX=diam/2.,srcRadY=diam/2.,srcAngle=0.0,raOffset=0.0,decOffset=0.0,central3x3=central3x3,extra=True)
     central3x3=True
     e2p9 = specExtendedToPointCorrection(c9,calTree=calTree,srcRadX=diam/2.,srcRadY=diam/2.,srcAngle=0.0,raOffset=0.0,decOffset=0.0,central3x3=central3x3,extra=True)

  c1e2pCorr = divide(ds1=c1, ds2=e2p1, overwrite=False)
  c9e2pCorr = divide(ds1=c9, ds2=e2p9, overwrite=False)

  # Concatenate all bands (slices) to one spectrum1d, keeping them as separate segments
  spec = c1.spectrum1d #c1 is a SimpleSpectrum and we convert it to Spectrum1d
  fullSpec.concatenate(spec)
  segments.append(Int1d(spec.flux.length(), slice))
  flg = Int1d(spec.wave.length(), 0)
  flg=flagSet(band,1,flg,spec.wave)
  flags.append(flg)
  spec = c1e2pCorr
  fullSpecCorr.concatenate(spec)
  spec9 = c9.spectrum1d
  fullSpec9.concatenate(spec9)
  spec9 = c9e2pCorr
  fullSpec9Corr.concatenate(spec9)
  corr = e2p1.spectrum1d
  fullCorrection1.concatenate(corr)
  corr = e2p9.spectrum1d
  fullCorrection9.concatenate(corr)

# Include the column segment and flags in the final Spectrum1d
fullSpec.setSegment(segments)
fullSpec.setFlag(flags)
fullSpec9.setSegment(segments)
fullSpec9.setFlag(flags)
fullSpecCorr.setSegment(segments)
fullSpecCorr.setFlag(flags)
fullSpec9Corr.setSegment(segments)
fullSpec9Corr.setFlag(flags)
fullCorrection1.setSegment(segments)
fullCorrection1.setFlag(flags)
fullCorrection9.setSegment(segments)
fullCorrection9.setFlag(flags)

openVariable("fullSpec", "Spectrum Explorer")

pC1=PlotXY()

npt=fullSpec.getSegmentCount()
col=java.awt.Color.ORANGE
colShade=colorShade(col,npt+1) 
for i in range(npt):
   try:
      wave,flux,flag = fullSpec.getSegment(i).wave, fullSpec.getSegment(i).flux, fullSpec.getSegment(i).flag
      idx= flag.where(flag == 0) 
      if i == npt/2:
         pC1.addLayer(LayerXY(wave[idx], flux[idx], line=1, color=colShade[i], inLegend=True, name="C1"))
      else:
         pC1.addLayer(LayerXY(wave[idx], flux[idx], line=1, color=colShade[i], inLegend=False))
   except:
      pass
col=java.awt.Color.RED
colShade=colorShade(col,npt) 
for i in range(npt):
   try:
      wave,flux,flag = fullSpecCorr.getSegment(i).wave, fullSpecCorr.getSegment(i).flux, fullSpecCorr.getSegment(i).flag
      idx= flag.where(flag == 0) 
      if i == npt/2:
         pC1.addLayer(LayerXY(wave[idx], flux[idx], line=1, color=colShade[i], inLegend=True, name="C1/e2p1"))
      else:
         pC1.addLayer(LayerXY(wave[idx], flux[idx], line=1, color=colShade[i], inLegend=False))
   except:
      pass
col=java.awt.Color(255,0, 255)
colShade=colorShade(col,npt) 
for i in range(npt):
   try:
      wave,flux,flag = fullSpec9.getSegment(i).wave, fullSpec9.getSegment(i).flux, fullSpec9.getSegment(i).flag
      idx= flag.where(flag == 0) 
      if i == npt/2:
         pC1.addLayer(LayerXY(wave[idx], flux[idx], line=1, color=colShade[i], inLegend=True, name="C9"))
      else:
         pC1.addLayer(LayerXY(wave[idx], flux[idx], line=1, color=colShade[i], inLegend=False))
   except:
      pass
col=java.awt.Color(127,0, 255)
colShade=colorShade(col,npt) 
for i in range(npt):
   try:
      wave,flux,flag = fullSpec9Corr.getSegment(i).wave, fullSpec9Corr.getSegment(i).flux, fullSpec9Corr.getSegment(i).flag
      idx= flag.where(flag == 0) 
      if i == npt/2:
         pC1.addLayer(LayerXY(wave[idx], flux[idx], line=1, color=colShade[i], inLegend=True, name="C9/e2p9"))
      else:
         pC1.addLayer(LayerXY(wave[idx], flux[idx], line=1, color=colShade[i], inLegend=False))
   except:
      pass
pC1.legend.visible=1
pC1.legend.position = Position.RIGHTMIDDLE
pC1.titleText=sourceName
pC1.setXtitle("Wavelength ($\micro$m)")
pC1.setYtitle("Observed Flux (Jy)")
pC1.xaxis.tick.gridLines = 1
pC1.yaxis.tick.gridLines = 1
pC1.xaxis.range=[50, 192]
pC1.setPlotSize(8, 3.0)
c = 299792458.0 # m/s
xaux = ReciprocalAuxAxis(1.0e-3*c)
pC1.getXaxis().removeAuxAxis(0)
pC1.getXaxis().addAuxAxis(xaux)
xaux.setTitleText("Frequency (GHz)")
pC1.xaxis.type=Axis.LOG
pC1.yaxis.type=Axis.LOG

#Display corrections E2P1 and E2P9 in plot
pCor=PlotXY()
npt=fullSpec.getSegmentCount()
for i in range(npt):
   try:
      wave1,flux1,flag1 = fullCorrection1.getSegment(i).wave, fullCorrection1.getSegment(i).flux, fullCorrection1.getSegment(i).flag
      wave9,flux9,flag9 = fullCorrection9.getSegment(i).wave, fullCorrection9.getSegment(i).flux, fullCorrection9.getSegment(i).flag
      idx= flag1.where(flag1 == 0) 
      if i == 0:
         pCor.addLayer(LayerXY(wave1[idx], flux1[idx], line=1, color=java.awt.Color.BLUE, inLegend=True, name="e2p1"))
         pCor.addLayer(LayerXY(wave9[idx], flux9[idx], line=1, color=java.awt.Color.RED, inLegend=True, name="e2p9"))
      else:
         pCor.addLayer(LayerXY(wave1[idx], flux1[idx], line=1, color=java.awt.Color.BLUE, inLegend=False))     
         pCor.addLayer(LayerXY(wave9[idx], flux9[idx], line=1, color=java.awt.Color.RED, inLegend=False))
   except:
      pass
pCor.legend.visible=1
pCor.legend.position = Position.RIGHTMIDDLE
pCor.titleText='Source: %s, diam: %.1f \"' % (sourceName, diam)
pCor.setXtitle("Wavelength ($\micro$m)")
pCor.setYtitle("Factor")
pCor.xaxis.tick.gridLines = 1
pCor.yaxis.tick.gridLines = 1
pCor.xaxis.range=[50, 192]
pCor.setPlotSize(8, 3.0)

#Tool:
#Concentration plot to assess if line and continuum emission come from same angular position on the sky.
#Only useful for strong continuum sources
def pacsRingCubes (rcube, sourceName, pr=[], ratio=True):
    if not pr:
      pr=PlotXY()
    wc = rcube.wave
    fc = rcube.flux[:,2,2]
    f9 = StatWithNaN(MEAN)(StatWithNaN(MEAN)(rcube.flux[:,1:4,1:4],1),1)
    f8 = ((f9*9.) - fc) / 8.
    f25 = StatWithNaN(MEAN)(StatWithNaN(MEAN)(rcube.flux,1),1)
    f16 = ((f25*25.) - f9*9.) / 16.
    
    if ratio == True:
       if pr:
          pr.addLayer(LayerXY(wc, fc/(StatWithNaN(MAX)(fc))*100., color=java.awt.Color.black, name='Central',stroke=1))
          pr.addLayer(LayerXY(wc,f8/fc*100., color=java.awt.Color.blue, name='Ring 1',stroke=1))
          pr.addLayer(LayerXY(wc,f16/fc*100., color=java.awt.Color.red, name='Ring 2',stroke=1))
       else:
          pr.addLayer(LayerXY(wc, fc/(StatWithNaN(MAX)(fc))*100., color=java.awt.Color.black,stroke=1))
          pr.addLayer(LayerXY(wc,f8/fc*100., color=java.awt.Color.blue,stroke=1))
          pr.addLayer(LayerXY(wc,f16/fc*100., color=java.awt.Color.red,stroke=1))
       pr.ytitle="Ratio over central spaxel (%)"
    else:
       if pr:
          pr.addLayer(LayerXY(wc, fc, color=java.awt.Color.black, name='Central',stroke=1))
          pr.addLayer(LayerXY(wc,f8, color=java.awt.Color.blue, name='Ring 1',stroke=1))
          pr.addLayer(LayerXY(wc,f16, color=java.awt.Color.red, name='Ring 2',stroke=1))
       else:
          pr.addLayer(LayerXY(wc, fc, color=java.awt.Color.black,stroke=1))
          pr.addLayer(LayerXY(wc,f8, color=java.awt.Color.blue,stroke=1))
          pr.addLayer(LayerXY(wc,f16, color=java.awt.Color.red,stroke=1))
       pr.ytitle="Surface Brightness (Jy/pixel)"
    pr.xtitle="Wavelength ($\micro$m)"
    pr.titleText = sourceName
    pr.legend.visible=1
    
    return pr

pC=pacsRingCubes(rcube, sourceName, ratio=True)
for slice in range(len(cubes.refs)):
  rcube = cubes.get(slice)
  pacsRingCubes(rcube, sourceName, pC, ratio=True)

#pC.saveAsPNG("")


#ADDING PHOTOMETRY
#This table with the PACS photometry was created with the script addPhot.py, also shown in the Spectrometer Workshop PACS hands-on sessions 
table=asciiTableReader(file='/home/epuga/PACS_ICC/SpecWorkshop/PACSPhoto.tbl')
sourceNamesP = table['Source'].data
bandP = table['Band'].data
fluxP = table['Flux'].data
efluxP = table['eFlux'].data

#Adding photometry to plot pC1
idxp=sourceNamesP.where(sourceNamesP == sourceName)
pC1.addLayer(LayerXY(bandP[idxp], fluxP[idxp], errorY=[efluxP[idxp], efluxP[idxp]], line=0, symbol=Style.FSQUARE, color=java.awt.Color.blue, name='PACS Phot'))