#
# HIPE HCSS
# 
# Photometric Convolution script: Convolve a selected map1 at certain waveband with an appropriate 
# kernel to the PSF of a map2 at another waveband and regrid.
#
# Rutger De Nutte
#

from javax.swing import *

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

# Prompt for map1.
print 'Now selecting map1; the map to be convolved.'
possibleMap1select = ['local map', 'HSA observation', 'Scanamorphos map', ]
map1select = JOptionPane.showInputDialog(None, 'Please enter the map1 format', \
"Convolution", JOptionPane.QUESTION_MESSAGE, None, possibleMap1select, possibleMap1select[0])

if map1select == 'Scanamorphos map':
    # Load the Scanamorphos map from local disk.
    chooseMap1 = JFileChooser()
    filter1 = filechooser.FileNameExtensionFilter("fits files", ["fits"])
    chooseMap1.addChoosableFileFilter(filter1)
    cube1full = chooseMap1.showDialog(None, 'Please select map1')

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

    cube1 = simpleFitsReader(file = map1id)

    # Convert the Scanamorphos map to hipe readable map.
    map1 = SimpleImage()
    map1.description = 'Map from Scanamorphos - '+ cube1.meta['filter'].value
    map1.image = Double2d(cube1['PrimaryImage'].data[0,:,:])
    map1['image'].description = cube1.meta['PLANE0'].value
    map1.error = Double2d(cube1['PrimaryImage'].data[1,:,:])
    map1['error'].description = cube1.meta['PLANE1'].value
    map1.coverage = Double2d(cube1['PrimaryImage'].data[3,:,:])
    map1['coverage'].description = cube1.meta['PLANE3'].value
    map1.exposure = Double2d(cube1['PrimaryImage'].data[2,:,:])
    map1['exposure'].description = cube1.meta['PLANE2'].value
    map1.wcs.ctype1 = cube1.meta['ctype1'].value
    map1.wcs.ctype2 = cube1.meta['ctype2'].value
    map1.wcs.equinox = cube1.meta['equinox'].value
    map1.wcs.crpix1 = cube1.meta['crpix1'].value
    map1.wcs.crpix2 = cube1.meta['crpix2'].value
    map1.wcs.crval1 = cube1.meta['crval1'].value
    map1.wcs.crval2 = cube1.meta['crval2'].value
    map1.wcs.cdelt1 = cube1.meta['cdelt1'].value
    map1.wcs.cdelt2 = cube1.meta['cdelt2'].value
    map1.wcs.crota2 = cube1.meta['crota2'].value
    map1.meta = cube1.meta
    print 'Map1 successfully selected as a ' + map1select + '.'

elif map1select == 'HSA observation':
    # Acquire the observation from the Herschel Science Archive and select data level and waveband.
    obs1id = int(raw_input('Please enter observation id for map1 (e.g. 1342201436): '))
    instr1list = ['PACS', 'SPIRE', ]
    instr1 = JOptionPane.showInputDialog(None, 'Please select the instrument for map1.', \
    "Convolution", JOptionPane.QUESTION_MESSAGE, None, instr1list, instr1list[0])
    obs1 = getObservation(obs1id, useHsa=True, instrument = instr1)
    print 'Observation layout: '
    print obs1.refs
    obs1lvl = raw_input('Data level to use (e.g. level2): ')
    print 'Data level layout: '
    print obs1.refs[obs1lvl].product.refs
    obs1wb = raw_input('Waveband to use (e.g. PLW): ')
    if obs1.refs[obs1lvl].product.refs[obs1wb].product.class == SimpleImage().class:
        map1 = obs1.refs[obs1lvl].product.refs[obs1wb].product
    else:
        print 'Waveband level layout: '
        print obs1.refs[obs1lvl].product.refs[obs1wb].product.refs
        obs1map = raw_input('Map level to use (e.g. madmap): ')
        map1 = obs1.refs[obs1lvl].product.refs[obs1wb].product.refs[obs1map].product
    map1id = raw_input('Where do you want to save the convolved map (e.g. /home/user_name/hipe_convolution/)?')
    print 'Map1 successfully selected as a ' + map1select + '.'

elif map1select == 'local map':
    # Load the map from disk.
    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 a ' + map1select + '.'

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


# Prompt for map2.
print 'Now selecting map2; the map to which map1 is to be convolved to.'
possibleMap2select = ['local map', 'HSA observation', 'Scanamorphos map', ]
map2select = JOptionPane.showInputDialog(None, 'Please enter the map2 format', \
"Convolution", JOptionPane.QUESTION_MESSAGE, None, possibleMap2select, possibleMap2select[0])

