__init__.py

# -*- coding: utf-8; -*-
# stacker, Python module for stacking of interferometric data.
# Copyright (C) 2014  Lukas Lindroos
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor,
# Boston, MA 02110-1301, USA.

"""
    Library to stack interferometric images.
"""

import math
import os
from ctypes import cdll
import re
import glob
import platform

__author__ = 'Lukas Lindroos'
__copyright__ = 'Copyright 2014'
__license__ = 'GPL'
__version__ = '1.0.3'
__maintainer__ = 'Lukas Lindroos'
__email__ = 'lindroos@chalmers.se'

PB_CONST = 0
PB_MS = 1
PB_FITS = 2

FILETYPENAME = {}
FILE_TYPE_NONE = 0
FILETYPENAME[FILE_TYPE_NONE] = 'none'
FILE_TYPE_MS = 1
FILETYPENAME[FILE_TYPE_MS] = 'ms'
FILE_TYPE_FITS = 2
FILETYPENAME[FILE_TYPE_FITS] = 'fits'

MS_DATACOLUMN_DATA = 1
MS_MODELCOLUMN_DATA = 2


clib_path = os.path.join(os.path.abspath(__path__[0]),
                         'stacker_clib')

if platform.system() == 'Linux':
    lib_suffix = '.so'
elif platform.system() == 'Darwin':
    lib_suffix = '.dylib'


def _read_casadef(casapath, param):
    """Read and return param from casadef.py inside specified casapath"""
    import re
    if platform.system() == 'Linux':
        casadef = casapath + '/lib/python2.7/casadef.py'
    elif platform.system() == 'Darwin':
        casadef = casapath + '/Resources/python/casadef.py'
    else:
        raise ValueError('Unsupported platform: ' + platform,system())
    try:
        casadeffile = open(casadef)
        for line in casadeffile:
            match = re.match(param.strip() + ' *= *"(.*)"', line)
            if match:
                return match.groups()[0]
    except IOError:
        raise ValueError('Could not open ' + casadef)


def _casa_svnversion(casapath):
    import re
    svnversion = None

    svnversion = _read_casadef(casapath, "subversion_revision")

    if svnversion is None:
        raise IOError('Can not find casa version.')

    return svnversion


def _libs_matching_svn_version(svnversion):
    """Return the list of stacker versions for the specified svn version"""
    return sorted(glob.glob(os.path.join(clib_path, 'libstacker-r{0}*{1}'.format(svnversion, lib_suffix))), reverse=True)


def _load_stacker_casa_svn_version(svnversion):
    """Sequentially attempt to load the options for the supplied svn version until successful"""
    libpath_options = _libs_matching_svn_version(svnversion)
    print('Loading stacking library for casapy svn revision {0}'.format(svnversion))
    for libpath in libpath_options:
        try:
            libstacker = cdll.LoadLibrary(libpath)
            print("{0} loaded".format(libpath))
            return libstacker
        except OSError:
            continue
    else:
        print("Could not load library for svn version {0}".format(svnversion))

def _load_stacker_lib():
    """Attempt to load a suitable stacker lib"""
    try:
        # If we are inside CASA - use libs linked against NRAO's CASA releases
        from taskinit import casa
        svnversion = _casa_svnversion(casa['dirs']['root'])
        lib_list = _libs_matching_svn_version(svnversion)
        if len(lib_list) > 0:
            libstacker = _load_stacker_casa_svn_version(svnversion)
        else:
            print('warning, no precompiled library compatible with your version of casa exists.')
            print('It is recommended to recompile stacker for your casa version.')
            stacker_clib_ls = os.listdir(clib_path)
            stackerlibs = [((re.match('libstacker-r([0-9]*).*.so', f).group(1)), f)
                           for f in stacker_clib_ls
                           if re.match('libstacker-r[0-9]*.so', f)]
            print(stackerlibs)
            vdiff = [(abs(int(svnversion)-int(v)), lib, v)
                     for (v, lib) in stackerlibs]
            vdiff.sort()
            print(vdiff)
            print('Trying to use svn revision {0}.'.format(vdiff[0][2]))
            libpath = os.path.join(clib_path, vdiff[0][1])
            try:
                libstacker = cdll.LoadLibrary(libpath)
            except OSError, e:
                print(e)
                print("Loading libstacker failed. You may need to build stacker for your version of CASA.")
                return
    except ImportError:
        # We are in a pure python session.
        # Not that there is anything wrong with that.
        libpath = os.path.join(clib_path, 'libstacker.so')
        libstacker = cdll.LoadLibrary(libpath)
    return libstacker


