from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import math
import pdb


def imread(filename):
    """
    Matlab like function for reading an image file.
    Returns:
    im_array (numpy.ndarray): h x w x 3 ndarray for color images and 
        h x w for grayscale images
    """
    im = Image.open(filename)

    err_str = \
        "imread only supports 'RGB' and 'L' modes, found '{}'".format(im.mode)
    assert (im.mode == 'RGB' or im.mode == 'L'), err_str

    im_array = np.array(im)

    return im_array


def imshow(np_im):
    """
    Matlab like function for displaying a numpy ndarray as an image
    Args:
    np_im (numpy.ndarray): h x w x 3 ndarray for color images and 
        h x w for grayscale images with pixels stored in uint8 format
    """
    err_str = 'imshow expects ndarray of dimension h x w x c (RGB) or h x w (L)'
    assert (len(np_im.shape) == 3 or len(np_im.shape) == 2), err_str

    if len(np_im.shape) == 3:
        assert (np_im.shape[2] == 3), 'imshow expected 3 channels'
        im = Image.fromarray(np_im, 'RGB')
    else:
        im = Image.fromarray(np_im, 'L')

    im.show()

def imshow2(np_im):
    imgplot = plt.imshow(np_im)
    plt.show()
    return imgplot


def imwrite(np_im, filename):
    """
    Matlab like function to save numpy image.
    Args:
    np_im (numpy.ndarray): h x w x 3 ndarray for color images and 
        h x w for grayscale images with pixels stored in uint8 format
    """
    err_str = 'imshow expects ndarray of dimension h x w x c (RGB) or h x w (L)'
    assert (len(np_im.shape) == 3 or len(np_im.shape) == 2), err_str

    im = Image.fromarray(np_im)
    im.save(filename)


def imresize(np_im, method=Image.BILINEAR, **kwargs):
    assert_str = "Only 1 keyword argument expected with key either" + \
                 "'output_size' or 'scale'"
    assert (len(kwargs)==1), assert_str
    im_h = np_im.shape[0]
    im_w = np_im.shape[1]
    if 'output_size' in kwargs:
        h, w = kwargs['output_size']
    elif 'scale' in kwargs:
        h = scale*im_h
        w = scale*im_w
    else:
        assert_str = "Variable argument must be one of {'output_size','scale'}"
        assert (False), assert_str
    h = int(math.ceil(h))
    w = int(math.ceil(w))
    im = Image.fromarray(np_im)
    return np.array(im.resize((w,h)))