#
# HIPE HCSS
# 
# Spectroscopic convolution script: Convolve a selected map1 made around a certain line with
# an appropriate kernel to the PSF of a map2 at another band/wavelength and regrid.
#
# Rutger De Nutte
#

from javax.swing import *
from bisect import bisect_left as bisect

# Load the maps
print 'Input two line maps, convolution with kernel: map1 -> map2'

# Prompt for map1
print 'Now selecting map1; the map to be convolved.'
chooseMap1 = JFileChooser()
filter1 = filechooser.FileNameExtensionFilter("fits files", ["fits"])
chooseMap1.addChoosableFileFilter(filter1)
map1full = chooseMap1.showDialog(None, 'Please select map1')

if map1full == JFileChooser.APPROVE_OPTION:
    map1id = chooseMap1.getSelectedFile().getAbsolutePath()
else:
    print 'ERROR. Please enter a valid option.'
    raise SystemExit(0) 

map1 = simpleFitsReader(file = map1id)
print 'Map1 successfully selected as ' + map1id + '.'

# Prompt for map2
print 'Now selecting map2; the map to which map1 is to be convolved to.'
chooseMap2 = JFileChooser()
filter2 = filechooser.FileNameExtensionFilter("fits files", ["fits"])
chooseMap2.addChoosableFileFilter(filter2)
map2full = chooseMap2.showDialog(None, 'Please select map2')

if map2full == JFileChooser.APPROVE_OPTION:
    map2id = chooseMap2.getSelectedFile().getAbsolutePath()
else:
    print 'ERROR. Please enter a valid option.'
    raise SystemExit(0) 

map2 = simpleFitsReader(file = map2id)
print 'Map2 successfully selected as ' + map2id + '.'

# Determine the spectroscopic bands and lines the observations were made in.
# Create a dictionary containing the available band/line combinations for premade kernels.
# Give the option to automatically select the closest kernel online
# or create a Gaussian kernel.
print 'Successfully selected both maps. Now reading the metadata for map bands and lines.'

try:
    band1 = map1.meta['band'].value
except (ValueError, IndexError):
    band1 = str(raw_input('Could not find the spectral band in the header file. \
    Please insert the band for map1 manually (PACS: B2A, B2B, B3A, R1).'))

try:
    band2 = map2.meta['band'].value
except (ValueError, IndexError):
    band2 = str(raw_input('Could not find the spectral band in the header file. \
    Please insert the band for map2 manually (PACS: B2A, B2B, B3A, R1).'))

try:
    line1 = float(map1.meta['Line 1'].value.split(' micron')[0])
except (ValueError, IndexError):
    line1 = raw_input('Could not find the spectral line in the header file. \
    Please insert the wavelength for which for map1 was created manually (in micron).')

try:
    line2 = float(map2.meta['Line 1'].value.split(' micron')[0])
except (ValueError, IndexError):
    line2 = raw_input('Could not find the spectral line in the header file. \
    Please insert the wavelength for which for map2 was created manually (in micron).')

possibleKernelSelect = ['Closest premade kernel (online)', 'Create Gaussian kernel', 'User input kernel', ]
kernelSelect = JOptionPane.showInputDialog(None, 'Please select the kind of kernel to use', \
"Convolution", JOptionPane.QUESTION_MESSAGE, None, possibleKernelSelect, possibleKernelSelect[0])

