import collections
import numpy as np

import tetrominos

FIELD_WIDTH = 10
FIELD_HEIGHT = 20

COMMAND_NONE = 0
COMMAND_LEFT = 1
COMMAND_RIGHT = 2
COMMAND_ROTATE = 3

def tetromino_to_field_at(tetromino, row, col):
    col_offset = col + tetromino.col_offset
    row_offset = row + tetromino.row_offset
    if col_offset < 0 or row_offset < 0:
        return None
    if col_offset + tetromino.data.shape[1] > FIELD_WIDTH:
        return None
    if row_offset + tetromino.data.shape[0] > FIELD_HEIGHT:
        return None

    array = tetromino.data
    left_zeros = np.zeros(shape=(array.shape[0], col_offset),
                          dtype=np.int32, order="C")
    right_zeros = np.zeros(shape=(array.shape[0],
                                  FIELD_WIDTH-col_offset-array.shape[1]),
                           dtype=np.int32, order="C")
    array = np.concatenate((left_zeros, array, right_zeros), axis=1)

    top_zeros = np.zeros(shape=(row_offset,FIELD_WIDTH),
                         dtype=np.int32, order="C")
    bottom_zeros = np.zeros(shape=(FIELD_HEIGHT - row_offset - array.shape[0],
                                   FIELD_WIDTH),
                            dtype=np.int32, order="C")
    array = np.concatenate((top_zeros, array, bottom_zeros), axis=0)
    return array


GameState = collections.namedtuple('GameState',
                                   ['grid', 'current_piece', 'next_piece',
                                    'rotation', 'row', 'col'])


def initial_game_state(current_piece, next_piece):
    return GameState(grid=np.zeros(shape=(FIELD_HEIGHT,FIELD_WIDTH),
                                   dtype=np.int32, order="C"),
                     current_piece=current_piece,
                     next_piece=next_piece,
                     rotation=0, row=0, col=0)


def is_game_state_valid(game_state):
    tetromino_array = tetromino_to_field_at(
            tetrominos.SHAPES[game_state.current_piece][game_state.rotation],
            game_state.row, game_state.col)
    if tetromino_array is None:
        return False

    total_field = np.add(game_state.grid, tetromino_array)
    if np.any(np.vectorize(lambda x: x > 1)(total_field)):
        return False
    return True


def game_state_field(game_state):
    tetromino_array = tetromino_to_field_at(
            tetrominos.SHAPES[game_state.current_piece][game_state.rotation],
            game_state.row, game_state.col)
    total_field = np.add(game_state.grid, tetromino_array)
    total_field = np.add(total_field, tetromino_array)
    return total_field


def advance_game_state_shift(game_state, col_offset):
    next_game_state = GameState(grid=game_state.grid,
                                current_piece=game_state.current_piece,
                                next_piece=game_state.next_piece,
                                rotation=game_state.rotation,
                                row=game_state.row,
                                col=game_state.col + col_offset)
    if is_game_state_valid(next_game_state):
        return next_game_state
    return None


def advance_game_state_rotate(game_state):
    next_game_state = GameState(grid=game_state.grid,
                                current_piece=game_state.current_piece,
                                next_piece=game_state.next_piece,
                                rotation=(game_state.rotation + 1) % 4,
                                row=game_state.row,
                                col=game_state.col)
    if is_game_state_valid(next_game_state):
        return next_game_state
    return None


def advance_game_state_down(game_state):
    next_game_state = GameState(grid=game_state.grid,
                                current_piece=game_state.current_piece,
                                next_piece=game_state.next_piece,
                                rotation=game_state.rotation,
                                row=game_state.row + 1,
                                col=game_state.col)
    if is_game_state_valid(next_game_state):
        return next_game_state
    return None


def advance_game_state_piece(game_state, next_piece):
    tetromino_array = tetromino_to_field_at(
            tetrominos.SHAPES[game_state.current_piece][game_state.rotation],
            game_state.row, game_state.col)
    next_grid = np.add(game_state.grid, tetromino_array)
    next_game_state = GameState(grid=next_grid,
                                current_piece=game_state.next_piece,
                                next_piece=next_piece,
                                rotation=0, row=0, col=0)
    return next_game_state


def advance_game_state(game_state, command):
    command_game_state = None
    if command == COMMAND_LEFT:
        command_game_state = advance_game_state_shift(game_state, -1)
    elif command == COMMAND_RIGHT:
        command_game_state = advance_game_state_shift(game_state, 1)
    elif command == COMMAND_ROTATE:
        command_game_state = advance_game_state_rotate(game_state)

    if command_game_state is None:
        command_game_state = game_state

    next_game_state = advance_game_state_down(command_game_state)
    if next_game_state is None:
        next_game_state = advance_game_state_piece(command_game_state, 0)
    return next_game_state