import sys
import uuid
import tempfile

import re

if sys.version_info[0] == 2:  # noqa
    from io import BytesIO as StringIO
else:
    from io import StringIO

import warnings
from base64 import b64encode

import numpy as np
from astropy.io import fits
from matplotlib.pyplot import cm
from matplotlib.colors import Colormap
from astropy import units as u

from traitlets import HasTraits, validate, observe
from .traits import Any, Unicode, Float, Color, Bool, to_hex
from .utils import sanitize_image

__all__ = ['LayerManager', 'TableLayer', 'ImageLayer']

VALID_FRAMES = ['sky', 'ecliptic', 'galactic', 'sun', 'mercury', 'venus',
                'earth', 'mars', 'jupiter', 'saturn', 'uranus', 'neptune',
                'pluto', 'moon', 'io', 'europa', 'ganymede', 'callisto']

VALID_LON_UNITS = {u.deg: 'degrees',
                   u.hour: 'hours',
                   u.hourangle: 'hours'}


# NOTE: for cartesian coordinates, we can also allow custom units
VALID_ALT_UNITS = {u.m: 'meters',
                   u.imperial.foot: 'feet',
                   u.imperial.inch: 'inches',
                   u.imperial.mile: 'miles',
                   u.km: 'kilometers',
                   u.au: 'astronomicalUnits',
                   u.lyr: 'lightYears',
                   u.pc: 'parsecs',
                   u.Mpc: 'megaParsecs'}

VALID_ALT_TYPES = ['depth', 'altitude', 'distance', 'seaLevel', 'terrain']

VALID_MARKER_TYPES = ['gaussian', 'point', 'circle', 'square', 'pushpin']

VALID_MARKER_SCALES = ['screen', 'world']

VALID_STRETCHES = ['linear', 'log', 'power', 'sqrt', 'histeq']

# The following are columns that we add dynamically and internally, so we need
# to make sure they have unique names that won't clash with existing columns
SIZE_COLUMN_NAME = str(uuid.uuid4())
CMAP_COLUMN_NAME = str(uuid.uuid4())


def guess_lon_lat_columns(colnames):
    """
    Given column names in a table, return the columns to use for lon/lat, or
    None/None if no high confidence possibilities.
    """

    # Do all the checks in lowercase
    colnames_lower = [colname.lower() for colname in colnames]

    for lon, lat in [('ra', 'dec'), ('lon', 'lat'), ('lng', 'lat')]:

        # Check first for exact matches
        if colnames_lower.count(lon) == 1 and colnames_lower.count(lat) == 1:
            return lon, lat

        # Next check for columns that start with specified names

        lon_match = [colname.startswith(lon) for colname in colnames_lower]
        lat_match = [colname.startswith(lat) for colname in colnames_lower]

        if sum(lon_match) == 1 and sum(lat_match) == 1:
            return colnames[lon_match.index(True)], colnames[lat_match.index(True)]

        # We don't check for cases where lon/lat are inside the name but not at
        # the start since that might be e.g. for proper motions (pm_ra) or
        # errors (dlat).

    return None, None


def pick_unit_if_available(unit, valid_units):
    # Check for equality rather than just identity
    for valid_unit in valid_units:
        if unit == valid_unit:
            return valid_unit
    return unit

def csv_table_win_newline(table):
    '''
    Helper function to get Astropy tables as ASCII CSV with Windows line endings
    '''
    s = StringIO()
    table.write(s, format='ascii.basic', delimiter=',', comment=False)
    s.seek(0)
    #Replace single \r or \n characters with \r\n
    return re.sub(r"(?<![\r\n])(\r|\n)(?![\r\n])", "\r\n", s.read())


