#!/usr/bin/env python3

from random import shuffle
from copy import deepcopy

# Example of environment, for the game of fifteen.

# Maximum rows and columns size
MAX_SIZE = 4

class State:
    """Represents a state of the environment"""

    def __init__(self):
        elements = [x for x in range(MAX_SIZE**2)]
        shuffle(elements)
        self.matrix = [[elements[x+(y*MAX_SIZE)] for x in range(MAX_SIZE)] for y in range(MAX_SIZE)]

    def __iter__(self):
        self.row = 0
        self.col = 0
        return self

    def __next__(self):
        if self.row == MAX_SIZE:
            self.row = 0
            self.col = self.col+1
        if self.col == MAX_SIZE:
            raise StopIteration
        elem = self.matrix[self.col][self.row]
        self.row = self.row + 1
        return elem

    def printState(self):
        for i in range(MAX_SIZE):
            print(self.matrix[i])
        print("\n")

def up(state):
    (x, y) = findZero(state)
    (state.matrix[y][x], state.matrix[y-1][x]) = (state.matrix[y-1][x], state.matrix[y][x])

def down(state):
    (x, y) = findZero(state)
    (state.matrix[y][x], state.matrix[y+1][x]) = (state.matrix[y+1][x], state.matrix[y][x])

def left(state):
    (x, y) = findZero(state)
    (state.matrix[y][x], state.matrix[y][x-1]) = (state.matrix[y][x-1], state.matrix[y][x])

def right(state):
    (x, y) = findZero(state)
    (state.matrix[y][x], state.matrix[y][x+1]) = (state.matrix[y][x+1], state.matrix[y][x])

def findZero(state):
    for row in state.matrix:
        try:
            x = row.index(0)
            y = state.matrix.index(row)
            break
        except ValueError:
            pass
    return (x, y)

def actions(state):
    (x, y) = findZero(state)
    actions = []
    if x > 0:
        actions.append(left)
    if x < MAX_SIZE-1:
        actions.append(right)
    if y > 0:
        actions.append(up)
    if y < MAX_SIZE-1:
        actions.append(down)
    return actions

def result(state, action):
    newState = deepcopy(state)
    action(newState)
    return newState

def goalTest(state):
    el = []
    for cell in state:
        el.append(cell)
    return all(el[i] < el[i+1] for i in range(len(el)-1))

def equals(firstState, secondState):
    if firstState == secondState:
        return True
    return all(x[0]==x[1] for x in zip(firstState, secondState))

if __name__ == "__main__":
    pass