aiSandbox/environment.py

#!/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 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 row in state.matrix:
        el.extend(row)
    return all(el[i] < el[i+1] for i in range(len(el)-1))

if __name__ == "__main__":
    s0 = State()
    for action in actions(s0):
        result(s0, action).printState()