#
# AstroImage.py -- Abstraction of an astronomical data image.
#
# This is open-source software licensed under a BSD license.
# Please see the file LICENSE.txt for details.
#
import sys
import traceback
import warnings
from collections import OrderedDict

import numpy as np

from ginga.util import wcs, wcsmod
from ginga.BaseImage import BaseImage, ImageError, Header
from ginga.misc import Bunch


class AstroHeader(Header):
    pass


class AstroImage(BaseImage):
    """
    Abstraction of an astronomical data (image).

    NOTE: this module is NOT thread-safe!
    """
    # class variables for WCS and IO can be set
    wcsClass = None
    ioClass = None

    @classmethod
    def set_wcsClass(cls, klass):
        cls.wcsClass = klass

    @classmethod
    def set_ioClass(cls, klass):
        cls.ioClass = klass

    def __init__(self, data_np=None, metadata=None, logger=None,
                 name=None, wcsclass=None, ioclass=None,
                 inherit_primary_header=False, save_primary_header=True):

        BaseImage.__init__(self, data_np=data_np, metadata=metadata,
                           logger=logger, name=name)

        # wcsclass specifies a pluggable WCS module
        if wcsclass is None:
            if self.wcsClass is None:
                wcsclass = wcsmod.WCS
            else:
                wcsclass = self.wcsClass
        self.wcs = wcsclass(self.logger)

        # ioclass specifies a pluggable IO module
        if ioclass is None:
            if self.ioClass is None:
                from ginga.util import io_fits
                ioclass = io_fits.fitsLoaderClass
            else:
                ioclass = self.ioClass

        self.io = ioclass(self.logger)

        self.inherit_primary_header = inherit_primary_header
        self.save_primary_header = inherit_primary_header or save_primary_header
        if self.save_primary_header:
            # User wants to save/inherit from primary header--this will hold it
            self._primary_hdr = AstroHeader()
        else:
            self._primary_hdr = None

        if metadata is not None:
            header = self.get_header()
            self.wcs.load_header(header)

        # For navigating multidimensional data
        self.naxispath = []
        self.revnaxis = []
        self._md_data = None

    def setup_data(self, data, naxispath=None):
        # initialize data attribute to something reasonable
        if data is None:
            data = np.zeros((0, 0))
        elif not isinstance(data, np.ndarray):
            data = np.zeros((0, 0))
        elif 0 in data.shape:
            data = np.zeros((0, 0))

        # this is a handle to the full data array
        self._md_data = data

        # this will get reset in set_naxispath() if array is
        # multidimensional
        self._data = self._md_data

        if naxispath is None:
            naxispath = []

        # Set naxispath to drill down to first 2D data slice
        if len(naxispath) == 0:
            naxispath = ([0] * (len(data.shape) - 2))

        self.set_naxispath(naxispath)

    def load_hdu(self, hdu, fobj=None, naxispath=None,
                 inherit_primary_header=None):

        self.clear_metadata()

        # collect HDU header
        ahdr = self.get_header()
        self._copy_hdu_header(hdu, ahdr)

        # Set PRIMARY header
        if inherit_primary_header is None:
            inherit_primary_header = self.inherit_primary_header
        else:  # This ensures get_header() is consistent
            self.inherit_primary_header = inherit_primary_header

        save_primary_header = (self.save_primary_header or
                               inherit_primary_header)
        if save_primary_header and (fobj is not None):
            if self._primary_hdr is None:
                self._primary_hdr = AstroHeader()

            self._copy_hdu_header(fobj[0], self._primary_hdr)

        self.setup_data(hdu.data, naxispath=naxispath)

        # Try to make a wcs object on the header
        if hasattr(self, 'wcs') and self.wcs is not None:
            self.wcs.load_header(hdu.header, fobj=fobj)

    def load_nddata(self, ndd, naxispath=None):
        """Load from an astropy.nddata.NDData object.
        """
        self.clear_metadata()

        # Make a header based on any NDData metadata
        ahdr = self.get_header()
        ahdr.update(ndd.meta)

        self.setup_data(ndd.data, naxispath=naxispath)

        if ndd.wcs is None:
            # no wcs in ndd obj--let's try to make one from the header
            self.wcs = wcsmod.WCS(logger=self.logger)
            self.wcs.load_header(ahdr)
        else:
            # already have a valid wcs in the ndd object
            # we assume it needs an astropy compatible wcs
            wcsinfo = wcsmod.get_wcs_class('astropy')
            self.wcs = wcsinfo.wrapper_class(logger=self.logger)
            self.wcs.load_nddata(ndd)

    def load_file(self, filespec, **kwargs):

        if self.io is None:
            raise ImageError("No IO loader defined")

        self.io.load_file(filespec, dstobj=self, **kwargs)

    def load_data(self, data_np, naxispath=None, metadata=None):

        self.clear_metadata()

        self.setup_data(data_np, naxispath=naxispath)

        if metadata is not None:
            self.update_metadata(metadata)

    def load_buffer(self, buf, dims, dtype, byteswap=False,
                    naxispath=None, metadata=None):
        data = np.fromstring(buf, dtype=dtype)
        if byteswap:
            data.byteswap(True)
        data = data.reshape(dims)

        self.load_data(data, naxispath=naxispath, metadata=metadata)

    def get_mddata(self):
        return self._md_data

    def set_naxispath(self, naxispath):
        """Choose a slice out of multidimensional data.
        """
        revnaxis = list(naxispath)
        revnaxis.reverse()

        # construct slice view and extract it
        ndim = min(self.ndim, 2)
        view = tuple(revnaxis + [slice(None)] * ndim)
        data = self.get_mddata()[view]

        if len(data.shape) not in (1, 2):
            raise ImageError(
                "naxispath does not lead to a 1D or 2D slice: {}".format(naxispath))

        self.naxispath = naxispath
        self.revnaxis = revnaxis

        self.set_data(data)

    def set_wcs(self, wcs):
        self.wcs = wcs

    def set_io(self, io):
        self.io = io

    def get_data_size(self):
        return self.get_size()

    def get_header(self, create=True, include_primary_header=None):
        try:
            # By convention, the fits header is stored in a dictionary
            # under the metadata keyword 'header'
            hdr = self.metadata['header']

            if include_primary_header is None:
                include_primary_header = self.inherit_primary_header

            if include_primary_header and self._primary_hdr is not None:
                # Inherit PRIMARY header for display but keep metadata intact
                displayhdr = AstroHeader()
                for key in hdr.keyorder:
                    card = hdr.get_card(key)
                    bnch = displayhdr.__setitem__(card.key, card.value)
                    bnch.comment = card.comment
                for key in self._primary_hdr.keyorder:
                    if key not in hdr:
                        card = self._primary_hdr.get_card(key)
                        bnch = displayhdr.__setitem__(card.key, card.value)
                        bnch.comment = card.comment
            else:
                # Normal, separate header
                displayhdr = hdr

        except KeyError as e:
            if not create:
                raise e
            hdr = AstroHeader()
            self.metadata['header'] = hdr
            displayhdr = hdr

        return displayhdr

    def get_keyword(self, kwd, *args):
        """Get an item from the fits header, if any."""
        try:
            kwds = self.get_header()
            return kwds[kwd]
        except KeyError:
            # return a default if there is one
            if len(args) > 0:
                return args[0]
            raise KeyError(kwd)

    def get_keywords_list(self, *args):
        return list(map(self.get_keyword, args))

    def set_keyword(self, kwd, value, create=True):
        kwds = self.get_header(create=create)
        kwd = kwd.upper()
        if not create:
            prev = kwds[kwd]  # noqa, this raises KeyError
        kwds[kwd] = value

    def update_keywords(self, keyDict):
        hdr = self.get_header()
        # Upcase all keywords
        for kwd, val in keyDict.items():
            hdr[kwd.upper()] = val

        # Try to make a wcs object on the header
        if hasattr(self, 'wcs') and self.wcs is not None:
            self.wcs.load_header(hdr)

    def set_keywords(self, **kwds):
        """Set an item in the fits header, if any."""
        return self.update_keywords(kwds)

    def update_data(self, data_np, metadata=None, astype=None):
        """DO NOT USE: this method will be deprecated!
        """
        self.set_data(data_np.copy(), metadata=metadata,
                      astype=astype)

    def update_metadata(self, key_dict):
        for key, val in key_dict.items():
            self.metadata[key] = val

        # refresh the WCS
        if hasattr(self, 'wcs') and self.wcs is not None:
            header = self.get_header()
            self.wcs.load_header(header)

    def has_primary_header(self):
        return self._primary_hdr is not None

    def clear_all(self):
        # clear metadata and data
        super(AstroImage, self).clear_all()

        # unreference full data array
        self._md_data = self._data

    def transfer(self, other, astype=None):
        data = self._get_data()
        other.update_data(data, astype=astype)
        other.update_metadata(self.metadata)

    def copy(self, astype=None):
        data = self._get_data()
        other = AstroImage(data, logger=self.logger)
        self.transfer(other, astype=astype)
        return other

    def as_nddata(self, nddata_class=None):
        "Return a version of ourself as an astropy.nddata.NDData object"
        if nddata_class is None:
            from astropy.nddata import NDData
            nddata_class = NDData

        # transfer header, preserving ordering
        ahdr = self.get_header()
        header = OrderedDict(ahdr.items())
        data = self.get_mddata()

        wcs = None
        if hasattr(self, 'wcs') and self.wcs is not None:
            # for now, assume self.wcs wraps an astropy wcs object
            wcs = self.wcs.wcs

        ndd = nddata_class(data, wcs=wcs, meta=header)
        return ndd

    def as_hdu(self):
        "Return a version of ourself as an astropy.io.fits.PrimaryHDU object"
        from astropy.io import fits

        # transfer header, preserving ordering
        ahdr = self.get_header()
        header = fits.Header(ahdr.items())
        data = self.get_mddata()

        hdu = fits.PrimaryHDU(data=data, header=header)
        return hdu

    def _copy_hdu_header(self, hdu, ahdr):
        """Copy a FITS header from an astropy.io.fits.PrimaryHDU object
        into a ginga.AstroImage.AstroHeader object.
        """
        header = hdu.header
        for card in header.cards:
            if len(card.keyword) == 0:
                continue
            ahdr.set_card(card.keyword, card.value, comment=card.comment)

    def astype(self, type_name):
        """Convert AstroImage object to some other kind of object.
        """
        if type_name == 'nddata':
            return self.as_nddata()

        if type_name == 'hdu':
            return self.as_hdu()

        raise ValueError("Unrecognized conversion type '%s'" % (type_name))

    def save_as_file(self, filepath, **kwdargs):
        data = self._get_data()
        header = self.get_header()
        self.io.save_as_file(filepath, data, header, **kwdargs)

    def pixtocoords(self, x, y, system=None, coords='data'):
        args = [x, y] + self.revnaxis
        return self.wcs.pixtocoords(args, system=system, coords=coords)

    def spectral_coord(self, coords='data'):
        args = [0, 0] + self.revnaxis
        return self.wcs.spectral_coord(args, coords=coords)

    def pixtoradec(self, x, y, format='deg', coords='data'):
        args = [x, y] + self.revnaxis
        ra_deg, dec_deg = self.wcs.pixtoradec(args, coords=coords)

        if format == 'deg':
            return ra_deg, dec_deg
        return wcs.deg2fmt(ra_deg, dec_deg, format)

    def radectopix(self, ra_deg, dec_deg, format='deg', coords='data'):
        if format != 'deg':
            # convert coordinates to degrees
            ra_deg = wcs.lon_to_deg(ra_deg)
            dec_deg = wcs.lat_to_deg(dec_deg)
        return self.wcs.radectopix(ra_deg, dec_deg, coords=coords,
                                   naxispath=self.revnaxis)

    # -----> TODO:
    #   This section has been merged into ginga.util.wcs or
    #   ginga.util.mosaic .  Deprecate it here.
    #
    def get_starsep_XY(self, x1, y1, x2, y2):
        warnings.warn("This function has been deprecated--"
                      "use the version in ginga.util.wcs",
                      PendingDeprecationWarning)
        return wcs.get_starsep_XY(self, x1, y1, x2, y2)

    def calc_radius_xy(self, x, y, radius_deg):
        warnings.warn("This function has been deprecated--"
                      "use the version in ginga.util.wcs",
                      PendingDeprecationWarning)
        return wcs.calc_radius_xy(self, x, y, radius_deg)

    def calc_radius_deg2pix(self, ra_deg, dec_deg, delta_deg,
                            equinox=None):
        warnings.warn("This function has been deprecated--"
                      "use the version in ginga.util.wcs",
                      PendingDeprecationWarning)
        return wcs.calc_radius_deg2pix(self, ra_deg, dec_deg, delta_deg,
                                       equinox=equinox)

    def add_offset_xy(self, x, y, delta_deg_x, delta_deg_y):
        warnings.warn("This function has been deprecated--"
                      "use the version in ginga.util.wcs",
                      PendingDeprecationWarning)
        return wcs.add_offset_xy(self, x, y, delta_deg_x, delta_deg_y)

    def calc_radius_center(self, delta_deg):
        warnings.warn("This function has been deprecated--"
                      "use the version in ginga.util.wcs",
                      PendingDeprecationWarning)
        return wcs.calc_radius_center(self, delta_deg)

    def calc_compass(self, x, y, len_deg_e, len_deg_n):
        warnings.warn("This function has been deprecated--"
                      "use the version in ginga.util.wcs",
                      PendingDeprecationWarning)
        return wcs.calc_compass(self, x, y, len_deg_e, len_deg_n)

    def calc_compass_radius(self, x, y, radius_px):
        warnings.warn("This function has been deprecated--"
                      "use the version in ginga.util.wcs",
                      PendingDeprecationWarning)
        return wcs.calc_compass_radius(self, x, y, radius_px)

    def calc_compass_center(self):
        warnings.warn("This function has been deprecated--"
                      "use the version in ginga.util.wcs",
                      PendingDeprecationWarning)
        return wcs.calc_compass_center(self)

    def get_wcs_rotation_deg(self):
        warnings.warn("This function has been deprecated--"
                      "use get_rotation_and_scale in ginga.util.wcs",
                      PendingDeprecationWarning)
        header = self.get_header()
        (rot, cdelt1, cdelt2) = wcs.get_rotation_and_scale(header)
        return rot

    def mosaic_inline(self, imagelist, bg_ref=None, trim_px=None,
                      merge=False, allow_expand=True, expand_pad_deg=0.01,
                      max_expand_pct=None,
                      update_minmax=True, suppress_callback=False):
        warnings.warn("This function has been deprecated--"
                      "use the version in ginga.util.mosaic",
                      PendingDeprecationWarning)
        from ginga.util import mosaic
        return mosaic.mosaic_inline(self, imagelist, bg_ref=bg_ref,
                                    trim_px=trim_px, merge=merge,
                                    allow_expand=allow_expand,
                                    expand_pad_deg=expand_pad_deg,
                                    max_expand_pct=max_expand_pct,
                                    update_minmax=update_minmax,
                                    suppress_callback=suppress_callback)
    #
    # <----- TODO: deprecate

    def info_xy(self, data_x, data_y, settings):
        # Get the value under the data coordinates
        try:
            # We report the value across the pixel, even though the coords
            # change halfway across the pixel
            value = self.get_data_xy(int(data_x + 0.5), int(data_y + 0.5))

        except Exception as e:
            value = None

        system = settings.get('wcs_coords', None)
        format = settings.get('wcs_display', 'sexagesimal')
        ra_lbl, dec_lbl = chr(945), chr(948)

        # Calculate WCS coords, if available
        try:
            if self.wcs is None:
                self.logger.debug("No WCS for this image")
                ra_txt = dec_txt = 'NO WCS'

            elif self.wcs.coordsys == 'raw':
                self.logger.debug("No coordinate system determined")
                ra_txt = dec_txt = 'NO WCS'

            elif self.wcs.coordsys == 'pixel':
                args = [data_x, data_y] + self.revnaxis
                x, y = self.wcs.pixtosystem(args, system=system, coords='data')
                ra_txt = "%+.3f" % (x)
                dec_txt = "%+.3f" % (y)
                ra_lbl, dec_lbl = "X", "Y"

            else:
                args = [data_x, data_y] + self.revnaxis

                lon_deg, lat_deg = self.wcs.pixtosystem(
                    args, system=system, coords='data')

                if format == 'sexagesimal':
                    if system in ('galactic', 'ecliptic'):
                        sign, deg, min, sec = wcs.degToDms(lon_deg,
                                                           isLatitude=False)
                        ra_txt = '+%03d:%02d:%06.3f' % (deg, min, sec)
                    else:
                        deg, min, sec = wcs.degToHms(lon_deg)
                        ra_txt = '%02d:%02d:%06.3f' % (deg, min, sec)

                    sign, deg, min, sec = wcs.degToDms(lat_deg)
                    if sign < 0:
                        sign = '-'
                    else:
                        sign = '+'
                    dec_txt = '%s%02d:%02d:%06.3f' % (sign, deg, min, sec)

                else:
                    ra_txt = '%+10.7f' % (lon_deg)
                    dec_txt = '%+10.7f' % (lat_deg)

                if system == 'galactic':
                    ra_lbl, dec_lbl = "l", "b"
                elif system == 'ecliptic':
                    ra_lbl, dec_lbl = chr(0x03BB), chr(0x03B2)
                elif system == 'helioprojective':
                    ra_txt = "%+5.3f" % (lon_deg * 3600)
                    dec_txt = "%+5.3f" % (lat_deg * 3600)
                    ra_lbl, dec_lbl = "x-Solar", "y-Solar"

        except Exception as e:
            self.logger.warning("Bad coordinate conversion: %s" % (
                str(e)))
            ra_txt = dec_txt = 'BAD WCS'
            try:
                # log traceback, if possible
                (type_, value_, tb) = sys.exc_info()
                tb_str = "".join(traceback.format_tb(tb))
                self.logger.error("Traceback:\n%s" % (tb_str))
            except Exception:
                tb_str = "Traceback information unavailable."
                self.logger.error(tb_str)

        info = Bunch.Bunch(itype='astro', data_x=data_x, data_y=data_y,
                           x=data_x, y=data_y,
                           ra_txt=ra_txt, dec_txt=dec_txt,
                           ra_lbl=ra_lbl, dec_lbl=dec_lbl,
                           value=value)
        return info

# END
