MAX_SIZE = 1<<24    # 16Mio
RED = b"\xff\x00\x00"
GREEN = b"\x00\xff\x00"


def next_token(data, start):
    # skip blanks and comments
    try:
        while data[start] in b" \t\n\r#":
            if data[start] in b" \t\n\r":
                # blank char
                start += 1
                continue
            # skip comment
            while data[start] not in b"\n\r":
                start += 1

        end = start
        while data[end] not in b" \t\n\r":
            end += 1
        return start, end
    except IndexError:
        raise ValueError("unexpected end of file")


def read_image(filename):
    f = open(filename, mode="rb")
    data = f.read(MAX_SIZE)

    i = j = 0
    i, j = next_token(data, j)
    if data[i:j] != b"P6":
        raise ValueError("can only read binary ppm files")

    # size
    i, j = next_token(data, j)
    width = int(data[i:j])
    i, j = next_token(data, j)
    height = int(data[i:j])

    # max intensity
    i, j = next_token(data, j)
    if int(data[i:j]) != 255:
        raise ValueError("can only read ppm files with maximal intensity of 255")

    size = width * height
    if 3*size != len(data) - j - 1:
        raise ValueError("wrong size")

    img = []
    for y in range(height):
        L = []
        for x in range(width):
            idx = j+1 + 3 * (x + y*width)
            L.append(data[idx:idx+3])
        img.append(L)

    return img


class Image(list):

    height = 0
    width = 0

    def __init__(self, filename, mask=None):
        list.__init__(self, read_image(filename))
        self.height = len(self)
        self.width = len(self[0])

        # all pixels are offset so that we can have some
        #  - "keep" pixels (green pixels in the mask) : those have energy
        #  2*self.energy_offset
        #  - "remove" pixels (red pixels in the mask): those have energy 0
        #  - "normal" pixels (all other pixels): those have their normal
        #  energy, offset by self.energy_offset
        # The following should guarantee that "keep" pixels are neither part
        # of a minimal path and that minimal paths always contain at least one
        # "remove" pixel (provided there are some).
        self.energy_offset = (Image.max_energy() + 1) * self.height

        if mask is not None:
            self.mask = read_image(mask)
            assert self.height == len(self.mask)
            assert self.width == len(self.mask[0])
        else:
            self.mask = None
            self.energy_offset = 0

        self.E = [[0]*self.width for _ in range(self.height)]
        for i in range(self.width):
            for j in range(self.height):
                self.E[j][i] = self._energy(i, j)

    @staticmethod
    def _dist(p0, p1):
        r = p0[0] - p1[0]
        g = p0[1] - p1[1]
        b = p0[2] - p1[2]
        return (r*r + g*g + b*b)

    @staticmethod
    def max_energy():
        return 2*Image._dist([0, 0, 0], [255, 255, 255]) + 1

    def _energy(self, x, y):
        if self.mask is not None and self.mask[y][x] == RED:
            return 0
        elif self.mask is not None and self.mask[y][x] == GREEN:
            return self.max_energy() + 2*self.energy_offset

        if x == 0:
            left = self[y][x+1]     # to negate the right pixel
        else:
            left = self[y][x-1]

        if x == self.width-1:
            right = self[y][x-1]    # to negate the left pixel
        else:
            right = self[y][x+1]

        if y == 0:
            top = self[y+1][x]      # to negate the bottom pixel
        else:
            top = self[y-1][x]

        if y == self.height-1:
            bot = self[y-1][x]      # to negate the top pixel
        else:
            bot = self[y+1][x]

        e = self._dist(left, right) + self._dist(top, bot)

        return e + self.energy_offset

    def draw_seam(self, S, color=(255, 0, 0)):
        for y, x in enumerate(S):
            self[y][x] = color

    def carve_seam(self, S):
        for y, x in enumerate(S):
            del self[y][x]
            if self.mask is not None:
                del self.mask[y][x]
            del self.E[y][x]
        self.width -= 1

        for y, x in enumerate(S):
            for v in range(max(0, x-1), min(self.width, x+2)):
                self.E[y][v] = self._energy(v, y)

    def save(self, filename):
        if filename[-4:] != ".ppm":
            filename += ".ppm"
        f = open(filename, 'wb')
        f.write(b"P6\n")
        f.write(bytes("{} {}\n255\n".format(self.width, self.height), encoding="ASCII"))
        for y in range(self.height):
            for x in range(self.width):
                f.write(bytes(self[y][x]))
        f.close()

    def save_mask(self, filename):
        if self.mask is None:
            return
        if filename[-4:] != ".ppm":
            filename += ".ppm"
        f = open(filename, 'wb')
        f.write(b"P6\n")
        f.write(bytes("{} {}\n255\n".format(self.width, self.height), encoding="ASCII"))
        for y in range(self.height):
            for x in range(self.width):
                f.write(bytes(self.mask[y][x]))
        f.close()

    def save_E(self, filename):
        if filename[-4:] != ".ppm":
            filename += ".ppm"
        f = open(filename, 'wb')
        f.write(b"P6\n")
        f.write(bytes("{} {}\n255\n".format(self.width, self.height), encoding="ASCII"))
        m = self.max_energy()
        for y in range(self.height):
            for x in range(self.width):
                if self.mask is not None and self.mask[y][x] == RED:
                    f.write(RED)
                elif self.mask is not None and self.mask[y][x] == GREEN:
                    f.write(GREEN)
                else:
                    e = (self.E[y][x]-self.energy_offset) / m
                    g = int(256 * e)
                    f.write(bytes([g, g, g]))
        f.close()