"""
image HDU classes for fitslib, part of the fitsio package.

See the main docs at https://github.com/esheldon/fitsio

  Copyright (C) 2011  Erin Sheldon, BNL.  erin dot sheldon at gmail dot com

    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 St, Fifth Floor, Boston, MA  02110-1301  USA

"""

from __future__ import with_statement, print_function

import numpy

from math import floor
from .base import HDUBase, IMAGE_HDU
from ..util import (
    IS_PY3,
    array_to_native,
    copy_if_needed,
    _nonfinite_as_cfitsio_floating_null_value,
)

# for python3 compat
if IS_PY3:
    xrange = range


class ImageHDU(HDUBase):
    def _update_info(self):
        """
        Call parent method and make sure this is in fact a
        image HDU.  Set dims in C order
        """
        super(ImageHDU, self)._update_info()

        if self._info['hdutype'] != IMAGE_HDU:
            mess = "Extension %s is not a Image HDU" % self.ext
            raise ValueError(mess)

        # convert to c order
        if 'dims' in self._info:
            self._info['dims'] = list(reversed(self._info['dims']))

    def has_data(self):
        """
        Determine if this HDU has any data

        For images, check that the dimensions are not zero.

        For tables, check that the row count is not zero
        """
        ndims = self._info.get('ndims', 0)
        if ndims == 0:
            return False
        else:
            return True

    def is_compressed(self):
        """
        returns true of this extension is compressed
        """
        return self._info['is_compressed_image'] == 1

    def get_comptype(self):
        """
        Get the compression type.

        None if the image is not compressed.
        """
        return self._info['comptype']

    def get_dims(self):
        """
        get the shape of the image.  Returns () for empty
        """
        if self._info['ndims'] != 0:
            dims = self._info['dims']
        else:
            dims = ()

        return dims

    def reshape(self, dims):
        """
        reshape an existing image to the requested dimensions

        If the new shape is bigger than the current shape,
        the existing values in the image are "wrapped" around in C
        order, via the process of

            1. flattening the image in C order
            2. appending zeros to the image so that it matches the new
               total size
            3. reshaping the image to the new dimensions

        If the new shape is smaller than the current image, the current
        image is flattened, trunctaed to the new total length, and then
        reshaped to the new shape.

        parameters
        ----------
        dims: sequence
            Any sequence convertible to i8
        """

        adims = numpy.array(dims, ndmin=1, dtype='i8')
        # we have to reverse the dimensions here since cfitsio
        # uses fortran order
        self._FITS.reshape_image(self._ext + 1, adims[::-1])
        self._cached_info = None  # invalidate info cache

    def write(self, img, start=0, **keys):
        """
        Write the image into this HDU

        If data already exist in this HDU, they will be overwritten.  If the
        image to write is larger than the image on disk, or if the start
        position is such that the write would extend beyond the existing
        dimensions, the on-disk image is expanded.

        parameters
        ----------
        img: ndarray
            A simple numpy ndarray
        start: integer or sequence
            Where to start writing data.  Can be an integer offset
            into the entire array, or a sequence determining where
            in N-dimensional space to start.
        """

        if keys:
            import warnings

            warnings.warn(
                "The keyword arguments '%s' are being ignored! This warning "
                "will be an error in a future version of `fitsio`!" % keys,
                DeprecationWarning,
                stacklevel=2,
            )

        if img.dtype.fields is not None:
            raise ValueError("got recarray, expected regular ndarray")
        if img.size == 0:
            raise ValueError("data must have at least 1 row")

        # data must be c-contiguous and native byte order
        if not img.flags['C_CONTIGUOUS']:
            # this always makes a copy
            img_send = numpy.ascontiguousarray(img)
            img_send = array_to_native(img_send, inplace=True)
        else:
            img_send = array_to_native(img, inplace=False)

        if IS_PY3 and img_send.dtype.char == 'U':
            # for python3, we convert unicode to ascii
            # this will error if the character is not in ascii
            img_send = img_send.astype('S', copy=copy_if_needed)

        # see if we need to resize the image
        if self.has_data():
            self._expand_if_needed(self.get_dims(), img.shape, start)
            dims = self.get_dims()

            if numpy.isscalar(start):
                start = numpy.unravel_index(start, dims)

            if all(od == nd for od, nd in zip(dims, img.shape)) and all(
                st == 0 for st in start
            ):
                # we are replacing the whole image, so no need to
                # write a subset
                write_subset = False
            else:
                write_subset = True
        else:
            write_subset = False

        with _nonfinite_as_cfitsio_floating_null_value(
            img_send, self.is_compressed()
        ) as img_send_any_nan:
            img_send, any_nan = img_send_any_nan
            if not write_subset:
                # write in image at start in a single pass
                offset = 0
                self._FITS.write_image(
                    self._ext + 1,
                    img_send,
                    offset + 1,
                    1 if any_nan else 0,
                )
            else:
                if not any_nan and not self.is_compressed():
                    firstpixel = numpy.array(start, ndmin=1, dtype='i8')
                    # lastpixel is the index of the lastpixel so subtract 1
                    lastpixel = (
                        firstpixel
                        + numpy.array(img_send.shape, ndmin=1, dtype='i8')
                        - 1
                    )

                    # we have to reverse the dimensions here since cfitsio
                    # uses fortran order and offset by 1 for fortan indexing
                    firstpixel = firstpixel[::-1] + 1
                    lastpixel = lastpixel[::-1] + 1

                    self._FITS.write_subset(
                        self._ext + 1, img_send, firstpixel, lastpixel
                    )
                else:
                    # the C API doesn't support nan handling w/ rectangular
                    # subsets, so emulate in python
                    # go "row by row" but in more than two dimensions
                    ndims = len(dims)
                    for index in numpy.ndindex(*(img_send.shape[:-1])):
                        new_start = [
                            start[i] + index[i] for i in range(ndims - 1)
                        ]
                        new_start += [start[-1]]
                        offset = _convert_full_start_to_offset(dims, new_start)
                        img_slice = tuple(
                            [slice(ns, ns + 1) for ns in index]
                        ) + (slice(None),)
                        self._FITS.write_image(
                            self._ext + 1,
                            img_send[img_slice],
                            offset + 1,
                            1 if any_nan else 0,
                        )

        self._cached_info = None  # invalidate info cache

    def read(self, **keys):
        """
        Read the image.

        If the HDU is an IMAGE_HDU, read the corresponding image.  Compression
        and scaling are dealt with properly.
        """

        if keys:
            import warnings

            warnings.warn(
                "The keyword arguments '%s' are being ignored! This warning "
                "will be an error in a future version of `fitsio`!" % keys,
                DeprecationWarning,
                stacklevel=2,
            )

        if not self.has_data():
            return None

        dtype, shape = self._get_dtype_and_shape()
        array = numpy.zeros(shape, dtype=dtype)
        self._FITS.read_image(self._ext + 1, array)
        return array

    def _get_dtype_and_shape(self):
        """
        Get the numpy dtype and shape for image
        """
        npy_dtype = self._get_image_numpy_dtype()

        if self._info['ndims'] != 0:
            shape = self._info['dims']
        else:
            raise IOError("no image present in HDU")

        return npy_dtype, shape

    def _get_image_numpy_dtype(self):
        """
        Get the numpy dtype for the image
        """
        try:
            ftype = self._info['img_equiv_type']
            npy_type = _image_bitpix2npy[ftype]
        except KeyError:
            raise KeyError("unsupported fits data type: %d" % ftype)

        return npy_type

    def __getitem__(self, arg):
        """
        Get data from an image using python [] slice notation.

        e.g., [2:25, 4:45].
        """
        return self._read_image_slice(arg)

    def _read_image_slice(self, arg):
        """
        workhorse to read a slice
        """
        if 'ndims' not in self._info:
            raise ValueError("Attempt to slice empty extension")

        if isinstance(arg, slice):
            # one-dimensional, e.g. 2:20
            return self._read_image_slice((arg,))

        if not isinstance(arg, tuple):
            raise ValueError(
                "arguments must be slices, one for each "
                "dimension, e.g. [2:5] or [2:5,8:25] etc."
            )

        # should be a tuple of slices, one for each dimension
        # e.g. [2:3, 8:100]
        nd = len(arg)
        if nd != self._info['ndims']:
            raise ValueError(
                "Got slice dimensions %d, "
                "expected %d" % (nd, self._info['ndims'])
            )

        targ = arg
        arg = []
        for a in targ:
            if isinstance(a, slice):
                arg.append(a)
            elif isinstance(a, int):
                arg.append(slice(a, a + 1, 1))
            else:
                raise ValueError("arguments must be slices, e.g. 2:12")

        dims = self._info['dims']
        arrdims = []
        first = []
        last = []
        steps = []
        npy_dtype = self._get_image_numpy_dtype()

        # check the args and reverse dimensions since
        # fits is backwards from numpy
        dim = 0
        for slc in arg:
            start = slc.start
            stop = slc.stop
            step = slc.step

            if start is None:
                start = 0
            if stop is None:
                stop = dims[dim]
            if step is None:
                # Ensure sane defaults.
                if start <= stop:
                    step = 1
                else:
                    step = -1

            # Sanity checks for proper syntax.
            if (
                (step > 0 and stop < start)
                or (step < 0 and start < stop)
                or (start == stop)
            ):
                return numpy.empty(0, dtype=npy_dtype)

            if start < 0:
                start = dims[dim] + start
                if start < 0:
                    raise IndexError("Index out of bounds")

            if stop < 0:
                stop = dims[dim] + start + 1

            # move to 1-offset
            start = start + 1

            if stop > dims[dim]:
                stop = dims[dim]

            if stop < start:
                # A little black magic here.  The stop is offset by 2 to
                # accommodate the 1-offset of CFITSIO, and to move past the end
                # pixel to get the complete set after it is flipped along the
                # axis.  Maybe there is a clearer way to accomplish what this
                # offset is glossing over.
                # @at88mph 2019.10.10
                stop = stop + 2

            first.append(start)
            last.append(stop)

            # Negative step values are not used in CFITSIO as the dimension is
            # already properly calcualted.
            # @at88mph 2019.10.21
            steps.append(abs(step))
            arrdims.append(int(floor((stop - start) / step)) + 1)

            dim += 1

        first.reverse()
        last.reverse()
        steps.reverse()
        first = numpy.array(first, dtype='i8')
        last = numpy.array(last, dtype='i8')
        steps = numpy.array(steps, dtype='i8')

        array = numpy.zeros(arrdims, dtype=npy_dtype)
        self._FITS.read_image_slice(
            self._ext + 1, first, last, steps, self._ignore_scaling, array
        )
        return array

    def _expand_if_needed(self, dims, write_dims, start):
        """
        expand the on-disk image if the indended write will extend
        beyond the existing dimensions
        """
        ndim = len(dims)
        idim = len(write_dims)

        if idim != ndim:
            raise ValueError(
                "When expanding "
                "an existing image while writing, the input image "
                "must have the same number of dimensions "
                "as the original.  "
                "Got %d instead of %d" % (idim, ndim)
            )

        if numpy.isscalar(start):
            if len(dims) > 1:
                try:
                    _start = numpy.unravel_index(start, dims)
                except Exception:
                    # the unravel_index call fails when start is beyond
                    # end of the existing array.
                    # this means we are expanding the image and so we should
                    # error
                    raise ValueError(
                        "When expanding "
                        "an existing image while writing, the start keyword "
                        "must have the same number of dimensions "
                        "as the image or be exactly 0, got %s " % start
                    )
            else:
                _start = [start]
        else:
            _start = start

        new_dims = []
        for i in xrange(ndim):
            required_dim = _start[i] + write_dims[i]

            if required_dim < dims[i]:
                # careful not to shrink the image!
                dimsize = dims[i]
            else:
                dimsize = required_dim

            new_dims.append(dimsize)

        if any(nd != od for nd, od in zip(new_dims, dims)):
            if numpy.isscalar(start) and len(dims) > 1:
                if start != 0:
                    raise ValueError(
                        "When expanding "
                        "an existing image while writing, the start keyword "
                        "must have the same number of dimensions "
                        "as the image or be exactly 0, got %s " % start
                    )
            self.reshape(new_dims)

    def __repr__(self):
        """
        Representation for ImageHDU
        """
        text, spacing = self._get_repr_list()
        text.append("%simage info:" % spacing)
        cspacing = ' ' * 4

        # need this check for when we haven't written data yet
        if 'ndims' in self._info:
            if self._info['comptype'] is not None:
                text.append(
                    "%scompression: %s" % (cspacing, self._info['comptype'])
                )

            if self._info['ndims'] != 0:
                dimstr = [str(d) for d in self._info['dims']]
                dimstr = ",".join(dimstr)
            else:
                dimstr = ''

            dt = _image_bitpix2npy[self._info['img_equiv_type']]
            text.append("%sdata type: %s" % (cspacing, dt))
            text.append("%sdims: [%s]" % (cspacing, dimstr))

        text = '\n'.join(text)
        return text


def _convert_full_start_to_offset(dims, start):
    # convert to scalar offset
    # note we use the on-disk data type to get itemsize
    ndim = len(dims)

    # convert sequence to pixel start
    if len(start) != ndim:
        m = "start has len %d, which does not match requested dims %d"
        raise ValueError(m % (len(start), ndim))

    # MRB: I changed this to use the numpy util below.
    #      I have left the old code here for posterity.
    #      I checked that they give the same answer.
    # # this is really strides / itemsize
    # strides = [1]
    # for i in xrange(1, ndim):
    #     strides.append(strides[i - 1] * dims[ndim - i])

    # strides.reverse()
    # s = start
    # start_index = sum([s[i] * strides[i] for i in xrange(ndim)])

    # return start_index

    return numpy.ravel_multi_index(start, dims)


# remember, you should be using the equivalent image type for this
_image_bitpix2npy = {
    8: 'u1',
    10: 'i1',
    16: 'i2',
    20: 'u2',
    32: 'i4',
    40: 'u4',
    64: 'i8',
    80: 'u8',
    -32: 'f4',
    -64: 'f8',
}