""" 
Basic types:
file - a png file on disk
image - a list of list of pixels. pixels can be triples of
RGB intensities, or single grayscale values.
vec - a vector with domain {0..width-1}x{0..height-1}
display - not a type per se, but rather causing the type to
be shown on screen

Functions convert between these formats, and also can write
to temporary files and display them with a web browser.
"""
import png
import vec

# Native imports

import tempfile
import os
import atexit

# Round color coordinate to nearest int and clamp to [0, 255]
def color_int(col):
    return max(min(round(col), 255), 0)

# utility conversions, between boxed pixel and flat pixel formats
# the png library uses flat, we use boxed.
def boxed2flat(row):
    return [color_int(x) for box in row for x in box]

def flat2boxed(row):
    # Note we skip every 4th element, thus eliminating the alpha
    # channel
    return [tuple(row[i:i+3]) for i in range(0, len(row), 4)]

## Image conversions
def isgray(image):
    return type(image[0][0]) == int

def color2gray(image):
    """ Converts a color image to grayscale """
    # we use HDTV grayscale conversion as per
    # https://en.wikipedia.org/wiki/Grayscale
    return [[int(0.2126*p[0] + 0.7152*p[1] + 0.0722*p[2])
             for p in row]
             for row in image]

def gray2color(image):
    """ Converts a grayscale image to color """
    return [[(p,p,p) for p in row] for row in image]

#extracting and combining color channels    
def rgbsplit(image):
	""" Converts an RGB image to a 3-element list of grayscale
  images, one for each color channel"""
	return [[[pixel[i] for pixel in row] for row in image]
         for i in (0,1,2)]

def rgpsplice(R,G,B):
	return [[(R[row][col],G[row][col],B[row][col])
          for col in range(len(R[0]))] for row in range(len(R))]
	
## To and from files
def file2image(path):
    """ Reads an image into a list of lists of pixel values
    (tuples with three values). This is a color image. """
    (w, h, p, m) = png.Reader(filename = path).asRGBA() # force RGB and alpha
    return [flat2boxed(r) for r in p]


def image2file(image, path):
    """ Writes an image in list of lists format to a file.
    Will work with either color or grayscale. """
    if isgray(image):
        img = gray2color(image)
    else:
        img = image
    with open(path, 'wb') as f:
        png.Writer(width=len(image[0]), height=len(image)).write(f, 
            [boxed2flat(r) for r in img])

## To and from vecs
def image2vec(image):
    """ Converts an image in list of lists format to a vector.
    Will work with either color or grayscale. """
    if isgray(image):
        D = {(x,y) for x in range(len(image[0]))
                   for y in range(len(image))}
        F = {(x,y):image[y][x] for (x,y) in D}
    else:
        D = {(x,y,c) for c in ['r','g','b']
                     for x in range(len(image[0]))
                     for y in range(len(image))}
        F = dict()
        for y in range(len(image)):
            for x in range(len(image[y])):
                F[(x,y,'r')] = image[y][x][0]
                F[(x,y,'g')] = image[y][x][1]
                F[(x,y,'b')] = image[y][x][2]
    return vec.Vec(D, F)

def vec2image(vec):
    """ Converts a vector to an image in list of lists format """
    image = []
    width = max(vec.D, key=lambda p: p[0])[0]
    height = max(vec.D, key=lambda p: p[1])[1]
    # check if grayscale
    e = vec.D.pop()
    vec.D.add(e)
    gray = len(e) == 2
    for y in range(height):
      row = []
      for x in range(width):
        if gray:
          row += [vec[(x,y)]]
        else:
          row += [(vec[(x,y,'r')], vec[(x,y,'g')], vec[(x,y,'b')])]
      image += [row]
    return image
    
## Shortcuts - files to vecs and vice-versa

def file2vec(path):
    """ Reads an image from a file and turns it into a vector """
    return image2vec(file2image(path))

def vec2file(vec, path):
    """ Reads an image from a file and turns it into a vector """
    image2file(vec2image(vec), path)