libstacker = _load_stacker_lib()


class CoordList(list):
    """
        Extended list to contain list of coordinates.

    """
    def __init__(self, imagenames=[], coord_type='physical', unit='rad'):
        """
        Requires an image list in case of pixel coordinates to work properly.
        """

        super(CoordList, self).__init__()

        if isinstance(imagenames, str):
            imagenames = [imagenames]

        self.coords = []
        self.imagenames = imagenames
        self.coord_type = coord_type
        self.unit = unit

    def __getitem__(self, i):
        return self.coords[i]

    def __setitem__(self, i, x):
        self.coords[i] = x

    def append(self, x):
        self.coords.append(x)

    def __len__(self):
        return len(self.coords)

    def __iter__(self):
        for x in self.coords:
            yield x

    def __getslice__(self, i, j):
        new_a = CoordList(self.imagenames, self.coord_type, self.unit)
        new_a.coords = self.coords.__getslice__(i, j)
        return new_a

    def __repr__(self):
        ret = []
        for x in self.coords:
            ret.append('(' + x.__str__() + ')')
        return '\n'.join(ret)

    def __str__(self):
        ret = []
        for x in self.coords:
            ret.append('(' + x.__str__() + ')')
        return '\n'.join(ret)
#         return '{0}, {1}'.format(self.x, self.y)


class Coord:
    """
        Describes a stacking position.

        Class used internally to represent coordinates. May describe a
        physical coordinate or a pixel coordinate.
    """

    def __init__(self, x, y, weight=1., image=0):
        """
            Create a coordinate. A pixel coordinate should always specify
            to which image it belongs. Physical coordinates should be in
            J2000 radians.
        """
        self.x = x
        self.y = y
        self.weight = weight
        self.image = image

    def __str__(self):
        return '{0}, {1}'.format(self.x, self.y)


def readCoords(coordfile, unit='deg'):
    """
        Reads a coordinate file from disk and produces a list.

        coordfile:
            Path to coordinate file. A file in csv format. x and y should
            be in J2000 . A weight may be added in third column
            to weight positions for stacking. If no weight are wanted
            put no third column in coordinate file.
        unit:
            Unit of input coordinates. Allows two values, 'deg' and 'rad'.
    """
    import csv

    coordreader = csv.reader(open(coordfile, 'rb'), delimiter=',')
    coords = CoordList()
    for row in coordreader:
        x = float(row[0])
        y = float(row[1])

        if unit == 'deg':
            x = math.pi/180.*x
            y = math.pi/180.*y

        if len(row) > 2:
            weight = float(row[2])
        else:
            weight = 1.

        if x > 2*math.pi:
            x -= 2*math.pi
        if y > math.pi:
            y -= 2*math.pi

        coords.append(Coord(x, y, weight))

    return coords


def writeCoords(coordpath, coords, unit='deg'):
    import csv

    with open(coordpath, 'wb') as coordfile:
        coordwriter = csv.writer(coordfile, delimiter=',')
        for coord in coords:
            x = coord.x
            y = coord.y
            weight = coord.weight

            if unit == 'deg':
                x = x*180./math.pi
                y = y*180./math.pi

            coordwriter.writerow([str(x), str(y), str(weight)])




def _checkfile(filename, datacolumn):
    import re
    # Currently this supports only ms files
    # As such there is no reason to check filetype.
    # If it cannot be opened as ms it will not be supported.
#     if re.match('^.*[mM][sS]/*$', filename) is not None:
    try:
        from taskinit import ms
        ms.open(filename)
        ms.done()
    except ImportError:
        # This probably means that it was run from a pure python session.
        # We will relegate any checks that it is a valid ms file to the 
        # stacker.
        if not os.access(filename, os.F_OK):
            raise RuntimeError('Could not find data file "{}".'.format(
                filename))

    filename = filename
    filetype = FILE_TYPE_MS
    fileoptions = 0
    if datacolumn == 'data':
        fileoptions = MS_DATACOLUMN_DATA
    elif datacolumn == 'model' or datacolumn == 'model_data':
        fileoptions = MS_MODELCOLUMN_DATA
