---

import os
import random
import numpy as np

from tricky_sim import TrickySimulator


# ---------------------------------------------------------
# НАСТРОЙКИ
# ---------------------------------------------------------

NUM_ENVS = max(
    4,
    min(32, (os.cpu_count() or 8))
)

TOTAL_STEPS = 2_000_000

PRINT_EVERY = 10_000


# ---------------------------------------------------------
# СОЗДАЕМ ИГРЫ
# ---------------------------------------------------------

envs = [
    TrickySimulator(seed=i)
    for i in range(NUM_ENVS)
]

states = [
    env.reset()
    for env in envs
]


# ---------------------------------------------------------
# ПРОСТАЯ ТЕСТОВАЯ ПОЛИТИКА
# ---------------------------------------------------------
#
# Пока НЕ нейросеть.
#
# Она просто случайно выбирает действие.
# Наша задача сейчас проверить скорость симулятора.
# ---------------------------------------------------------

total_pieces = 0
total_games = 0

episode_rewards = [
    0.0
    for _ in range(NUM_ENVS)
]

best_pieces = 0


# ---------------------------------------------------------
# TRAINING LOOP
# ---------------------------------------------------------

for step in range(1, TOTAL_STEPS + 1):

    for i, env in enumerate(envs):

        # Случайное действие
        action = random.randrange(5)

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

        states[i] = next_state

        episode_rewards[i] += reward

        if done:

            total_pieces += env.pieces
            total_games += 1

            best_pieces = max(
                best_pieces,
                env.pieces
            )

            states[i] = env.reset()

            episode_rewards[i] = 0.0

    if step % PRINT_EVERY == 0:

        avg_pieces = (
            total_pieces / total_games
            if total_games
            else 0
        )

        print(
            f"steps={step:,} | "
            f"envs={NUM_ENVS} | "
            f"games={total_games:,} | "
            f"avg pieces={avg_pieces:.2f} | "
            f"best={best_pieces}"
        )


print()
print("Finished.")