"""
.. todo::

    WRITEME
"""
import logging
import sys
import numpy
from theano.compat.six.moves import xrange
from pylearn2.utils.image import Image, ensure_Image


logger = logging.getLogger(__name__)


def scale_to_unit_interval(ndar,eps=1e-8):
    """
    .. todo::

        WRITEME
    """
    ndar = ndar.copy()
    ndar -= ndar.min()
    ndar *= 1.0 / max(ndar.max(),eps)
    return ndar


def tile_raster_images(X, img_shape,
        tile_shape=None, tile_spacing=(1,1),
        scale_rows_to_unit_interval=True,
        output_pixel_vals=True,
        min_dynamic_range=1e-4,
        ):
    """
    Transform an array with one flattened image per row, into an array in which
    images are reshaped and layed out like tiles on a floor.

    This function is useful for visualizing datasets whose rows are images, and
    also columns of matrices for transforming those rows (such as the first
    layer of a neural net).

    Parameters
    ----------
    X : numpy.ndarray or tuple of 4 channels or None
        A 2-D array in which every row is a flattened image.
    img_shape : tuple
        The original shape of each image
    tile_shape: tuple
        The number of images to tile (rows, cols). Defaults to a square-ish \
        shape with the right area for the number of images.
    min_dynamic_range: float, positive
        Dynamic range of each image is used in scaling to the unit interval, \
        but images with less dynamic range than this will be scaled as if \
        this were the dynamic range.

    Returns
    -------
    out_array : 2D array with same dtype as X
        Array suitable for viewing as an image (See:`PIL.Image.fromarray`).
    """
    # This is premature when tile_slices_to_image is not documented at all yet,
    # but ultimately true:
    #print >> sys.stderr, "WARN: tile_raster_images sucks, use tile_slices_to_image"
    if len(img_shape)==3 and img_shape[2]==3:
        # make this save an rgb image
        if scale_rows_to_unit_interval:
            logger.warning("tile_raster_images' scaling routine "
                           "messes up colour - try tile_slices_to_image")
        return tile_raster_images(
                (X[:,0::3], X[:,1::3], X[:,2::3], None),
                img_shape=img_shape[:2],
                tile_shape=tile_shape,
                tile_spacing=tile_spacing,
                scale_rows_to_unit_interval=scale_rows_to_unit_interval,
                output_pixel_vals=output_pixel_vals,
                min_dynamic_range=min_dynamic_range)

    if isinstance(X, tuple):
        n_images_in_x = X[0].shape[0]
    else:
        n_images_in_x = X.shape[0]

    if tile_shape is None:
        tile_shape = most_square_shape(n_images_in_x)

    assert len(img_shape) == 2
    assert len(tile_shape) == 2
    assert len(tile_spacing) == 2

    #out_shape is the shape in pixels of the returned image array
    out_shape = [(ishp + tsp) * tshp - tsp for ishp, tshp, tsp
        in zip(img_shape, tile_shape, tile_spacing)]

    if isinstance(X, tuple):
        if scale_rows_to_unit_interval:
            raise NotImplementedError()
        assert len(X) == 4
        if output_pixel_vals:
            out_array = numpy.zeros((out_shape[0], out_shape[1], 4), dtype='uint8')
        else:
            out_array = numpy.zeros((out_shape[0], out_shape[1], 4), dtype=X.dtype)

        #colors default to 0, alpha defaults to 1 (opaque)
        if output_pixel_vals:
            channel_defaults = [0,0,0,255]
        else:
            channel_defaults = [0.,0.,0.,1.]

        for i in xrange(4):
            if X[i] is None:
                out_array[:,:,i] = numpy.zeros(out_shape,
                        dtype='uint8' if output_pixel_vals else out_array.dtype
                        )+channel_defaults[i]
            else:
                out_array[:,:,i] = tile_raster_images(X[i], img_shape, tile_shape, tile_spacing, scale_rows_to_unit_interval, output_pixel_vals)
        return out_array

    else:
        H, W = img_shape
        Hs, Ws = tile_spacing

        out_scaling = 1
        if output_pixel_vals and str(X.dtype).startswith('float'):
            out_scaling = 255

        out_array = numpy.zeros(out_shape, dtype='uint8' if output_pixel_vals else X.dtype)
        for tile_row in xrange(tile_shape[0]):
            for tile_col in xrange(tile_shape[1]):
                if tile_row * tile_shape[1] + tile_col < X.shape[0]:
                    if scale_rows_to_unit_interval:
                        try:
                            this_img = scale_to_unit_interval(
                                    X[tile_row * tile_shape[1] + tile_col].reshape(img_shape),
                                    eps=min_dynamic_range)
                        except ValueError:
                            raise ValueError('Failed to reshape array of shape %s to shape %s'
                                    % (
                                        X[tile_row*tile_shape[1] + tile_col].shape
                                        , img_shape
                                        ))
                    else:
                        this_img = X[tile_row * tile_shape[1] + tile_col].reshape(img_shape)
                    out_array[
                        tile_row * (H+Hs):tile_row*(H+Hs)+H,
                        tile_col * (W+Ws):tile_col*(W+Ws)+W
                        ] \
                        = this_img * out_scaling
        return out_array


