__all__ = ['imread', 'imread_collection']

import numpy as np
import scikits.image.io as io

try:
    import pyfits
except ImportError:
    raise ImportError("PyFITS could not be found. Please refer to\n"
        "http://www.stsci.edu/resources/software_hardware/pyfits\n"
        "for further instructions.")


def imread(fname, as_grey=True, dtype=None):
    """Load an image from a FITS file.

    Parameters
    ----------
    fname : string
        Image file name, e.g. ``test.fits``.
    as_grey : bool
        For FITS images, this is ignored (treated as True).
    dtype : dtype, optional
        For FITS, this argument is ignored because Stefan is planning on
        removing the dtype argument from imread anyway.

    Returns
    -------
    img_array : ndarray
        Unlike other plugins where the different colour bands/channels are
        stored in the third dimension, here the native dimensionality defined
        in the FITS file is used.

        Currently if no image is found in the file, None will be returned

    Notes
    -----

    Currently FITS imread() always returns the first image extension when
    given a Multi-Extension FITS file; use imread_collection() (which does
    lazy loading) to get all the extensions at once.

    """

    hdulist = pyfits.open(fname)

    # Iterate over FITS image extensions, ignoring any other extension types
    # such as binary tables, and get the first image data array:
    img_array = None
    for hdu in hdulist:
        if isinstance(hdu, pyfits.ImageHDU) or \
           isinstance(hdu, pyfits.PrimaryHDU):
           if hdu.data is not None:
               img_array = hdu.data
               break
    hdulist.close()

    return img_array


def imread_collection(load_pattern, conserve_memory=True):
    """Load a collection of images from one or more FITS files

       Parameters
       ----------
       load_pattern : str or list
           List of extensions to load. Filename globbing is currently
           unsupported.
       converve_memory : bool
           If True, never keep more than one in memory at a specific
           time. Otherwise, images will be cached once they are loaded.

       Returns
       -------

       ic : ImageCollection
           Collection of images.

    """

    intype = type(load_pattern)
    if intype is not list and intype is not str:
        raise TypeError("Input must be a file or list of files")

    # Ensure we have a list, otherwise we'll end up iterating over the string:
    if intype is not list:
        load_pattern = [load_pattern]

    # Generate a list of filename/extension pairs by opening the list of
    # files and finding the image extensions in each one:
    ext_list = []
    for filename in load_pattern:
        hdulist = pyfits.open(filename)
        for n, hdu in zip(range(len(hdulist)), hdulist):
            if isinstance(hdu, pyfits.ImageHDU) or \
               isinstance(hdu, pyfits.PrimaryHDU):
                # Ignore (primary) header units with no data (use '.size'
                # rather than '.data' to avoid actually loading the image):
                if hdu.size() > 0:
                    ext_list.append((filename, n))
        hdulist.close()

    return io.ImageCollection(ext_list, load_func=FITSFactory,
                              conserve_memory=conserve_memory)


    imread(fname, as_grey=as_grey, dtype=dtype,
                       plugin=plugin, **plugin_args)


def FITSFactory(image_ext):
    """Load an image extension from a FITS file and return a NumPy array
    
    Parameters
    ----------

       image_ext : tuple
           FITS extension to load, in the format (filename, ext_num).
           The (extname, extver) format is unsupported, since this function
           is not called directly by the user and imread_collection() does
           the work of figuring out which extensions need loading.

    """

    # Expect a length-2 tuple with a filename as the first element:
    if not isinstance(image_ext, tuple):
        raise TypeError("Expected a tuple")

    if len(image_ext) != 2:
        raise ValueError("Expected a tuple of length 2")

    filename = image_ext[0]
    extnum = image_ext[1]

    if type(filename) is not str or type(extnum) is not int:
        raise ValueError("Expected a (filename, extension) tuple")

    hdulist = pyfits.open(filename)
        
    data = hdulist[extnum].data

    hdulist.close()

    if data is None:
        raise RuntimeError("Extension %d of %s has no data" %
                  (extnum, filename))

    return data