if map2select == 'Scanamorphos map':
    # Load the Scanamorphos map from local disk.
    chooseMap2 = JFileChooser()
    filter2 = filechooser.FileNameExtensionFilter("fits files", ["fits"])
    chooseMap2.addChoosableFileFilter(filter2)
    cube2full = chooseMap2.showDialog(None, 'Please select map2')

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

    cube2 = simpleFitsReader(file = map2id)

    # Convert the Scanamorphos map to hipe readable map.
    map2 = SimpleImage()
    map2.description = 'Map from Scanamorphos - '+ cube2.meta['filter'].value
    map2.image = Double2d(cube2['PrimaryImage'].data[0,:,:])
    map2['image'].description = cube2.meta['PLANE0'].value
    map2.error = Double2d(cube2['PrimaryImage'].data[1,:,:])
    map2['error'].description = cube2.meta['PLANE1'].value
    map2.coverage = Double2d(cube2['PrimaryImage'].data[3,:,:])
    map2['coverage'].description = cube2.meta['PLANE3'].value
    map2.exposure = Double2d(cube2['PrimaryImage'].data[2,:,:])
    map2['exposure'].description = cube2.meta['PLANE2'].value
    map2.wcs.ctype1 = cube2.meta['ctype1'].value
    map2.wcs.ctype2 = cube2.meta['ctype2'].value
    map2.wcs.equinox = cube2.meta['equinox'].value
    map2.wcs.crpix1 = cube2.meta['crpix1'].value
    map2.wcs.crpix2 = cube2.meta['crpix2'].value
    map2.wcs.crval1 = cube2.meta['crval1'].value
    map2.wcs.crval2 = cube2.meta['crval2'].value
    map2.wcs.cdelt1 = cube2.meta['cdelt1'].value
    map2.wcs.cdelt2 = cube2.meta['cdelt2'].value
    map2.wcs.crota2 = cube2.meta['crota2'].value
    map2.meta = cube2.meta
    print 'Map2 successfully selected as a ' + map2select + '.'

elif map2select == 'HSA observation':
    # Acquire the observation from the Herschel Science Archive and select data level and waveband.
    obs2id = int(raw_input('Please enter observation id for map2 (e.g. 1342201436): '))
    instr2list = ['PACS', 'SPIRE', ]
    instr2 = JOptionPane.showInputDialog(None, 'Please select the instrument for map2.', \
    "Convolution", JOptionPane.QUESTION_MESSAGE, None, instr2list, instr2list[0])
    obs2 = getObservation(obs2id, useHsa=True, instrument = instr2)
    print 'Observation layout: '
    print obs2.refs
    obs2lvl = raw_input('Data level to use (e.g. level2): ')
    print 'Data level layout: '
    print obs2.refs[obs2lvl].product.refs
    obs2wb = raw_input('Waveband to use (e.g. PLW): ')
    if obs2.refs[obs2lvl].product.refs[obs2wb].product.class == SimpleImage().class:
        map2 = obs2.refs[obs2lvl].product.refs[obs2wb].product
    else:
        print 'Waveband level layout: '
        print obs2.refs[obs2lvl].product.refs[obs2wb].product.refs
        obs2map = raw_input('Map level to use (e.g. madmap): ')
        map2 = obs2.refs[obs2lvl].product.refs[obs2wb].product.refs[obs2map].product
    print 'Map2 successfully selected as a ' + map2select + '.'

elif map2select == 'local map':
    # Load the map from disk.
    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 a ' + map2select + '.'

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

# Determine which instrument/band combinations the observations were made in.
# Select the appropriate kernel and load it from the internet.
print 'Successfully selected both maps.'

try:
    instr1 = map1.meta['instrument'].value.strip()
except (ValueError, IndexError):
    instr1list = ['PACS', 'SPIRE', ]
    instr1 = JOptionPane.showInputDialog(None, 'Could not find the instrument for map1 in the metadata. \
    Please select manually.', "Convolution", JOptionPane.QUESTION_MESSAGE, None, instr1list, instr1list[0])

try:
    instr2 = map2.meta['instrument'].value.strip()
except (ValueError, IndexError):
    instr2list = ['PACS', 'SPIRE', ]
    instr2 = JOptionPane.showInputDialog(None, 'Could not find the instrument for map2 in the metadata. \
    Please select manually.', "Convolution", JOptionPane.QUESTION_MESSAGE, None, instr2list, instr2list[0])

try:
    if map1select == 'Scanamorphos map':
        band1 = map1.meta['filter'].value.strip()
    else:
        band1 = str(int(map1.meta['wavelength'].value))
except (ValueError, IndexError):
    band1 = str(int(raw_input('Could not find the waveband in the header file. \
    Please insert the wavelength for map1 manually (i.e. PACS: 70, 100, 160; SPIRE: 250, 350, 500).')))

try:
    if map2select == 'Scanamorphos map':
        band2 = map2.meta['filter'].value.strip()
    else:
        band2 = str(int(map2.meta['wavelength'].value)) 
except (ValueError, IndexError):
    band2 = str(int(raw_input('Could not find the waveband in the header file. \
    Please insert the wavelength for map2 manually (i.e. PACS: 70, 100, 160; SPIRE: 250, 350, 500).')))

possibleKernelSelect = ['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 == 'Premade kernel (online)':
    print 'Automatically selecting proper convolution kernel.'
    kernelName = '_aniano'
    kernelRoot = 'http://www.astro.princeton.edu/~ganiano/Kernels/Ker_2012_May/Kernels_fits_Files/Hi_Resolution/'
    kernelFile = 'Kernel_HiRes_%s_%s_to_%s_%s.fits.gz' % (instr1, band1, instr2, band2)
    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 dimensions 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 'Creating Gaussian kernel from the instrument beam sizes.'
    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 for PSF FWHM.
    fwhm = {'70' : 5.55, '100' : 6.75, '160' : 11.3, '250' : 18.2, '350' : 24.9,'500' : 36.3}
    fwhm1 = fwhm[band1]
    sigma1 = ABS(fwhm1 / (SQRT(8. * LOG(2)) * 3600. * kerWcs.getCdelt1()))
    fwhm2 = fwhm[band2]
    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.'))
    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.
print 'Performing the convolution using the rescaled, rotated and normalized kernel.'
kernelRot.image = kernelRot.image / SUM(kernelRot.image)
convMap1 = imageConvolution(image = map1, kernel = kernelRot)
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 + kernelName + '.fits'
fits.save(convid, convMap1Regrid)
print 'Saved the convolved map1 to ' + convid + '.'