if kernelSelect == 'Closest premade kernel (online)':
    print 'Automatically selecting proper convolution kernel.'
    kernelName = '_aniano'

    kernelLinesB2A = [55, 62, 68]
    kernelLinesB2B = [73, 75, 84, 94]
    kernelLinesB3A = [68]
    kernelLinesR1 = [110, 125, 136, 145, 150, 168, 187]
    kernelDict = {"B2A" : kernelLinesB2A, "B2B" : kernelLinesB2B, 'B3A' : kernelLinesB3A, \
    'R1' : kernelLinesR1}
    kernelLine1 = int(min([(abs(line1-x), x) for x in kernelDict[band1]])[1])
    kernelLine2 = int(min([(abs(line2-x), x) for x in kernelDict[band2]])[1])

    # Load the kernel from the web.
    kernelRoot = 'http://www.astro.princeton.edu/~ganiano/Kernels/Ker_2012_April_PACS_Spec/Kernels_fits_Files/Low_Resolution/'
    kernelFile = 'Kernel_LoRes_PACS_S_%s%i_to_PACS_S_%s%i.fits.gz' % (band1, kernelLine1, band2, kernelLine2)
    url = java.net.URL(kernelRoot + kernelFile)
    fitsStream = url.openStream()
    gzipStream = java.util.zip.GZIPInputStream(fitsStream)
    fa = FitsArchive(reader=FitsArchive.STANDARD_READER)
    print 'Loading ' + kernelFile + ' from the internet...'
    kernel = fa.load(gzipStream)
    gzipStream.close()
    fitsStream.close()

    # Convert the kernel to a SimpleImage.
    kernelImage = SimpleImage()
    kernelImage.image = kernel['PrimaryImage'].data
    kernelImage.wcs.cd1_1 = kernel.meta['cd1_1'].value
    kernelImage.wcs.cd1_2 = kernel.meta['cd1_2'].value
    kernelImage.wcs.cd2_1 = kernel.meta['cd2_1'].value
    kernelImage.wcs.cd2_2 = kernel.meta['cd2_2'].value
    kernelImage.wcs.crpix1 = kernel.meta['crpix1'].value
    kernelImage.wcs.crpix2 = kernel.meta['crpix2'].value
    kernelImage.wcs.crval1 = kernel.meta['crval1'].value
    kernelImage.wcs.crval2 = kernel.meta['crval2'].value
    kernelImage.wcs.ctype1 = kernel.meta['ctype1'].value
    kernelImage.wcs.ctype2 = kernel.meta['ctype2'].value
    kernelImage.wcs.epoch = kernel.meta['equinox'].value 
    kernelImage.meta = kernel.meta
    print 'Successfully loaded kernel ' + kernelFile + '.'

    # Rescale and resize the kernel to match map1 before convolving.
    print 'Rescaling and resizing the kernel to match map1 dimensions.'
    scalingKernelx = ABS(map1.wcs.cdelt1 / kernelImage.wcs.cdelt1)
    scalingKernely = ABS(map1.wcs.cdelt2 / kernelImage.wcs.cdelt2)
    kernelRegridSize = [int(round(kernelImage.image.dimensions[0]/scalingKernelx)), \
        int(round(kernelImage.image.dimensions[1]/scalingKernely))]
    kernelRegrid = Regrid(kernelRegridSize)(kernelImage.image[:,:])
    kernelRot = map1.copy()
    kernelRot.image[:,:] = 0.
    mXk = kernelRot.image.dimensions[0]
    mYk = kernelRot.image.dimensions[1]
    nXk = kernelRegrid.dimensions[0]
    nYk = kernelRegrid.dimensions[1]

    if mXk > nXk and mYk > nYk:
        kernelRot.image[mXk/2-nXk/2:mXk/2-nXk/2+nXk, mYk/2-nYk/2:mYk/2-nYk/2+nYk] = kernelRegrid
    else:
        kernelRot.image = kernelRegrid[nXk/2-mXk/2:nXk/2-mXk/2+mXk, nYk/2-mYk/2:nYk/2-mYk/2+mYk]

