"""Data loading backend for CASA Measurement Sets."""

import argparse
import logging

import casacore.tables
import casacore.quanta
import numpy as np
from astropy import units
import astropy.time
import astropy.coordinates

import katsdpimager.loader_core


_logger = logging.getLogger(__name__)


def _get(table, name, data, casacore_units=None, astropy_units=None,
         measinfo_type=None, measinfo_ref=None):
    """Apply scaling to data loaded from a column. Scaling is done manually
    rather than with :mod:`casacore.quanta`, because the latter only operates
    (slowly) on scalars while a direct scaling is vectorised.

    Because not all writers specify the units for all quantities, units are
    allowed to be missing.

    Parameters
    ----------
    table : `casacore.tables.table`
        Table from which to load
    name : str
        Column name
    data : array-like
        Raw data loaded from the column. The leading dimension must index
        rows.
    casacore_units : str, optional
        Expected unit. A column with compatible units will be converted.
        Use ``None`` for unitless quantities.
    astropy_units : :class:`astropy.units.Unit`, optional
        Astropy unit equivalent to `casacore_units`, used for the return
        value. Use ``None`` for unitless quantities (the return value will
        then be a plain numpy array).
    measinfo_type : str, optional
        If specified, requires the ``MEASINFO`` field keyword to have this
        type.
    measinfo_ref : str, optional
        If specified, requires the ``MEASINFO`` field keyword to have this
        reference.
    """
    keywords = table.getcolkeywords(name)
    quantum_units = keywords.get('QuantumUnits')
    if quantum_units is not None:
        if casacore_units is None:
            raise ValueError('Found unexpected QuantumUnits for column {}'.format(name))
        quantum_units = np.array(quantum_units)
        if len(data.shape) == 1:
            # Scalar column. Convert quantum_units from a 1D array of 1 element
            # to a numpy scalar
            if quantum_units.shape != (1,):
                raise ValueError('Column {} is scalar, but QuantumUnits has shape {}'
                                 .format(name, quantum_units.shape))
            quantum_units = quantum_units[0]
        # data[0] is included in the iteration so that broadcasting rules are
        # triggered if quantum_units has a smaller shape (e.g. one unit for both
        # receptors on a feed).
        it = np.nditer([quantum_units, data[0]], flags=['multi_index'])
        output_quantity = casacore.quanta.quantity(1, casacore_units)
        while not it.finished:
            input_quantity = casacore.quanta.quantity(1, it[0][()])
            if not input_quantity.conforms(output_quantity):
                raise ValueError('Expected {} in column {} but found {} instead'.format(
                                 casacore_units, name, input_quantity.get_unit()))
            ratio = input_quantity.get_value(casacore_units)
            if ratio != 1:
                # Special-case ratio=1 because it is the common case, and the scaling
                # is potentially expensive.
                data[np.index_exp[:] + it.multi_index] *= ratio
            it.iternext()
    if astropy_units is not None:
        data = units.Quantity(data, astropy_units, copy=False)
    measinfo = keywords.get('MEASINFO')
    if measinfo is not None:
        if (measinfo_ref is not None and measinfo['Ref'] != measinfo_ref) or \
                (measinfo_type is not None and measinfo['type'] != measinfo_type):
            raise ValueError('Unsupported MEASINFO for {}: {}'.format(name, measinfo))
    return data


def _getcol(table, name, start=0, count=-1,
            casacore_units=None, astropy_units=None,
            measinfo_type=None, measinfo_ref=None):
    """Wrap the getcol method to fetch a batch of data from a column, applying
    scaling to the expected units. See :func:`_get` for details.

    Parameters
    ----------
    table : `casacore.tables.table`
        Table from which to load
    name : str
        Column name
    start : int
        First row to load
    count : int
        Number of rows to load
    casacore_units : str, optional
        Expected unit. A column with compatible units will be converted.
        Use ``None`` for unitless quantities.
    astropy_units : :class:`astropy.units.Unit`, optional
        Astropy unit equivalent to `casacore_units`, used for the return
        value. Use ``None`` for unitless quantities (the return value will
        then be a plain numpy array).
    measinfo_type : str, optional
        If specified, requires the ``MEASINFO`` field keyword to have this
        type.
    measinfo_ref : str, optional
        If specified, requires the ``MEASINFO`` field keyword to have this
        reference.
    """
    data = table.getcol(name, start, count)
    return _get(table, name, data, casacore_units, astropy_units, measinfo_type, measinfo_ref)