#     elif re.match('^.*[fF][iI][tT][sS]$', filename) is not None:
#         raise NotImplementedError('FITS format is currently not supported.')
    return filetype, filename, fileoptions


def coordsTocl(name, flux, coords):
    from taskinit import cl, qa

    flux = qa.quantity(flux)
    cl.done()
    cl.rename(name)
    for coord in coords:
        clpars = {}
        clpars['flux'] = -flux['value']
        clpars['fluxunit'] = flux['unit']
        clpars['dir'] = ['J2000', str(coord.x)+'rad', str(coord.y)+'rad']
        clpars['shape'] = 'point'
        cl.addcomponent(**clpars)
    cl.done()


def randomCoords(imagenames, ncoords=10):
    import random
    from taskinit import ia, qa

    xmin, xmax = [], []
    ymin, ymax = [], []
    for image in imagenames:
        ia.open(image)
        print image, ia.boundingbox()
        trc = ia.boundingbox()['trcf'].split(', ')
        blc = ia.boundingbox()['blcf'].split(', ')
        xmin.append(qa.convert(qa.quantity(trc[0]), 'rad')['value'])
        xmax.append(qa.convert(qa.quantity(blc[0]), 'rad')['value'])
        ymin.append(qa.convert(qa.quantity(blc[1]), 'rad')['value'])
        ymax.append(qa.convert(qa.quantity(trc[1]), 'rad')['value'])
        ia.done()

    randomcoords = CoordList(imagenames)
    for i in range(ncoords):
        imageid = random.randint(0, len(imagenames)-1)
        x = random.uniform(xmin[imageid], xmax[imageid])
        y = random.uniform(ymin[imageid], ymax[imageid])
        c = Coord(x, y, 1.0)
        randomcoords.append(c)

    return randomcoords


def randomizeCoords(coords, beam):
    import random
    import math

    randomcoords = CoordList(coords.imagenames, coords.coord_type,
                             unit=coords.unit)

    for coord in coords:
        dr = random.uniform(beam, 5*beam)
        dphi = random.uniform(0, 2*math.pi)
        x = coord.x + dr*math.cos(dphi)
        y = coord.y + dr*math.sin(dphi)
        randomcoords.append(Coord(x, y, coord.weight, coord.image))

    return randomcoords


def _getPixelCoords1ImSimpleProj(coords, imagename):
    from taskinit import ia
    from interval import interval

    ia.open(imagename)
    cs = ia.coordsys()
    imshape = ia.shape()
    ia.done()

    pixcoords = []
    for coord in coords:
        p = cs.convert(coordin=[coord.x, coord.y, 0, 0], absin=[True]*4,
                       unitsin=[coords.unit, coords.unit, 'pix', 'pix'],
                       absout=[True]*4, unitsout=['pix']*4)
        x = p[0]
        y = p[1]

        if x in interval[0, imshape[0]-1] and y in interval[0., imshape[1]-1]:
            c = Coord(x, y)
            try:
                c.index = coord.index
            except AttributeError:
                pass
            pixcoords.append(c)

    return pixcoords


def _getPixelCoords1Im(coords, imagename):
    from interval import interval
    import math

    try:
        from taskinit import ia
        ia.open(imagename)
        cs = ia.coordsys()
        Nx = ia.shape()[0]
        Ny = ia.shape()[1]
        ia.done()
        x0 = cs.referencevalue()['numeric'][0]
        y0 = cs.referencevalue()['numeric'][1]
        x_pix_ref = cs.referencepixel()['numeric'][0]
        y_pix_ref = cs.referencepixel()['numeric'][1]
        x_pix_inc = cs.increment()['numeric'][0]
        y_pix_inc = cs.increment()['numeric'][1]

