import random
import numpy as np

from grid import Grid


class Vector(object):
    def __init__(self, x, y):
        self.x = x
        self.y = y


class Traversals(object):
    def __init__(self, x, y):
        self.x = x
        self.y = y
        
# env types:
# 0 - normal
# 1 - reward in [-1, 0, 1]
# 2 - finish game when grid was not moved
class Game2048(object):
    
    def __init__(self, size, env_type=0):
        self.size = size
        self.start_tiles = 2
        self.action_space = 4
        self.observation_space = 16
        self.env_type = env_type

        self.over = False
        self.won = False
        self.grid = Grid(self.size)
        self.score = 0
        self.win_count = 0

    def reset(self):
        self.over = False
        self.won = False
        self.grid = Grid(self.size)
        self.score = 0

        self._add_start_tiles()

        return self.grid.get_values()

    def _add_start_tiles(self):
        for x in range(self.start_tiles):
            self._add_rendom_tile()

    def _add_rendom_tile(self):
        value = 2 if random.uniform(0, 1) < 0.9 else 4
        self.grid.insert_value_in_random_available_cell(value)

    def _get_vector(self, direction):
        return [Vector(0, -1), Vector(1, 0), Vector(0, 1), Vector(-1, 0)][direction]

    def _build_traversals(self ,vector):
        x = range(self.size)[::-1] if vector.x == 1 else range(self.size)
        y = range(self.size)[::-1] if vector.y == 1 else range(self.size)

        return Traversals(x, y)

    def _prepare_tiles(self):
        for row in self.grid:
            for cell in row:
                cell.merged = False

    def _find_farthest_position(self, cell, vector):
        previous = None
        while True:
            previous = cell
            cell = self._get_next_tile(cell, vector)
            if cell is None or cell.value != 0: 
                break

        return previous, cell

    def _get_next_tile(self, tile, vector):
        newX = tile.x + vector.x
        newY = tile.y + vector.y
        if self._is_in_bounds(Vector(newX, newY)):
            return self.grid[newY][newX]

        return None

    def _is_in_bounds(self, vector):
        return 0 <= vector.x < self.size and 0 <= vector.y < self.size

    def _moves_available(self):
        for x in range(self.size):
            for y in range(self.size):
                tile = self.grid[y][x]
                if tile.value == 0 or self._is_tile_matches_available(tile):
                    return True

    def _is_tile_matches_available(self, tile):
        for direction in range(self.action_space):
            vector = self._get_vector(direction)
            neighbor = self._get_next_tile(tile, vector)
            if neighbor is not None and neighbor.value == tile.value:
                return True

        return False

    def step(self, direction): 
        start_score = self.score
        vector = self._get_vector(direction)
        traversals = self._build_traversals(vector)
        moved = False

        self._prepare_tiles()

        for x in traversals.x:
            for y in traversals.y:
                tile = self.grid[y][x]
                assert tile.x == x and tile.y == y
                if tile.value == 0:
                    continue
                farthest, next_cell = self._find_farthest_position(tile, vector)

                if next_cell is not None and next_cell.value == tile.value and not next_cell.merged:
                    new_value = tile.value * 2
                    next_cell.value = new_value
                    self.score += new_value
                    if new_value == 2048:
                        self.win_count += 1
                        self.won = True
                    moved = True

                    next_cell.merged = True
                    tile.value = 0
                else:
                    if farthest.x != tile.x or farthest.y != tile.y:
                        farthest.value = tile.value
                        tile.value = 0
                        moved = True

        if moved:
            self._add_rendom_tile()

            if not self._moves_available():
                self.over = True

        reward = self.score - start_score
        if self.env_type == 1:
            reward = -1 if not moved else min(reward, 1)

        is_terminated_state = self.over
        if self.env_type == 2:
            is_terminated_state = is_terminated_state or not moved

        return self.grid.get_values(), reward, is_terminated_state, moved