def _getcolslice(table, name, blc, trc, inc=[], start=0, count=-1,
                 casacore_units=None, astropy_units=None,
                 measinfo_type=None, measinfo_ref=None):
    """Wrap the getcol method to fetch a batch of data from a column with
    slicing, applying scaling to the expected units. See :func:`_get` for
    details.

    Parameters
    ----------
    table : `casacore.tables.table`
        Table from which to load
    name : str
        Column name
    blc : list
        Bottom left corner for array slice
    trc : list
        Top right corner for array slice (inclusive)
    ind : list
        Strides
    start : int
        First row to load
    count : int
        Number of rows to load
    casacore_units : str, optional
        Expected unit. A column with compatible units will be converted.
        Use ``None`` for unitless quantities.
    astropy_units : :class:`astropy.units.Unit`, optional
        Astropy unit equivalent to `casacore_units`, used for the return
        value. Use ``None`` for unitless quantities (the return value will
        then be a plain numpy array).
    measinfo_type : str, optional
        If specified, requires the ``MEASINFO`` field keyword to have this
        type.
    measinfo_ref : str, optional
        If specified, requires the ``MEASINFO`` field keyword to have this
        reference.
    """
    data = table.getcolslice(name, blc, trc, inc, start, count)
    return _get(table, name, data, casacore_units, astropy_units, measinfo_type, measinfo_ref)


def _getcolchannels(table, name, start_channel, stop_channel, start=0, count=-1,
                    casacore_units=None, astropy_units=None,
                    measinfo_type=None, measinfo_ref=None):
    """Convenience wrapper around :func:`_getcolslice` for columns where the
    shape is channels by polarizations, and a range of channels is desired.
    """
    return _getcolslice(table, name, [start_channel, -1], [stop_channel - 1, -1], [],
                        start, count, casacore_units, astropy_units,
                        measinfo_type, measinfo_ref)


def _getcell(table, name, row,
             casacore_units=None, astropy_units=None,
             measinfo_type=None, measinfo_ref=None):
    """Like :meth:`_getcol`, but for a single cell"""
    data = table.getcell(name, row)
    # Temporary add a leading dimension, required by _get
    data = data[np.newaxis, :]
    data = _getcol(table, name, row, 1, casacore_units, astropy_units,
                   measinfo_type, measinfo_ref)
    out = _get(table, name, data, casacore_units, astropy_units, measinfo_type, measinfo_ref)
    return out[0]


def _fix_cache_size(table, column):
    """Workaround for https://github.com/casacore/casacore/issues/581.
    We set the column cache size to be large enough for two tiles worth of
    rows.
    """
    cache_size = 0
    try:
        for item in table.getdminfo(column)['SPEC']['HYPERCUBES'].values():
            cube_shape = tuple(item['CubeShape'])
            tile_shape = tuple(item['TileShape'])
            cache_shape = cube_shape[:2] + tile_shape[2:]
            cache_size = max(cache_size, 2 * int(np.product(cache_shape)))
    except (TypeError, KeyError, IndexError, ValueError):
        pass
    if cache_size == 0:
        # No tile size info found - pick some default
        _logger.warning('Could not get tile info for column %s', column)
        cache_size = 1024 * 1024
    _logger.debug('Using cache size %d for column %s', cache_size, column)
    table.setmaxcachesize(column, cache_size)


