---

import random
import numpy as np


WIDTH = 10
HEIGHT = 20


# Фигуры
SHAPES = [
    # O
    np.array([
        [1, 1],
        [1, 1]
    ], dtype=np.uint8),

    # I
    np.array([
        [1, 1, 1, 1]
    ], dtype=np.uint8),

    # T
    np.array([
        [0, 1, 0],
        [1, 1, 1]
    ], dtype=np.uint8),

    # L
    np.array([
        [1, 0],
        [1, 0],
        [1, 1]
    ], dtype=np.uint8),

    # J
    np.array([
        [0, 1],
        [0, 1],
        [1, 1]
    ], dtype=np.uint8),

    # S
    np.array([
        [0, 1, 1],
        [1, 1, 0]
    ], dtype=np.uint8),

    # Z
    np.array([
        [1, 1, 0],
        [0, 1, 1]
    ], dtype=np.uint8),
]


class TrickySimulator:

    def __init__(self):
        self.reset()

    def reset(self):

        self.board = np.zeros(
            (HEIGHT, WIDTH),
            dtype=np.uint8
        )

        self.score = 0
        self.pieces_placed = 0
        self.lines = 0
        self.game_over = False

        self.spawn_piece()

        return self.get_state()

    # ==========================================
    # СОЗДАНИЕ НОВОЙ ФИГУРЫ
    # ==========================================

    def spawn_piece(self):

        self.piece = random.choice(
            SHAPES
        ).copy()

        self.piece_x = (
            WIDTH // 2
            - self.piece.shape[1] // 2
        )

        self.piece_y = 0

        if self.collision(
            self.piece,
            self.piece_x,
            self.piece_y
        ):
            self.game_over = True

    # ==========================================
    # ПРОВЕРКА СТОЛКНОВЕНИЯ
    # ==========================================

    def collision(self, piece, x, y):

        h, w = piece.shape

        # Левая/правая граница
        if x < 0:
            return True

        if x + w > WIDTH:
            return True

        # Верх/низ
        if y < 0:
            return True

        if y + h > HEIGHT:
            return True

        # Область поля под фигурой
        area = self.board[
            y:y + h,
            x:x + w
        ]

        return np.any(
            (area == 1) &
            (piece == 1)
        )

    # ==========================================
    # ДВИЖЕНИЕ ВЛЕВО
    # ==========================================

    def move_left(self):

        new_x = self.piece_x - 1

        if not self.collision(
            self.piece,
            new_x,
            self.piece_y
        ):

            self.piece_x = new_x
            return True

        return False

    # ==========================================
    # ДВИЖЕНИЕ ВПРАВО
    # ==========================================

    def move_right(self):

        new_x = self.piece_x + 1

        if not self.collision(
            self.piece,
            new_x,
            self.piece_y
        ):

            self.piece_x = new_x
            return True

        return False

    # ==========================================
    # ПОВОРОТ
    # ==========================================

    def rotate(self):

        rotated = np.rot90(
            self.piece,
            -1
        )

        # Пробуем немного сдвинуть фигуру,
        # если возле стены она не помещается
        for dx in [0, -1, 1, -2, 2]:

            new_x = self.piece_x + dx

            if not self.collision(
                rotated,
                new_x,
                self.piece_y
            ):

                self.piece = rotated
                self.piece_x = new_x

                return True

        return False

    # ==========================================
    # МОЖЕТ ЛИ ФИГУРА ПАДАТЬ
    # ==========================================

    def can_move_down(self):

        return not self.collision(
            self.piece,
            self.piece_x,
            self.piece_y + 1
        )

    # ==========================================
    # ОПУСТИТЬ ФИГУРУ
    # ==========================================

    def move_down(self):

        if self.can_move_down():

            self.piece_y += 1
            return True

        return False

    # ==========================================
    # ЗАФИКСИРОВАТЬ ФИГУРУ
    # ==========================================

    def lock_piece(self):

        h, w = self.piece.shape

        self.board[
            self.piece_y:self.piece_y + h,
            self.piece_x:self.piece_x + w
        ] |= self.piece

        self.pieces_placed += 1

        # Базовая награда
        self.score += 10

        self.clear_lines()

        self.spawn_piece()

    # ==========================================
    # УДАЛЕНИЕ ЗАПОЛНЕННЫХ ЛИНИЙ
    # ==========================================

    def clear_lines(self):

        full_rows = np.all(
            self.board == 1,
            axis=1
        )

        count = int(
            np.sum(full_rows)
        )

        if count == 0:
            return

        self.board = self.board[
            ~full_rows
        ]

        empty_rows = np.zeros(
            (count, WIDTH),
            dtype=np.uint8
        )

        self.board = np.vstack([
            empty_rows,
            self.board
        ])

        self.lines += count

        self.score += count * 100

    # ==========================================
    # ВЫСОТА КАЖДОГО СТОЛБЦА
    # ==========================================

    def get_column_heights(self):

        heights = []

        for x in range(WIDTH):

            column = self.board[:, x]

            occupied = np.where(
                column == 1
            )[0]

            if len(occupied) == 0:
                heights.append(0)

            else:
                heights.append(
                    HEIGHT - occupied[0]
                )

        return np.array(
            heights,
            dtype=np.float32
        )

    # ==========================================
    # ВЫСОТА БАШНИ
    # ==========================================

    def get_tower_height(self):

        heights = self.get_column_heights()

        return int(
            np.max(heights)
        )

    # ==========================================
    # КОЛИЧЕСТВО ДЫР
    # ==========================================

    def count_holes(self):

        holes = 0

        for x in range(WIDTH):

            block_found = False

            for y in range(HEIGHT):

                if self.board[y, x]:

                    block_found = True

                elif block_found:

                    holes += 1

        return holes

    # ==========================================
    # СТАБИЛЬНОСТЬ БАШНИ
    # ==========================================

    def get_stability(self):

        heights = self.get_column_heights()

        if np.max(heights) == 0:
            return 1.0

        differences = np.abs(
            np.diff(heights)
        )

        roughness = np.sum(
            differences
        )

        return float(
            1.0 / (1.0 + roughness)
        )

    # ==========================================
    # СОСТОЯНИЕ ДЛЯ AI
    # ==========================================

    def get_state(self):

        state = self.board.copy()

        if self.game_over:
            return state

        h, w = self.piece.shape

        y1 = self.piece_y
        y2 = min(
            self.piece_y + h,
            HEIGHT
        )

        x1 = self.piece_x
        x2 = min(
            self.piece_x + w,
            WIDTH
        )

        if (
            y1 >= 0
            and x1 >= 0
            and y1 < HEIGHT
            and x1 < WIDTH
        ):

            state[
                y1:y2,
                x1:x2
            ] |= self.piece[
                :y2 - y1,
                :x2 - x1
            ]

        return state

    # ==========================================
    # ОДИН ШАГ ИГРЫ
    # ==========================================

    def step(self, action):

        if self.game_over:

            return (
                self.get_state(),
                -100.0,
                True
            )

        old_height = self.get_tower_height()
        old_holes = self.count_holes()
        old_pieces = self.pieces_placed
        old_lines = self.lines

        # 0 = ничего
        # 1 = влево
        # 2 = вправо
        # 3 = поворот
        # 4 = вниз

        if action == 1:

            self.move_left()

        elif action == 2:

            self.move_right()

        elif action == 3:

            self.rotate()

        elif action == 4:

            self.move_down()

        # Обычная гравитация
        if self.can_move_down():

            self.piece_y += 1

        else:

            self.lock_piece()

        new_height = self.get_tower_height()
        new_holes = self.count_holes()

        reward = 0.0

        # Фигура успешно поставлена
        if self.pieces_placed > old_pieces:

            reward += 10.0

        # Рост башни
        if new_height > old_height:

            reward += (
                new_height - old_height
            ) * 1.0

        # Появились новые дыры
        if new_holes > old_holes:

            reward -= (
                new_holes - old_holes
            ) * 2.0

        # Очистили линии
        if self.lines > old_lines:

            reward += (
                self.lines - old_lines
            ) * 100.0

        # Проигрыш
        if self.game_over:

            reward -= 100.0

        return (
            self.get_state(),
            reward,
            self.game_over
        )