# If we fail to load ia, we will use pyrap instead.
# This probably means stacker was loaded from outside casapy.
    except ImportError:
        from pyrap.images import image
        im = image(imagename)
        cs = im.coordinates().get_coordinate('direction')
        dir_axis_index = im.coordinates().get_axes().index(cs.get_axes())
        imshape = im.shape()
        try:
            x_axis_index = cs.get_axes().index('Right Ascension')
        except ValueError:
            raise ValueError('Could not find direction coordinate: '\
                              'RightAscension')
        try:
            y_axis_index = cs.get_axes().index('Declination')
        except ValueError:
            raise ValueError('Could not find direction coordinate: '\
                              'Declination')



        Nx = im.shape()[dir_axis_index+x_axis_index]
        Ny = im.shape()[dir_axis_index+y_axis_index]
        x0 = cs.get_referencevalue()[x_axis_index]
        y0 = cs.get_referencevalue()[y_axis_index]
        x_pix_ref = cs.get_referencepixel()[x_axis_index]
        y_pix_ref = cs.get_referencepixel()[y_axis_index]
        x_pix_inc = cs.get_increment()[x_axis_index]
        y_pix_inc = cs.get_increment()[y_axis_index]

    pixcoords = []
    for coord in coords:
        dx = (coord.x - x0)*math.cos(coord.y)
        dy = math.asin(math.sin(coord.y)/math.cos(dx)) - y0
        x = dx/x_pix_inc+x_pix_ref
        y = dy/y_pix_inc+y_pix_ref

        if x in interval[0, Nx-1] and y in interval[0., Ny-1]:
#             pixcoords.append(Coord(x,y, coord.weight))
            c = Coord(x, y, coord.weight)

            try:
                c.index = coord.index
            except AttributeError:
                pass

            pixcoords.append(c)

    return pixcoords


def make_pbfile(vis, pbfile):
    from taskinit import im, ms, ia, qa, tb
    import numpy as np
    from scipy.constants import c

    ms.open(vis)
    fields = ms.range('field_id')['field_id']
    ms.done()
    im.open(vis)
    im.selectvis(field=fields[0])
    ms.open(vis)
    freq = np.mean(ms.range('chan_freq')['chan_freq'])
    phase_dir = ms.range('phase_dir')['phase_dir']['direction']
    ms.done()

    phase_dir = phase_dir[0][0], phase_dir[1][0]
    phase_dir = [qa.formxxx(str(phase_dir[0])+'rad', format='hms'),
                 qa.formxxx(str(phase_dir[1])+'rad', format='dms')]
    phase_dir = 'J2000 '+' '.join(phase_dir)

    tb.open(vis+'/ANTENNA/')
    dishdia = np.min(tb.getcol('DISH_DIAMETER'))
    tb.done()

    # pb of 512 pix cover pb down to 0.001
    # ensure largest pixel to pixel var to .01
    minpb = 0.001
    nx = 512
    cellconv = (nx*np.sqrt(np.log(2)/np.log(1/minpb)))**-1

    beam = c/freq/dishdia
    cell = {}
    cell['value'] = beam*cellconv
    cell['unit'] = 'rad'

#     nx = int(3*3e8/freq/dishdia*1.22*180/
#              math.pi*3600/qa.convert(advise['cell'],
#              'arcsec')['value'])
    # Chosen as to be 3 times fwhm of primary beam,
    # should include up to approximately .01 of peak flux

    im.defineimage(nx=nx, ny=nx, cellx=cell, celly=cell, phasecenter=phase_dir)
    im.setvp(dovp=True)
    im.makeimage(type='pb', image=pbfile)
    im.done()
    ia.open(pbfile)
    cs = ia.coordsys()
    cs.setreferencevalue(type='direction', value=[0., 0.])
    ia.setcoordsys(cs.torecord())
    ia.maskhandler('delete', 'mask0')
    ia.done()


def getPixelCoords(coords, imagenames):
    """
        Creates pixel coordinate list from a physical coordinate
        list and a list of images.
    """

    pixcoords = CoordList(imagenames, 'pixel', unit='pix')

    for (i, imagename) in enumerate(pixcoords.imagenames):
        for coord in _getPixelCoords1Im(coords, imagename):
            coord.image = i
            pixcoords.append(coord)

    return pixcoords