import os

import numpy as np
from PIL import Image

__all__ = ['list_images', 'ColorMode', 'color_convert', 'image_to_array', 'array_to_image']


def list_images(directory, file_format=('jpg', 'jpeg', 'bmp', 'png', 'ppm', 'tif', 'tiff')):
    """Lists all pictures in a directory, including all subdirectories.
    
    Args:
        directory: string, absolute path to the directory.
        file_format: tuple of strings or single string, extensions of the pictures.
    Returns:
        a list of image paths.
    """
    file_format = tuple('.%s' % e for e in ((file_format,) if isinstance(file_format, str) else file_format))
    return [os.path.join(root, f) for root, _, files in os.walk(directory) for f in files if f.lower().endswith(file_format)]


class color_mode:
    """image color type."""
    O     = {'mode':'1',      'description':'1-bit pixels, black and white, stored with one pixel per byte'}
    L     = {'mode':'L',      'description':'8-bit pixels, black and white'}
    LA    = {'mode':'LA',     'description':'L with alpha'}
    La    = {'mode':'La',     'description':'L with premultiplied alpha'}
    I     = {'mode':'I',      'description':'32-bit signed integer pixels'}
    F     = {'mode':'F',      'description':'32-bit floating point pixels'}
    P     = {'mode':'P',      'description':'8-bit pixels, mapped to any other mode using a color palette'}
    PA    = {'mode':'PA',     'description':'P with alpha'}
    RGB   = {'mode':'RGB',    'description':'3x8-bit pixels, true color'}
    RGBA  = {'mode':'RGBA',   'description':'4x8-bit pixels, true color with transparency mask'}
    RGBX  = {'mode':'RGBX',   'description':'true color with padding'}
    RGBa  = {'mode':'RGBa',   'description':'true color with premultiplied alpha'}
    CMYK  = {'mode':'CMYK',   'description':'4x8-bit pixels, color separation'}
    YCbCr = {'mode':'YCbCr',  'description':'3x8-bit pixels, color video format'}
    LAB   = {'mode':'LAB',    'description':'3x8-bit pixels, the L*a*b color space'}
    HSV   = {'mode':'HSV',    'description':'3x8-bit pixels, Hue, Saturation, Value color space'}
    I16   = {'mode':'I;16',   'description':'16-bit unsigned integer pixels'}
    I16L  = {'mode':'I;16L',  'description':'16-bit little endian unsigned integer pixels'}
    I16B  = {'mode':'I;16B',  'description':'16-bit big endian unsigned integer pixels'}
    I16N  = {'mode':'I;16N',  'description':'16-bit native endian unsigned integer pixels'}
    BGR15 = {'mode':'BGR;15', 'description':'15-bit reversed true colour'}
    BGR16 = {'mode':'BGR;16', 'description':'16-bit reversed true colour'}
    BGR24 = {'mode':'BGR;24', 'description':'24-bit reversed true colour'}
    BGR32 = {'mode':'BGR;32', 'description':'32-bit reversed true colour'}
    grayscale = {'mode':'grayscale',  'description':'8-bit pixels, black and white'}

ColorMode = color_mode()


def color_convert(image, color_mode=ColorMode.RGB):
    """Transform image color mode.
    
    Args:
        image: PIL Image instance.
        color_mode: Image color mode, more see api "la.image.ColorMode".
    Returns:
        PIL instance.
    Raises:
        ValueError: color_mode error.
    """
    if isinstance(color_mode, dict):
        if color_mode['mode'] == 'grayscale':
            if image.mode not in ('L', 'I;16', 'I'):
                image = image.convert('L')
        elif image.mode != color_mode['mode']:
            image = image.convert(color_mode['mode'])
    else:
        raise ValueError('color_mode error.')
    return image


def image_to_array(image, data_format='HWC', dtype='float32'):
    """Converts a PIL Image instance to a Numpy array.
    
    Args:
        image: PIL instance.
        data_format: array data format, eg.'HWC'.
        dtype: Dtype to use for the returned array.
    Returns:
        A Numpy array.
    Raises:
        ValueError: if invalid `img` or `data_format` is passed.
    """
    transpose = {'H':0, 'W':1, 'C':2}
    x = np.asarray(image, dtype=dtype)
    if len(x.shape) == 2:
        x = np.expand_dims(x, axis=-1)
    if len(x.shape) == 3:
        x = x.transpose(tuple(transpose[i] for i in data_format))
    else:
        raise ValueError('Unsupported image shape: %s' % (x.shape,))
    return x


def array_to_image(x, data_format='HWC'):
    """Converts a 3D Numpy array to a PIL Image instance.
    
    Args:
        x: Input Numpy array.
        data_format: array data format, eg.'HWC'.
    Returns:
        A PIL instance.
    Raises:
        ValueError: if invalid `x` or `data_format` is passed.
    """
    if x.ndim != 3:
        raise ValueError('Expected image array to have rank 3 (single image). '
                         'Got array with shape: %s' % (x.shape,))

    transpose = {i:r for r,i in enumerate(data_format)}
    if data_format != 'HWC':
        x = x.transpose(tuple(transpose[i] for i in 'HWC'))

    if x.shape[2] == 4:
        return Image.fromarray(x.astype('uint8'), 'RGBA')
    elif x.shape[2] == 3:
        return Image.fromarray(x.astype('uint8'), 'RGB')
    elif x.shape[2] == 1:
        if np.max(x) > 255:
            return Image.fromarray(x[:, :, 0].astype('int32'), 'I')
        return Image.fromarray(x[:, :, 0].astype('uint8'), 'L')
    else:
        raise ValueError('Unsupported channel number: %s' % (x.shape[2],))
