
from PIL import Image as Image

def compute_energy(image):
    """Inputs
         image: a color image of the dictionary form described
       
       Returns a map of the same size as `image`, containing values
       which describe the presence or absence of edges at each pixel
       location in the image (higher numerical values indicate the 
       presence of edges).
    """
    grey_map = get_greyscale_map(image)
    return edges(grey_map)

def get_greyscale_map(image):
    """
    Returns a greyscale version of the given color image. 

    Formula for rgb --> grey conversion adapted from:
    https://en.wikipedia.org/wiki/Grayscale#Converting_color_to_grayscale
    """
    return {
        'height': image['height'],
        'width': image['width'],
        'pixels': [round(.299*r + .587*g + .114*b) 
                        for r,g,b in image['pixels']],
    }

def edges(grey_map):
    """
    Return a new map representing the result of applying the Sobel 
    edge detection operator to the given greyscale input map.

    This edge detection is described here, for those curious:
    https://py.mit.edu/spring20/labs/lab0#_edge_detection
    """
    sobel_x = [[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]
    sobel_y = [[-1, -2, -1], [0, 0, 0], [1, 2, 1]]
    # Correlating with the two kernels above results in two images: one
    # that has the vertical edges only, and one that has the horizontal
    # edges only.
    rx = correlate(grey_map, sobel_x)
    ry = correlate(grey_map, sobel_y)
    # Now loop over the pixels and combine
    out = empty_like(grey_map)
    for x, y in get_coords(grey_map):
        val = (get_pixel(rx, x, y)**2 + get_pixel(ry, x, y)**2)**0.5
        set_pixel(out, x, y, clip(round(val), 0, 255))
    return out

def correlate(grey_map, kernel):
    """
    Compute the result of correlating the given image with the given kernel.

    Correlation is described here, for those curious: 
    https://py.mit.edu/spring20/labs/lab0#_image_filtering_via_correlation

    Does not mutate the input map; rather, creates a separate structure to 
    represent the output.

        kernel: 2-d array (list of lists).
    """
    out = empty_like(grey_map)
    n = len(kernel)
    kern_size = n//2  # the 'half-width' of the kernel
    for x, y in get_coords(grey_map):
        # for each pixel, loop over the surrounding pixels, multiplying
        # each by the associated value in the kernel and accumulating
        # the results.
        v = 0
        for kx in range(n):
            for ky in range(n):
                px = x - kern_size + kx
                py = y - kern_size + ky
                v += get_pixel_extended(grey_map, px, py) * kernel[ky][kx]
        # finally, set the pixel to the accumulated value.
        set_pixel(out, x, y, v)
    return out

def empty_like(map_, val=0):
    """
    Helper function.  Given a map or an image, create a new map of the same 
    shape, but filled with all val.
    """
    return {
        'width': map_['width'],
        'height': map_['height'],
        'pixels': [val for _ in map_['pixels']],
    }


def get_pixel(image, x, y):
    return image['pixels'][flat_index(image, x, y)]


def set_pixel(image, x, y, c):
    image['pixels'][flat_index(image, x, y)] = c


def get_coords(image):
    """
    Helper function. Yields all the (x, y) tuple coordinates on the image.
    
    If you're unfamiliar with `yield` or, more generally, generators, you 
    may find them interesting to investigate, for example, here: 
    https://wiki.python.org/moin/Generators
    """
    for y in range(image['height']):
        for x in range(image['width']):
            yield x, y

def get_pixel_extended(image, x, y):
    """
    Helper function.  Returns the pixel value at the given coordinate or,
    if the coordinate is out of bounds, the closest coordinate in bounds.
    """
    return get_pixel(image,
                     clip(x, 0, image['width']-1),
                     clip(y, 0, image['height']-1))

def flat_index(image, x, y):
    """
    Helper function.  Given an image and x and y coordinates, return the
    integer index into image['pixels'] associated with that (x,y) location in
    the image.
    """
    return y * image['width'] + x

def clip(x, lo, hi):
    """
    Helper function. Returns `x` if it's within the range from
    `lo` to `hi` (inclusive); otherwise returns the cloeser of 
    `lo` and `hi`.
    """
    return max(min(x, hi), lo)

################################################################
# HELPER FUNCTIONS FOR LOADING AND SAVING COLOR IMAGES

def load_color_image(filename):
    """
    Loads an image from the given file and returns a dictionary
    representing that image.

    Invoked as, for example:
       i = load_image('test_images/cat.png')
    """
    with open(filename, 'rb') as img_handle:
        img = Image.open(img_handle)
        img = img.convert('RGB')  # in case we were given a greyscale image
        img_data = img.getdata()
        pixels = list(img_data)
        w, h = img.size
        return {'height': h, 'width': w, 'pixels': pixels}


def save_color_image(image, filename, mode='PNG'):
    """
    Saves the given image to disk or to a file-like object.  If filename is
    given as a string, the file type will be inferred from the given name.  If
    filename is given as a file-like object, the file type will be determined
    by the 'mode' parameter.
    """
    out = Image.new(mode='RGB', size=(image['width'], image['height']))
    out.putdata(image['pixels'])
    if isinstance(filename, str):
        out.save(filename)
    else:
        out.save(filename, mode)
    out.close()