class LayerManager(object):
    """
    A simple container for layers.
    """

    def __init__(self, parent=None):
        self._layers = []
        self._parent = parent

    def add_image_layer(self, image=None, **kwargs):
        """
        Add an image layer to the current view

        Parameters
        ----------
        image : str or :class:`~astropy.io.fits.ImageHDU` or tuple
            The image to show, which should be given either as a filename,
            an :class:`~astropy.io.fits.ImageHDU` object, or a tuple of the
            form ``(array, wcs)`` where ``array`` is a Numpy array and ``wcs``
            is an astropy :class:`~astropy.wcs.WCS` object
        kwargs
            Additional keyword arguments can be used to set properties on the
            image layer.

        Returns
        -------
        layer : :class:`~pywwt.layers.ImageLayer`
        """
        layer = ImageLayer(self._parent, image=image, **kwargs)
        self._add_layer(layer)
        return layer

    def add_table_layer(self, table=None, frame='Sky', **kwargs):
        """
        Add a data layer to the current view

        Parameters
        ----------
        table : :class:`~astropy.table.Table`
            The table containing the data to show.
        frame : str
            The reference frame to use for the data. This should be either
            ``'Sky'``, ``'Ecliptic'``, ``Galactic``, or the name of a planet
            or a natural satellite in the solar system.
        kwargs
            Additional keyword arguments can be used to set properties on the
            table layer.

        Returns
        -------
        layer : :class:`~pywwt.layers.TableLayer`
        """

        # Validate frame
        if frame.lower() not in VALID_FRAMES:
            raise ValueError('frame should be one of {0}'.format('/'.join(sorted(str(x) for x in VALID_FRAMES))))
        frame = frame.capitalize()

        if table is not None:
            layer = TableLayer(self._parent, table=table, frame=frame, **kwargs)
        else:
            # NOTE: in future we may allow different arguments such as e.g.
            # orbit=, hence why we haven't made this a positional argument.
            raise ValueError("The table argument is required")
        self._add_layer(layer)
        return layer

    def add_data_layer(self, *args, **kwargs):
        """
        Deprecated, use ``add_table_layer`` instead.
        """
        warnings.warn('add_data_layer has been deprecated, use add_table_layer '
                      'instead', DeprecationWarning)
        return self.add_table_layer(*args, **kwargs)

    def _add_layer(self, layer):
        if layer in self._layers:
            raise ValueError("layer already exists in layer manager")
        self._layers.append(layer)
        layer._manager = self

    def remove_layer(self, layer):
        if layer not in self._layers:
            raise ValueError("layer not in layer manager")
        layer.remove()
        # By this point, the call to remove() above may already have resulted
        # in the layer getting removed, so we check first if it's still present.
        if layer in self._layers:
            self._layers.remove(layer)

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

    def __iter__(self):
        for layer in self._layers:
            yield layer

    def __getitem__(self, item):
        return self._layers[item]

    def __str__(self):
        if len(self) == 0:
            return 'Layer manager with no layers'
        else:
            s = 'Layer manager with {0} layers:\n\n'.format(len(self))
            for ilayer, layer in enumerate(self._layers):
                s += '  [{0}]: {1}\n'.format(ilayer, layer)
            return s

    __repr__ = __str__