elif kernelSelect == 'Create Gaussian kernel':
    print 'Downloading beam information, interpolating for line wavelength and creating Gaussian kernel.'
    kernelName = '_gauss'

    # Create kernel WCS.
    kerWcs   = map1.wcs.copy()
    N1 = kerWcs.getNaxis1()
    N2 = kerWcs.getNaxis2()
    x0 = N1/2
    y0 = N2/2
    x1 = Double1d.range(N1)
    x2 = Double1d.range(N2)

    # Create Gaussian kernel.
    # Look up beam information on the calTree, interpolate to correct wavelengths from the two closest ones.
    calTree = getCalTree()
    w = calTree.spectrometer.beamSize["beamSize"]["wavelength"].data
    fwhm = calTree.spectrometer.beamSize["beamSize"]["fwhm"].data

    i1 = bisect(w, line1)
    if i1 == len(w):
        i1 -= 1
    if i1 == 0:
        i1 += 1
    w1min = w[i1-1]
    w1max = w[i1]
    fwhm1min = fwhm[i1-1]
    fwhm1max = fwhm[i1]
    fwhm1 = fwhm1min + (line1 - w1min) * (fwhm1max - fwhm1min) / (w1max - w1min)
    sigma1 = ABS(fwhm1 / (SQRT(8. * LOG(2)) * 3600. * kerWcs.getCdelt1()))

    i2 = bisect(w, line2)
    if i2 == len(w):
        i2 -= 1
    if i2 == 0:
        i2 += 1
    w2min = w[i2-1]
    w2max = w[i2]
    fwhm2min = fwhm[i2-1]
    fwhm2max = fwhm[i2]
    fwhm2 = fwhm2min + (line2 - w2min) * (fwhm2max - fwhm2min) / (w2max - w2min)
    sigma2 = ABS(fwhm2 / (SQRT(8. * LOG(2)) * 3600. * kerWcs.getCdelt1()))

    # Determine the kernel sigma as the square difference of the beams sigma.
    # Fill the kernel array as a Gaussian with correct normalization. 
    sigma = SQRT(sigma2**2 - sigma1**2)
    ampl  = 1. / (2. * Math.PI * sigma**2)
    y1    = EXP(-0.5 * SQUARE((x1 - x0) / sigma))
    y2    = EXP(-0.5 * SQUARE((x2 - y0) / sigma))
    y2d   = Double2d(N2, N1)
    for k in range(N2):
        for i in range(N1):
            y2d.set(k, i, ampl * y1[i] * y2[k])

    # Create kernel and set WCS and image data.
    kernelRot = SimpleImage()
    kernelRot.setImage(y2d)
    kernelRot.setWcs(kerWcs)

elif kernelSelect == 'User input kernel':
    # Load the kernel as a simpleImage.
    print 'Loading the custom kernel.'
    kernelName = '_custom'
    chooseKernel = JFileChooser()
    filter3 = filechooser.FileNameExtensionFilter("fits files", ["fits"])
    chooseKernel.addChoosableFileFilter(filter3)
    kernelfull = chooseKernel.showDialog(None, 'Please select the kernel.')

    if kernelfull == JFileChooser.APPROVE_OPTION:
        kernelid = chooseKernel.getSelectedFile().getAbsolutePath()
    else:
        print 'ERROR. Please enter a valid option.'
        raise SystemExit(0) 

    bigKernel = simpleFitsReader(file = kernelid)

    # Rotate the kernel
    print 'Rotating the selected kernel to match the observation scan angle.'
    kernelRotAngle = float(raw_input('Input the angle (in decimal degrees) by which the\
    selected kernel is to be rotated (from the observation scan angle).'))
    kernelRot = rotate(image = bigKernel, angle = kernelRotAngle, subsampleBits = 32,\
    interpolation = rotate.INTERP_BICUBIC)

else:
    print 'ERROR. Please enter a valid option.'
    raise SystemExit(0)

# Replace NaN values in kernel with 0. to facilitate convolution.
w2 = kernelRot.image.where(IS_FINITE(kernelRot.image).not())
kernelRot.image[w2] = 0.

# Normalize the kernel and do the convolution.
# Substract the background from map1 before convolving to minimize edge effects.
print 'Performing the convolution using the rescaled, rotated and normalized kernel.'
kernelRot.image = kernelRot.image / SUM(kernelRot.image)
map1NoBg = map1.copy()
w3 = map1NoBg.image.where(map1NoBg.image != 0.)
bg = MEDIAN(map1NoBg.image[w3])
map1NoBg.image[w3] -= bg

convMap1 = imageConvolution(image = map1NoBg, kernel = kernelRot)
convMap1.image[w3] += bg
convMap1.meta = map1.meta
convMap1.wcs = map1.wcs

# Regrid the convolved map to match map2 size/scale
print 'Regridding the convolved map1 to match the WCS of map2.'
convMap1Regrid = regrid(source=convMap1, target=map2)
convMap1Regrid.meta = map1.meta
print 'Convolution script executed successfully!'

# Save the convolved map to disk.
fits = FitsArchive()
convid = map1id.rstrip('.fits') + '_convTo_' + band2 + str(int(line2)) + kernelName + '.fits'
fits.save(convid, convMap1Regrid)
print 'Saved the convolved map1 to ' + convid + '.'