class LoaderMS(katsdpimager.loader_core.LoaderBase):
    def __init__(self, filename, options):
        super().__init__(filename, options)
        parser = argparse.ArgumentParser(
            prog='Measurement set options',
            usage='Measurement set options: [-i data=COLUMN] [-i field=FIELD] ...')
        parser.add_argument('--data', type=str, metavar='COLUMN', default='DATA',
                            help='Column containing visibilities to image [%(default)s]')
        parser.add_argument('--data-desc', type=int, default=0,
                            help='Data description ID to image [%(default)s]')
        parser.add_argument('--field', type=int, default=0,
                            help='Field to image [%(default)s]')
        parser.add_argument('--pol-frame', choices=['sky', 'feed'], default='sky',
                            help='Reference frame for polarization [%(default)s]')
        parser.add_argument('--uvw', choices=['casa', 'strict'], default='casa',
                            help='UVW sign convention [%(default)s]')
        args = parser.parse_args(options)
        self._strict_uvw = (args.uvw == 'strict')
        self._main = casacore.tables.table(filename, ack=False)
        _fix_cache_size(self._main, 'FLAG')
        self._antenna = casacore.tables.table(self._main.getkeyword('ANTENNA'), ack=False)
        self._data_description = casacore.tables.table(self._main.getkeyword('DATA_DESCRIPTION'),
                                                       ack=False)
        self._field = casacore.tables.table(self._main.getkeyword('FIELD'), ack=False)
        self._spectral_window = casacore.tables.table(self._main.getkeyword('SPECTRAL_WINDOW'),
                                                      ack=False)
        self._polarization = casacore.tables.table(self._main.getkeyword('POLARIZATION'), ack=False)
        self._feed = casacore.tables.table(self._main.getkeyword('FEED'), ack=False)
        self._data_col = args.data
        self._field_id = args.field
        self._data_desc_id = args.data_desc
        if self._data_col not in self._main.colnames():
            raise ValueError('{} has no column named {}'.format(
                filename, self._data_col))
        if args.field < 0 or args.field >= self._field.nrows():
            raise ValueError('Field {} is out of range'.format(args.field))
        if args.data_desc < 0 or args.data_desc >= self._data_description.nrows():
            raise ValueError('Data description {} is out of range'.format(args.data_desc))
        self._polarization_id = self._data_description.getcell(
            'POLARIZATION_ID', args.data_desc)
        self._spectral_window_id = self._data_description.getcell(
            'SPECTRAL_WINDOW_ID', args.data_desc)
        # Detect whether we have a usable WEIGHT_SPECTRUM column. Measurement
        # sets in the wild sometimes have the column but with 0x0 shape.
        try:
            self._main.getcellslice('WEIGHT_SPECTRUM', 0, [0, 0], [0, 0])
        except RuntimeError:
            self._has_weight_spectrum = False
        else:
            self._has_weight_spectrum = True

        self._feed_angle_correction = args.pol_frame == 'feed'
        if self._feed_angle_correction:
            # Load per-antenna feed angles. For now, only a constant value per
            # antenna is supported, rather than per feed or per receptor within a
            # feed. This avoids the need for more per-visibility indexing.
            antenna_id = _getcol(self._feed, 'ANTENNA_ID')
            receptor_angle = _getcol(self._feed, 'RECEPTOR_ANGLE', 0, -1, 'rad', units.rad)
            antenna_angle = [None] * (np.max(antenna_id) + 1)
            for i in range(len(antenna_id)):
                for angle in receptor_angle[i]:
                    if (antenna_angle[antenna_id[i]] is not None
                            and not np.allclose(antenna_angle[antenna_id[i]],
                                                angle, atol=1e-8 * units.rad)):
                        raise ValueError('Multiple feed angles for one antenna is not supported')
                    antenna_angle[antenna_id[i]] = angle
            self._antenna_angle = units.Quantity(antenna_angle)
        else:
            self._antenna_angle = None

    @classmethod
    def match(cls, filename):
        return filename.lower().endswith('.ms')

    def antenna_diameters(self):
        return _getcol(self._antenna, 'DISH_DIAMETER', 0, -1, 'm', units.m)

    def antenna_positions(self):
        return _getcol(self._antenna, 'POSITION', 0, -1, 'm', units.m, 'position', 'ITRF')

    def phase_centre(self):
        value = _getcell(self._field, 'PHASE_DIR', self._field_id,
                         'rad', units.rad, 'direction', 'J2000')
        if tuple(value.shape) != (1, 2):
            raise ValueError('Unsupported shape for PHASE_DIR: {}'.format(value.shape))
        return value[0, :]

    def num_channels(self):
        return len(_getcell(self._spectral_window, 'CHAN_FREQ', self._spectral_window_id,
                            'Hz', units.Hz))

    def frequency(self, channel):
        return _getcell(self._spectral_window, 'CHAN_FREQ', self._spectral_window_id,
                        'Hz', units.Hz)[channel]

    def polarizations(self):
        return _getcell(self._polarization, 'CORR_TYPE', self._polarization_id)

    def has_feed_angles(self):
        return self._feed_angle_correction

    def data_iter(self, start_channel, stop_channel, max_chunk_vis=None):
        if max_chunk_vis is None:
            max_chunk_vis = self._main.nrows()
        num_channels = stop_channel - start_channel
        max_chunk_rows = self._main.nrows()
        if max_chunk_vis is not None:
            max_chunk_rows = max(1, max_chunk_vis // num_channels)
        # PHASE_DIR is used as an approximation to the antenna pointing
        # direction, for the purposes of parallactic angle correction. This
        # probably doesn't generalise well beyond dishes with single-pixel
        # feeds.
        pointing = self.phase_centre()
        pointing = astropy.coordinates.SkyCoord(ra=pointing[0], dec=pointing[1], frame='fk5')
        pos = self.antenna_positions()
        pos = astropy.coordinates.EarthLocation.from_geocentric(pos[:, 0], pos[:, 1], pos[:, 2])
        # Construct a point that is displaced from the pointing by a small
        # quantity northwards. It's necessary to use a small finite difference
        # rather than the pole itself, because the transformation to AzEl is
        # not rigid (does not preserve great circles).
        pole = pointing.directional_offset_by(0 * units.rad, 1e-5 * units.rad)
        for start in range(0, self._main.nrows(), max_chunk_rows):
            end = min(self._main.nrows(), start + max_chunk_rows)
            nrows = end - start
            flag_row = _getcol(self._main, 'FLAG_ROW', start, nrows)
            field_id = _getcol(self._main, 'FIELD_ID', start, nrows)
            data_desc_id = _getcol(self._main, 'DATA_DESC_ID', start, nrows)
            antenna1 = _getcol(self._main, 'ANTENNA1', start, nrows)
            antenna2 = _getcol(self._main, 'ANTENNA2', start, nrows)
            valid = np.logical_not(flag_row)
            valid = np.logical_and(valid, field_id == self._field_id)
            valid = np.logical_and(valid, data_desc_id == self._data_desc_id)
            valid = np.logical_and(valid, antenna1 != antenna2)
            antenna1 = antenna1[valid]
            antenna2 = antenna2[valid]
            data = _getcolchannels(self._main, self._data_col, start_channel, stop_channel,
                                   start, nrows, 'Jy', units.Jy)[valid, ...]
            feed_angle1 = None
            feed_angle2 = None
            if self._feed_angle_correction:
                time_full = _getcol(self._main, 'TIME_CENTROID', start, nrows, 's', None,
                                    'epoch', 'UTC')[valid, ...]
                # Each time will be repeated per baseline, but we do not need to repeat all the
                # calculations for each time. Extract just the unique times.
                time, inverse = np.unique(time_full, return_inverse=True)
                # Convert time from MJD seconds to MJD. We do this here rather
                # than by passing 'd' to _getcol, because not all measurement sets
                # specify the units and we want to assume seconds if not specified.
                time = astropy.time.Time(time / 86400.0, format='mjd', scale='utc')
                # Convert to CIRS, which is closer to AltAz in astropy's
                # conversion graph. This avoids duplicating precession-nutation
                # calculations for every antenna.
                cirs_frame = astropy.coordinates.CIRS(obstime=time)
                pole_cirs = pole.transform_to(cirs_frame)
                pointing_cirs = pointing.transform_to(cirs_frame)
                feed_angle = units.Quantity(
                    np.empty((len(time), len(pos))), unit=units.rad, copy=False)
                for i in range(len(pos)):
                    altaz_frame = astropy.coordinates.AltAz(location=pos[i], obstime=time)
                    pole_altaz = pole_cirs.transform_to(altaz_frame)
                    pointing_altaz = pointing_cirs.transform_to(altaz_frame)
                    pa = pointing_altaz.position_angle(pole_altaz)
                    feed_angle[:, i] = pa + self._antenna_angle[i]
                feed_angle1 = feed_angle[inverse, antenna1].astype(np.float32)
                feed_angle2 = feed_angle[inverse, antenna2].astype(np.float32)
            uvw = _getcol(self._main, 'UVW', start, nrows, 'm', units.m, 'uvw')[valid, ...]
            if not self._strict_uvw:
                uvw = -uvw
            if self._has_weight_spectrum:
                weight = _getcolchannels(self._main, 'WEIGHT_SPECTRUM', start_channel, stop_channel,
                                         start, nrows)[valid, ...]
            else:
                weight = _getcol(self._main, 'WEIGHT', start, nrows)[valid, ...]
                # Fake a channel axis
                weight = weight[:, np.newaxis, :]
                weight = np.broadcast_to(weight, data.shape)
            flag = _getcolchannels(self._main, 'FLAG', start_channel, stop_channel,
                                   start, nrows)[valid, ...]
            weight = weight * np.logical_not(flag)
            baseline = (antenna1 * self._antenna.nrows() + antenna2)
            ret = dict(uvw=uvw,
                       weights=np.swapaxes(weight, 0, 1),
                       baselines=baseline,
                       vis=np.swapaxes(data, 0, 1),
                       progress=end,
                       total=self._main.nrows())
            if self._feed_angle_correction:
                ret['feed_angle1'] = feed_angle1
                ret['feed_angle2'] = feed_angle2
            yield ret

    @property
    def raw_data(self):
        return self._main

    def close(self):
        self._main.close()
        self._antenna.close()
        self._data_description.close()
        self._field.close()
        self._spectral_window.close()
        self._polarization.close()
        self._feed.close()