class TableLayer(HasTraits):
    """
    A layer where the data is stored in an :class:`~astropy.table.Table`
    """

    lon_att = Unicode(help='The column to use for the longitude').tag(wwt='lngColumn')
    lon_unit = Any(help='The units to use for longitude').tag(wwt='raUnits')
    lat_att = Unicode(help='The column to use for the latitude').tag(wwt='latColumn')

    alt_att = Unicode(help='The column to use for the altitude').tag(wwt='altColumn')
    alt_unit = Any(help='The units to use for the altitude').tag(wwt='altUnit')
    alt_type = Unicode(help='The type of altitude').tag(wwt='altType')

    size_scale = Float(10, help='The factor by which to scale the size of the points').tag(wwt='scaleFactor')

    # NOTE: we deliberately don't link size_att to sizeColumn because we need to
    # compute the sizes ourselves based on the min/max and then use the
    # resulting column.
    size_att = Unicode(help='The column to use for the size')
    size_vmin = Float(None, allow_none=True)
    size_vmax = Float(None, allow_none=True)

    # NOTE: we deliberately don't link cmap_att to colorMapColumn because we
    # need to compute the colors ourselves based on the min/max and then use the
    # resulting column.
    cmap_att = Unicode(help='The column to use for the colormap')
    cmap_vmin = Float(None, allow_none=True)
    cmap_vmax = Float(None, allow_none=True)
    cmap = Any(cm.viridis, help='The Matplotlib colormap')

    color = Color('white', help='The color of the markers').tag(wwt='color')
    opacity = Float(1, help='The opacity of the markers').tag(wwt='opacity')

    marker_type = Unicode('gaussian', help='The type of marker').tag(wwt='plotType')
    marker_scale = Unicode('screen', help='Whether the scale is defined in '
                           'world or pixel coordinates').tag(wwt='markerScale')

    far_side_visible = Bool(False, help='Whether markers on the far side are '
                            'visible').tag(wwt='showFarSide')

    # TODO: support:
    # xAxisColumn
    # yAxisColumn
    # zAxisColumn
    # xAxisReverse
    # yAxisReverse
    # zAxisReverse

    def __init__(self, parent=None, table=None, frame=None, **kwargs):

        # TODO: need to validate reference frame
        self.table = table
        self.frame = frame

        self.parent = parent
        self.id = str(uuid.uuid4())

        # Attribute to keep track of the manager, so that we can notify the
        # manager if a layer is removed.
        self._manager = None
        self._removed = False

        self._initialize_layer()

        # Force defaults
        self._on_trait_change({'name': 'alt_type', 'new': self.alt_type})
        self._on_trait_change({'name': 'size_scale', 'new': self.size_scale})
        self._on_trait_change({'name': 'color', 'new': self.color})
        self._on_trait_change({'name': 'opacity', 'new': self.opacity})
        self._on_trait_change({'name': 'marker_type', 'new': self.marker_type})
        self._on_trait_change({'name': 'marker_scale', 'new': self.marker_scale})
        self._on_trait_change({'name': 'far_side_visible', 'new': self.far_side_visible})
        self._on_trait_change({'name': 'size_att', 'new': self.size_att})
        self._on_trait_change({'name': 'cmap_att', 'new': self.cmap_att})

        self.observe(self._on_trait_change, type='change')

        if any(key not in self.trait_names() for key in kwargs):
            raise KeyError('a key doesn\'t match any layer trait name')

        super(TableLayer, self).__init__(**kwargs)

        lon_guess, lat_guess = guess_lon_lat_columns(self.table.colnames)

        if 'lon_att' not in kwargs:
            self.lon_att = lon_guess or self.table.colnames[0]

        if 'lat_att' not in kwargs:
            self.lat_att = lat_guess or self.table.colnames[1]

    @validate('lon_unit')
    def _check_lon_unit(self, proposal):
        # Pass the proposal to Unit - this allows us to validate the unit,
        # and allows strings to be passed.
        unit = u.Unit(proposal['value'])
        unit = pick_unit_if_available(unit, VALID_LON_UNITS)
        if unit in VALID_LON_UNITS:
            return unit
        else:
            raise ValueError('lon_unit should be one of {0}'.format('/'.join(sorted(str(x) for x in VALID_LON_UNITS))))

    @validate('alt_unit')
    def _check_alt_unit(self, proposal):
        # Pass the proposal to Unit - this allows us to validate the unit,
        # and allows strings to be passed.
        with u.imperial.enable():
            unit = u.Unit(proposal['value'])
        unit = pick_unit_if_available(unit, VALID_ALT_UNITS)
        if unit in VALID_ALT_UNITS:
            return unit
        else:
            raise ValueError('alt_unit should be one of {0}'.format('/'.join(sorted(str(x) for x in VALID_ALT_UNITS))))

    @validate('alt_type')
    def _check_alt_type(self, proposal):
        if proposal['value'] in VALID_ALT_TYPES:
            return proposal['value']
        else:
            raise ValueError('alt_type should be one of {0}'.format('/'.join(str(x) for x in VALID_ALT_TYPES)))

    @validate('marker_type')
    def _check_marker_type(self, proposal):
        if proposal['value'] in VALID_MARKER_TYPES:
            return proposal['value']
        else:
            raise ValueError('marker_type should be one of {0}'.format('/'.join(str(x) for x in VALID_MARKER_TYPES)))

    @validate('marker_scale')
    def _check_marker_scale(self, proposal):
        if proposal['value'] in VALID_MARKER_SCALES:
            return proposal['value']
        else:
            raise ValueError('marker_scale should be one of {0}'.format('/'.join(str(x) for x in VALID_MARKER_SCALES)))

    @validate('cmap')
    def _check_cmap(self, proposal):
        if isinstance(proposal['value'], str):
            return cm.get_cmap(proposal['value'])
        elif not isinstance(proposal['value'], Colormap):
            raise TypeError('cmap should be set to a Matplotlib colormap')
        else:
            return proposal['value']

    @observe('alt_att')
    def _on_alt_att_change(self, *value):
        # Check if we can set the unit of the altitude automatically
        if len(self.alt_att) == 0:
            return
        column = self.table[self.alt_att]
        unit = pick_unit_if_available(column.unit, VALID_ALT_UNITS)
        if unit in VALID_ALT_UNITS:
            self.alt_unit = unit
        elif unit is not None:
            warnings.warn('Column {0} has units of {1} but this is not a valid '
                          'unit of altitude - set the unit directly with '
                          'alt_unit'.format(self.alt_att, unit), UserWarning)

    @observe('lon_att')
    def _on_lon_att_change(self, *value):
        # Check if we can set the unit of the altitude automatically
        if len(self.lon_att) == 0:
            return
        column = self.table[self.lon_att]
        unit = pick_unit_if_available(column.unit, VALID_LON_UNITS)
        if unit in VALID_LON_UNITS:
            self.lon_unit = unit
        elif unit is not None:
            warnings.warn('Column {0} has units of {1} but this is not a valid '
                          'unit of longitude - set the unit directly with '
                          'lon_unit'.format(self.lon_att, unit), UserWarning)

    @observe('size_att')
    def _on_size_att_change(self, *value):

        # Set the min/max levels automatically based on the min/max values

        if len(self.size_att) == 0:
            self.parent._send_msg(event='table_layer_set', id=self.id,
                                  setting='sizeColumn', value=-1)
            return

        self.size_vmin = None
        self.size_vmax = None

        column = self.table[self.size_att]

        self.size_vmin = np.nanmin(column)
        self.size_vmax = np.nanmax(column)

    @observe('size_vmin', 'size_vmax')
    def _on_size_vmin_vmax_change(self, *value):

        # Update the size column in the table

        if len(self.size_att) == 0 or self.size_vmin is None or self.size_vmax is None:
            self.parent._send_msg(event='table_layer_set', id=self.id,
                                  setting='sizeColumn', value=-1)
            return

        column = self.table[self.size_att]

        size = (column - self.size_vmin) / (self.size_vmax - self.size_vmin) * 10

        self.table[SIZE_COLUMN_NAME] = size

        self.parent._send_msg(event='table_layer_update', id=self.id,
                              table=self._table_b64)

        self.parent._send_msg(event='table_layer_set', id=self.id,
                              setting='pointScaleType', value=0)

        self.parent._send_msg(event='table_layer_set', id=self.id,
                              setting='sizeColumn', value=SIZE_COLUMN_NAME)

    @observe('cmap_att')
    def _on_cmap_att_change(self, *value):

        # Set the min/max levels automatically based on the min/max values

        if len(self.cmap_att) == 0:

            self.parent._send_msg(event='table_layer_set', id=self.id,
                                  setting='colorMapColumn', value=-1)

            self.parent._send_msg(event='table_layer_set', id=self.id,
                                  setting='_colorMap', value=0)

            return

        self.cmap_vmin = None
        self.cmap_vmax = None

        column = self.table[self.cmap_att]

        self.cmap_vmin = np.nanmin(column)
        self.cmap_vmax = np.nanmax(column)

    @observe('cmap_vmin', 'cmap_vmax', 'cmap')
    def _on_cmap_vmin_vmax_change(self, *value):

        # Update the cmap column in the table

        if len(self.cmap_att) == 0 or self.cmap_vmin is None or self.cmap_vmax is None:

            self.parent._send_msg(event='table_layer_set', id=self.id,
                                  setting='colorMapColumn', value=-1)

            self.parent._send_msg(event='table_layer_set', id=self.id,
                                  setting='_colorMap', value=0)

            return

        column = self.table[self.cmap_att]

        values = (column - self.cmap_vmin) / (self.cmap_vmax - self.cmap_vmin)

        # PERF: vectorize the calculation of the hex strings
        rgb = self.cmap(values)[:, :-1]
        hex_values = [to_hex(x) for x in rgb]

        self.table[CMAP_COLUMN_NAME] = hex_values

        self.parent._send_msg(event='table_layer_update', id=self.id,
                              table=self._table_b64)

        self.parent._send_msg(event='table_layer_set', id=self.id,
                              setting='_colorMap', value=3)

        self.parent._send_msg(event='table_layer_set', id=self.id,
                              setting='colorMapColumn', value=CMAP_COLUMN_NAME)

    @property
    def _table_b64(self):

        # TODO: We need to make sure that the table has ra/dec columns since
        # WWT absolutely needs that upon creation.

        csv = csv_table_win_newline(self.table)

        return b64encode(csv.encode('ascii', errors='replace')).decode('ascii')

    def _initialize_layer(self):
        self.parent._send_msg(event='table_layer_create',
                              id=self.id, table=self._table_b64, frame=self.frame)

    def update_data(self, table=None):
        """
        Update the underlying data.
        """
        self.table = table.copy(copy_data=False)
        self.parent._send_msg(event='table_layer_update', id=self.id, table=self._table_b64)

        if len(self.alt_att) > 0:
            if self.alt_att in self.table.colnames:
                self._on_alt_att_change()
            else:
                self.alt_att = ''

        lon_guess, lat_guess = guess_lon_lat_columns(self.table.colnames)

        if self.lon_att in self.table.colnames:
            self._on_lon_att_change()
        else:
            self.lon_att = lon_guess or self.table.colnames[0]

        if self.lat_att not in self.table.colnames:
            self.lat_att = lat_guess or self.table.colnames[1]

    def remove(self):
        """
        Remove the layer.
        """
        if self._removed:
            return
        self.parent._send_msg(event='table_layer_remove', id=self.id)
        self._removed = True
        if self._manager is not None:
            self._manager.remove_layer(self)

    def _on_trait_change(self, changed):
        # This method gets called anytime a trait gets changed. Since this class
        # gets inherited by the Jupyter widgets class which adds some traits of
        # its own, we only want to react to changes in traits that have the wwt
        # metadata attribute (which indicates the name of the corresponding WWT
        # setting).
        wwt_name = self.trait_metadata(changed['name'], 'wwt')
        if wwt_name is not None:
            value = changed['new']
            if changed['name'] == 'alt_unit':
                value = VALID_ALT_UNITS[self._check_alt_unit({'value': value})]
            elif changed['name'] == 'lon_unit':
                value = VALID_LON_UNITS[self._check_lon_unit({'value': value})]
            self.parent._send_msg(event='table_layer_set',
                                  id=self.id,
                                  setting=wwt_name,
                                  value=value)

    def __str__(self):
        return 'TableLayer with {0} markers'.format(len(self.table))

    def __repr__(self):
        return '<{0}>'.format(str(self))