def most_square_shape(N):
    """
    Return a rectangle (height, width) with area N that is closest to square.

    Parameters
    ----------
    N : int
        WRITEME

    Returns
    -------
    WRITEME
    """
    for i in xrange(int(numpy.sqrt(N)),0, -1):
        if 0 == N % i:
            return (i, N/i)


def save_tiled_raster_images(tiled_img, filename):
    """
    Save a a return value from `tile_raster_images` to `filename`.

    Returns
    -------
    img : WRITEME
        The PIL image that was saved
    """
    if tiled_img.ndim==2:
        ensure_Image()
        img = Image.fromarray( tiled_img, 'L')
    elif tiled_img.ndim==3:
        ensure_Image()
        img = Image.fromarray(tiled_img, 'RGBA')
    else:
        raise TypeError('bad ndim', tiled_img)

    img.save(filename)
    return img


def tile_slices_to_image_uint8(X, tile_shape=None):
    """
    .. todo::

        WRITEME
    """
    if str(X.dtype) != 'uint8':
        raise TypeError(X)
    if tile_shape is None:
        #how many tile rows and cols
        (TR, TC) = most_square_shape(X.shape[0])
    H, W = X.shape[1], X.shape[2]

    Hs = H+1 #spacing between tiles
    Ws = W+1 #spacing between tiles

    trows, tcols= most_square_shape(X.shape[0])
    outrows = trows * Hs - 1
    outcols = tcols * Ws - 1
    out = numpy.zeros((outrows, outcols,3), dtype='uint8')
    tr_stride= 1+X.shape[1]
    for tr in range(trows):
        for tc in range(tcols):
            Xrc = X[tr*tcols+tc]
            if Xrc.ndim==2: # if no color channel make it broadcast
                Xrc=Xrc[:,:,None]
            #print Xrc.shape
            #print out[tr*Hs:tr*Hs+H,tc*Ws:tc*Ws+W].shape
            out[tr*Hs:tr*Hs+H,tc*Ws:tc*Ws+W] = Xrc
    ensure_Image()
    img = Image.fromarray(out, 'RGB')
    return img


def tile_slices_to_image(X,
        tile_shape=None,
        scale_each=True,
        min_dynamic_range=1e-4):
    """
    .. todo::

        WRITEME
    """
    #always returns an RGB image
    def scale_0_255(x):
        xmin = x.min()
        xmax = x.max()
        return numpy.asarray(
                255 * (x - xmin) / max(xmax - xmin, min_dynamic_range),
                dtype='uint8')

    if scale_each:
        uintX = numpy.empty(X.shape, dtype='uint8')
        for i, Xi in enumerate(X):
            uintX[i] = scale_0_255(Xi)
        X = uintX
    else:
        X = scale_0_255(X)
    return tile_slices_to_image_uint8(X, tile_shape=tile_shape)
