---

import random
import numpy as np

WIDTH = 10
HEIGHT = 24
CELL = 1.0

# 7 фигур Tricky Towers / стандартных тетрамино.
SHAPES = {
    "S": [
        [(1, 0), (2, 0), (0, 1), (1, 1)]
    ],

    "Z": [
        [(0, 0), (1, 0), (1, 1), (2, 1)]
    ],

    "O": [
        [(0, 0), (1, 0), (0, 1), (1, 1)]
    ],

    "I": [
        [(0, 0), (1, 0), (2, 0), (3, 0)]
    ],

    "J": [
        [(0, 0), (0, 1), (1, 1), (2, 1)]
    ],

    "L": [
        [(2, 0), (0, 1), (1, 1), (2, 1)]
    ],

    "T": [
        [(0, 0), (1, 0), (2, 0), (1, 1)]
    ],
}


class TrickySimulator:

    def __init__(self, seed=None):
        self.rng = random.Random(seed)
        self.reset()

    def reset(self):
        self.grid = np.zeros((HEIGHT, WIDTH), dtype=np.int8)

        self.score = 0
        self.pieces = 0
        self.game_over = False

        self.current_name = None
        self.current_shape = None

        self.x = 0
        self.y = 0
        self.rotation = 0

        self.previous_height = 0
        self.max_height = 0

        self.spawn_piece()

        return self.get_state()

    # ---------------------------------------------------------
    # PIECES
    # ---------------------------------------------------------

    def spawn_piece(self):

        self.current_name = self.rng.choice(list(SHAPES.keys()))

        self.current_shape = [
            list(p) for p in SHAPES[self.current_name][0]
        ]

        self.rotation = 0

        min_x = min(x for x, y in self.current_shape)
        max_x = max(x for x, y in self.current_shape)

        piece_width = max_x - min_x + 1

        self.x = (WIDTH - piece_width) // 2 - min_x
        self.y = 0

        if self.collision(self.current_shape, self.x, self.y):
            self.game_over = True

    def rotated_shape(self):

        # Поворот вокруг условного центра.
        points = self.current_shape

        xs = [p[0] for p in points]
        ys = [p[1] for p in points]

        cx = (min(xs) + max(xs)) / 2
        cy = (min(ys) + max(ys)) / 2

        result = []

        for x, y in points:

            dx = x - cx
            dy = y - cy

            nx = round(cx - dy)
            ny = round(cy + dx)

            result.append((nx, ny))

        min_x = min(x for x, y in result)
        min_y = min(y for x, y in result)

        result = [
            (x - min_x, y - min_y)
            for x, y in result
        ]

        return result

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

    def collision(self, shape, px, py):

        for x, y in shape:

            gx = px + x
            gy = py + y

            if gx < 0 or gx >= WIDTH:
                return True

            if gy >= HEIGHT:
                return True

            if gy >= 0 and self.grid[gy, gx]:
                return True

        return False

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

    def move_left(self):

        if not self.collision(
            self.current_shape,
            self.x - 1,
            self.y
        ):
            self.x -= 1
            return True

        return False

    def move_right(self):

        if not self.collision(
            self.current_shape,
            self.x + 1,
            self.y
        ):
            self.x += 1
            return True

        return False

    def rotate(self):

        new_shape = self.rotated_shape()

        if not self.collision(
            new_shape,
            self.x,
            self.y
        ):
            self.current_shape = new_shape
            self.rotation = (self.rotation + 1) % 4
            return True

        # Простое wall-kick поведение
        for dx in [-1, 1, -2, 2]:

            if not self.collision(
                new_shape,
                self.x + dx,
                self.y
            ):
                self.x += dx
                self.current_shape = new_shape
                self.rotation = (self.rotation + 1) % 4
                return True

        return False

    def move_down(self):

        if not self.collision(
            self.current_shape,
            self.x,
            self.y + 1
        ):
            self.y += 1
            return True

        self.lock_piece()
        return False

    # ---------------------------------------------------------
    # LOCK
    # ---------------------------------------------------------

    def lock_piece(self):

        for x, y in self.current_shape:

            gx = self.x + x
            gy = self.y + y

            if 0 <= gx < WIDTH and 0 <= gy < HEIGHT:
                self.grid[gy, gx] = 1

        self.pieces += 1

        self.clear_lines()

        self.spawn_piece()

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

    def clear_lines(self):

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

        count = int(np.sum(full))

        if count > 0:

            self.score += count * 100

            remaining = self.grid[~full]

            new_grid = np.zeros_like(self.grid)

            new_grid[-len(remaining):] = remaining

            self.grid = new_grid

    # ---------------------------------------------------------
    # TOWER ANALYSIS
    # ---------------------------------------------------------

    def column_heights(self):

        heights = []

        for x in range(WIDTH):

            ys = np.where(self.grid[:, x] > 0)[0]

            if len(ys) == 0:
                heights.append(0)
            else:
                heights.append(HEIGHT - ys[0])

        return heights

    def tower_height(self):

        heights = self.column_heights()
        return max(heights)

    def holes(self):

        result = 0

        for x in range(WIDTH):

            ys = np.where(self.grid[:, x] > 0)[0]

            if len(ys) == 0:
                continue

            top = ys[0]

            result += int(
                np.sum(self.grid[top:, x] == 0)
            )

        return result

    def bumpiness(self):

        h = self.column_heights()

        return sum(
            abs(h[i] - h[i + 1])
            for i in range(WIDTH - 1)
        )

    def center_of_mass(self):

        ys, xs = np.where(self.grid > 0)

        if len(xs) == 0:
            return WIDTH / 2

        return float(np.mean(xs))

    def stability(self):

        h = self.column_heights()

        if max(h) == 0:
            return 1.0

        # Чем сильнее разница между соседними колонками,
        # тем менее устойчива башня.
        b = self.bumpiness()

        value = 1.0 - b / (WIDTH * HEIGHT)

        return max(0.0, min(1.0, value))

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

    def get_state(self):

        heights = self.column_heights()

        normalized_heights = [
            h / HEIGHT for h in heights
        ]

        shape_mask = np.zeros(16, dtype=np.float32)

        for x, y in self.current_shape:

            ix = y * 4 + x

            if 0 <= ix < 16:
                shape_mask[ix] = 1.0

        state = np.array(
            normalized_heights
            + list(shape_mask)
            + [
                self.x / WIDTH,
                self.y / HEIGHT,
                self.tower_height() / HEIGHT,
                self.holes() / 100.0,
                self.bumpiness() / 50.0,
                self.stability(),
                self.rotation / 4.0,
            ],
            dtype=np.float32
        )

        return state

    # ---------------------------------------------------------
    # REWARD
    # ---------------------------------------------------------

    def calculate_reward(
        self,
        old_height,
        old_holes,
        old_bumpiness,
        old_stability
    ):

        height = self.tower_height()
        holes = self.holes()
        bump = self.bumpiness()
        stability = self.stability()

        reward = 0.0

        # Выживание
        reward += 0.02

        # Чем меньше дыр
        reward += (old_holes - holes) * 0.8

        # Стараемся не создавать резких ступеней
        reward += (old_bumpiness - bump) * 0.08

        # Устойчивость
        reward += (stability - old_stability) * 2.0

        # Слишком быстро растущая башня опасна
        if height > old_height:
            reward -= (height - old_height) * 0.04

        # Очистка линии
        if self.score > 0:
            reward += 0.1

        # Проигрыш
        if self.game_over:
            reward -= 15.0

        return reward

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

    def step(self, action):

        if self.game_over:
            return self.get_state(), -15.0, True

        old_height = self.tower_height()
        old_holes = self.holes()
        old_bumpiness = self.bumpiness()
        old_stability = self.stability()

        # 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:

            # ↓ ускоряет падение
            for _ in range(2):
                if not self.move_down():
                    break

        else:

            # Естественная гравитация
            self.move_down()

        reward = self.calculate_reward(
            old_height,
            old_holes,
            old_bumpiness,
            old_stability
        )

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

    # ---------------------------------------------------------
    # DEBUG
    # ---------------------------------------------------------

    def render(self):

        board = self.grid.copy()

        if not self.game_over:

            for x, y in self.current_shape:

                gx = self.x + x
                gy = self.y + y

                if (
                    0 <= gx < WIDTH
                    and 0 <= gy < HEIGHT
                ):
                    board[gy, gx] = 2

        print("\n" + "=" * (WIDTH + 2))

        for y in range(HEIGHT):

            line = ""

            for x in range(WIDTH):

                if board[y, x] == 2:
                    line += "@"

                elif board[y, x] == 1:
                    line += "#"

                else:
                    line += "."

            print("|" + line + "|")

        print("=" * (WIDTH + 2))

        print(
            "piece:",
            self.current_name,
            "x:",
            self.x,
            "y:",
            self.y,
            "rotation:",
            self.rotation
        )

        print(
            "score:",
            self.score,
            "pieces:",
            self.pieces,
            "height:",
            self.tower_height(),
            "holes:",
            self.holes(),
            "stability:",
            round(self.stability(), 3)
        )


if __name__ == "__main__":

    env = TrickySimulator()

    print("Simulator started.")

    for i in range(30):

        action = random.randint(0, 4)

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

        if i % 5 == 0:
            env.render()

        if done:
            print("GAME OVER")
            break