class ImageLayer(HasTraits):
    """
    An image layer.
    """

    vmin = Float(None, allow_none=True)
    vmax = Float(None, allow_none=True)
    stretch = Unicode('linear')
    opacity = Float(1, help='The opacity of the markers').tag(wwt='opacity')

    def __init__(self, parent=None, image=None, **kwargs):

        self._image = image

        self.parent = parent
        self.id = str(uuid.uuid4())

        # Attribute to keep track of the manager, so that we can notify the
        # manager if a layer is removed.
        self._manager = None
        self._removed = False

        # Transform the image so that it is always acceptable to WWT (Equatorial,
        # TAN projection, double values) and write out to a temporary file
        self._sanitized_image = tempfile.mktemp()
        sanitize_image(image, self._sanitized_image)

        # The first thing we need to do is make sure the image is being served.
        # For now we assume that image is a filename, but we could do more
        # detailed checks and reproject on-the-fly for example.

        self._image_url = self.parent._serve_file(self._sanitized_image, extension='.fits')

        # Because of the way the image loading works in WWT, we may end up with
        # messages being applied out of order (see notes in image_layer_stretch
        # in wwt_json_api.js)
        self._stretch_version = 0

        self._initialize_layer()

        # Force defaults
        self._on_trait_change({'name': 'opacity', 'new': self.opacity})

        self.observe(self._on_trait_change, type='change')

        if any(key not in self.trait_names() for key in kwargs):
            raise KeyError('a key doesn\'t match any layer trait name')

        super(ImageLayer, self).__init__(**kwargs)

        # Determine initial stretch parameters
        data = fits.getdata(self._sanitized_image)

        self.vmin = np.nanpercentile(data, 0.5)
        self.vmax = np.nanpercentile(data, 99.5)

    @validate('stretch')
    def _check_stretch(self, proposal):
        if proposal['value'] in VALID_STRETCHES:
            return proposal['value']
        else:
            raise ValueError('stretch should be one of {0}'.format('/'.join(str(x) for x in VALID_STRETCHES)))

    def _initialize_layer(self):
        self.parent._send_msg(event='image_layer_create',
                              id=self.id, url=self._image_url)

    def remove(self):
        """
        Remove the layer.
        """
        if self._removed:
            return
        self.parent._send_msg(event='image_layer_remove', id=self.id)
        self._removed = True
        if self._manager is not None:
            self._manager.remove_layer(self)

    def _on_trait_change(self, changed):

        if changed['name'] in ('stretch', 'vmin', 'vmax'):
            if self.vmin is not None and self.vmax is not None:
                stretch_id = VALID_STRETCHES.index(self.stretch)
                self._stretch_version += 1
                self.parent._send_msg(event='image_layer_stretch', id=self.id,
                                      stretch=stretch_id,
                                      vmin=self.vmin, vmax=self.vmax,
                                      version=self._stretch_version)

        # This method gets called anytime a trait gets changed. Since this class
        # gets inherited by the Jupyter widgets class which adds some traits of
        # its own, we only want to react to changes in traits that have the wwt
        # metadata attribute (which indicates the name of the corresponding WWT
        # setting).
        wwt_name = self.trait_metadata(changed['name'], 'wwt')
        if wwt_name is not None:
            self.parent._send_msg(event='image_layer_set',
                                  id=self.id,
                                  setting=wwt_name,
                                  value=changed['new'])

    def __str__(self):
        return 'ImageLayer'

    def __repr__(self):
        return '<{0}>'.format(str(self))