# ==========================================
# ПЕЧАТЬ ПОЛЯ
# ==========================================

def print_board(sim):

    board = sim.get_state()

    print()

    print(
        "+" + "-" * WIDTH + "+"
    )

    for row in board:

        print(
            "|" +
            "".join(
                "#" if cell else " "
                for cell in row
            ) +
            "|"
        )

    print(
        "+" + "-" * WIDTH + "+"
    )

    print(
        f"score={sim.score} | "
        f"pieces={sim.pieces_placed} | "
        f"lines={sim.lines} | "
        f"height={sim.get_tower_height()} | "
        f"holes={sim.count_holes()} | "
        f"stability={sim.get_stability():.3f}"
    )


# ==========================================
# РУЧНОЙ ТЕСТ
# ==========================================

if __name__ == "__main__":

    sim = TrickySimulator()

    print("================================")
    print("      TRICKY TOWERS SIM")
    print("================================")
    print()
    print("a = ←")
    print("d = →")
    print("w = ↑ поворот")
    print("s = ↓")
    print("Enter = ничего")
    print("q = выход")

    print_board(sim)

    while not sim.game_over:

        command = input("> ").lower().strip()

        if command == "q":
            break

        actions = {
            "": 0,
            "a": 1,
            "d": 2,
            "w": 3,
            "s": 4
        }

        action = actions.get(
            command,
            0
        )

        _, reward, done = sim.step(
            action
        )

        print(
            f"reward = {reward:.2f}"
        )

        print_board(sim)

        if done:

            print()
            print("==============================")
            print("          GAME OVER")
            print("==============================")
            print(
                "score:",
                sim.score
            )
            print(
                "pieces:",
                sim.pieces_placed
            )
            print(
                "lines:",
                sim.lines
            )
            break