---

import random
import numpy as np

WIDTH = 10
HEIGHT = 20

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

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

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

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

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

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

    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()

    # -------------------------
    # FIGURE
    # -------------------------

    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

    # -------------------------
    # COLLISION
    # -------------------------

    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

        region = self.board[
            y:y+h,
            x:x+w
        ]

        return np.any(
            (region & piece) != 0
        )

    # -------------------------
    # MOVEMENT
    # -------------------------

    def move_left(self):

        if not self.collision(
            self.piece,
            self.piece_x - 1,
            self.piece_y
        ):

            self.piece_x -= 1
            return True

        return False

    def move_right(self):

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

            self.piece_x += 1
            return True

        return False

    def rotate(self):

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

        # Небольшой wall kick
        for dx in [0, -1, 1, -2, 2]:

            if not self.collision(
                rotated,
                self.piece_x + dx,
                self.piece_y
            ):

                self.piece = rotated
                self.piece_x += dx

                return True

        return False

    # -------------------------
    # FALL
    # -------------------------

    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

    # -------------------------
    # PLACE
    # -------------------------

    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()

    # -------------------------
    # LINES
    # -------------------------

    def clear_lines(self):

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

        count = int(np.sum(full))

        if count == 0:
            return

        self.board = self.board[~full]

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

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

        self.lines += count

        self.score += count * 100

    # -------------------------
    # HEIGHT
    # -------------------------

    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)

    def get_tower_height(self):

        heights = self.get_column_heights()

        return int(np.max(heights))

    # -------------------------
    # HOLES
    # -------------------------

    def count_holes(self):

        holes = 0

        for x in range(WIDTH):

            found = False

            for y in range(HEIGHT):

                if self.board[y, x]:

                    found = True

                elif found:

                    holes += 1

        return holes

    # -------------------------
    # STABILITY
    # -------------------------

    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
        )

        stability = 1.0 / (
            1.0 + roughness
        )

        return float(stability)

    # -------------------------
    # STATE
    # -------------------------

    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

    # -------------------------
    # STEP
    # -------------------------

    def step(self, action):

        if self.game_over:

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

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

        # 0 = ничего
        # 1 = left
        # 2 = right
        # 3 = rotate
        # 4 = down

        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

        # Рост башни
        if new_height > old_height:
            reward += (
                new_height - old_height
            ) * 1.0

        # Новые дырки — плохо
        if new_holes > old_holes:
            reward -= (
                new_holes - old_holes
            ) * 2.0

        # Убрали линии
        # clear_lines уже добавил score,
        # здесь добавляем отдельную награду
        if self.lines > 0:
            reward += self.lines * 0.05

        if self.game_over:
            reward -= 100

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


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

def print_board(sim):

    board = sim.get_state()

    print("\n")

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

    for row in board:

        print(
            "|" +
            "".join(
                "#" if x else " "
                for x 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("TRICKY TOWERS SIMULATOR")
    print()
    print("a = ←")
    print("d = →")
    print("w = ↑ rotate")
    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)

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

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

        print_board(sim)

        if done:

            print()
            print("========== GAME OVER ==========")
            print(
                "Итоговый score:",
                sim.score
            )
            print(
                "Поставлено фигур:",
                sim.pieces_placed
            )
            print(
                "Линий:",
                sim.lines
            )
            break