---

import os

# Не позволяем каждому worker-процессу создавать дополнительные
# BLAS/OpenMP потоки. Один worker = один CPU-процесс.
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"

import random
import time
import pickle
import multiprocessing as mp
from collections import deque

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim

from tricky_sim import TrickySimulator


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

# i9-13900F:
NUM_WORKERS = 20
ENVS_PER_WORKER = 6

# Итого:
# 20 workers × 6 игр = 120 одновременно

TOTAL_STEPS = 5_000_000

# GPU
BATCH_SIZE = 2048
TRAIN_UPDATES = 4

LEARNING_RATE = 0.0003
GAMMA = 0.99

EPS_START = 1.0
EPS_END = 0.05
EPS_DECAY_STEPS = 1_000_000

WARMUP_STEPS = 20_000

TARGET_UPDATE = 10_000

REPLAY_SIZE = 500_000

SAVE_EVERY = 100_000

# Сколько действий выполняем между обменами с CPU
ROLLOUT_STEPS = 8

# ============================================================


DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")


# ============================================================
# DQN
# ============================================================

class DQN(nn.Module):

    def __init__(self, state_dim, action_dim):

        super().__init__()

        self.net = nn.Sequential(

            nn.Linear(state_dim, 512),
            nn.ReLU(),

            nn.Linear(512, 512),
            nn.ReLU(),

            nn.Linear(512, 256),
            nn.ReLU(),

            nn.Linear(256, 128),
            nn.ReLU(),

            nn.Linear(128, action_dim)
        )

    def forward(self, x):
        return self.net(x)


# ============================================================
# WORKER
# ============================================================

def worker_process(worker_id, command_queue, result_queue):

    # Каждый worker использует один CPU-поток.
    torch.set_num_threads(1)

    random.seed(10000 + worker_id)
    np.random.seed(10000 + worker_id)

    envs = []

    for i in range(ENVS_PER_WORKER):
        envs.append(TrickySimulator())

    # Получаем начальные состояния
    states = np.asarray(
        [env.get_state() for env in envs],
        dtype=np.float32
    )

    result_queue.put((
        "init",
        worker_id,
        states
    ))

    while True:

        command = command_queue.get()

        if command is None:
            break

        actions = command

        next_states = []
        rewards = []
        dones = []

        finished_games = 0
        finished_pieces = []

        for i, env in enumerate(envs):

            action = int(actions[i])

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

            rewards.append(float(reward))
            dones.append(bool(done))

            if done:

                finished_games += 1

                try:
                    finished_pieces.append(
                        int(getattr(env, "pieces", 0))
                    )
                except Exception:
                    finished_pieces.append(0)

                # Новая игра
                env.reset()

                next_state = env.get_state()

            next_states.append(next_state)

        states = np.asarray(
            next_states,
            dtype=np.float32
        )

        result_queue.put((
            "step",
            worker_id,
            states,
            np.asarray(rewards, dtype=np.float32),
            np.asarray(dones, dtype=np.bool_),
            finished_games,
            finished_pieces
        ))


# ============================================================
# REPLAY BUFFER
# ============================================================

class ReplayBuffer:

    def __init__(self, capacity, state_dim):

        self.capacity = capacity
        self.state_dim = state_dim

        self.states = np.empty(
            (capacity, state_dim),
            dtype=np.float32
        )

        self.next_states = np.empty(
            (capacity, state_dim),
            dtype=np.float32
        )

        self.actions = np.empty(
            capacity,
            dtype=np.int64
        )

        self.rewards = np.empty(
            capacity,
            dtype=np.float32
        )

        self.dones = np.empty(
            capacity,
            dtype=np.float32
        )

        self.position = 0
        self.size = 0

    def add_batch(
        self,
        states,
        actions,
        rewards,
        next_states,
        dones
    ):

        n = len(actions)

        for i in range(n):

            p = self.position

            self.states[p] = states[i]
            self.actions[p] = actions[i]
            self.rewards[p] = rewards[i]
            self.next_states[p] = next_states[i]
            self.dones[p] = dones[i]

            self.position += 1

            if self.position >= self.capacity:
                self.position = 0

            self.size = min(
                self.size + 1,
                self.capacity
            )

    def sample(self, batch_size):

        indexes = np.random.randint(
            0,
            self.size,
            size=batch_size
        )

        return (
            self.states[indexes],
            self.actions[indexes],
            self.rewards[indexes],
            self.next_states[indexes],
            self.dones[indexes]
        )


# ============================================================
# EPSILON
# ============================================================

def get_epsilon(step):

    if step >= EPS_DECAY_STEPS:
        return EPS_END

    ratio = step / EPS_DECAY_STEPS

    return EPS_START + ratio * (
        EPS_END - EPS_START
    )


# ============================================================
# TRAINING
# ============================================================

def train():

    print()
    print("======================================")
    print(" TRICKY TOWERS DQN V3")
    print("======================================")

    print("Device:", DEVICE)

    if torch.cuda.is_available():

        print("GPU:", torch.cuda.get_device_name(0))

        print(
            "CUDA:",
            torch.version.cuda
        )

        print(
            "VRAM:",
            round(
                torch.cuda.get_device_properties(0).total_memory
                / 1024**3,
                2
            ),
            "GB"
        )

    else:

        print("WARNING: CUDA OFF")
        print("Training will use CPU.")

    print()

    # --------------------------------------------------------
    # multiprocessing
    # --------------------------------------------------------

    ctx = mp.get_context("spawn")

    command_queues = []
    result_queue = ctx.Queue(
        maxsize=NUM_WORKERS * 3
    )

    workers = []

    print(
        f"Starting {NUM_WORKERS} CPU workers..."
    )

    for worker_id in range(NUM_WORKERS):

        q = ctx.Queue(
            maxsize=2
        )

        p = ctx.Process(
            target=worker_process,
            args=(
                worker_id,
                q,
                result_queue
            )
        )

        p.start()

        command_queues.append(q)
        workers.append(p)

    # --------------------------------------------------------
    # initial states
    # --------------------------------------------------------

    initial_states = [None] * NUM_WORKERS

    received = 0

    while received < NUM_WORKERS:

        msg = result_queue.get()

        if msg[0] == "init":

            _, worker_id, states = msg

            initial_states[worker_id] = states

            received += 1

    current_states = np.concatenate(
        initial_states,
        axis=0
    )

    state_dim = current_states.shape[1]

    action_dim = 5

    total_envs = NUM_WORKERS * ENVS_PER_WORKER

    print()
    print("======================================")
    print("CPU ENVIRONMENTS")
    print("======================================")
    print("Workers:", NUM_WORKERS)
    print("Games / worker:", ENVS_PER_WORKER)
    print("Parallel games:", total_envs)
    print("State size:", state_dim)
    print()

    # --------------------------------------------------------
    # GPU model
    # --------------------------------------------------------

    policy = DQN(
        state_dim,
        action_dim
    ).to(DEVICE)

    target = DQN(
        state_dim,
        action_dim
    ).to(DEVICE)

    target.load_state_dict(
        policy.state_dict()
    )

    target.eval()

    optimizer = optim.AdamW(
        policy.parameters(),
        lr=LEARNING_RATE,
        weight_decay=1e-5
    )

    # --------------------------------------------------------
    # Replay
    # --------------------------------------------------------

    replay = ReplayBuffer(
        REPLAY_SIZE,
        state_dim
    )

    # --------------------------------------------------------
    # GPU optimization
    # --------------------------------------------------------

    scaler = torch.amp.GradScaler(
        "cuda",
        enabled=(DEVICE.type == "cuda")
    )

    # --------------------------------------------------------
    # counters
    # --------------------------------------------------------

    global_step = 0
    games = 0
    best_pieces = 0

    episode_pieces = []

    start_time = time.time()

    last_print = time.time()
    last_print_step = 0

    # --------------------------------------------------------
    # MAIN LOOP
    # --------------------------------------------------------

    try:

        while global_step < TOTAL_STEPS:

            epsilon = get_epsilon(
                global_step
            )

            # ------------------------------------------------
            # GPU chooses actions
            # ------------------------------------------------

            with torch.no_grad():

                state_tensor = torch.from_numpy(
                    current_states
                )

                state_tensor = state_tensor.to(
                    DEVICE,
                    non_blocking=True
                )

                q_values = policy(
                    state_tensor
                )

                greedy_actions = (
                    q_values.argmax(
                        dim=1
                    )
                    .cpu()
                    .numpy()
                )

            # epsilon exploration

            actions = greedy_actions.copy()

            random_mask = (
                np.random.random(total_envs)
                < epsilon
            )

            actions[random_mask] = np.random.randint(
                0,
                action_dim,
                size=random_mask.sum()
            )

            # ------------------------------------------------
            # Send actions to CPU workers
            # ------------------------------------------------

            offset = 0

            for worker_id in range(NUM_WORKERS):

                worker_actions = actions[
                    offset:
                    offset + ENVS_PER_WORKER
                ]

                command_queues[
                    worker_id
                ].put(worker_actions)

                offset += ENVS_PER_WORKER

            # ------------------------------------------------
            # Receive CPU results
            # ------------------------------------------------

            new_states = [None] * NUM_WORKERS
            rewards = [None] * NUM_WORKERS
            dones = [None] * NUM_WORKERS

            batch_games = 0
            batch_pieces = []

            received = 0

            while received < NUM_WORKERS:

                msg = result_queue.get()

                if msg[0] != "step":
                    continue

                (
                    _,
                    worker_id,
                    states,
                    worker_rewards,
                    worker_dones,
                    worker_games,
                    worker_pieces
                ) = msg

                new_states[
                    worker_id
                ] = states

                rewards[
                    worker_id
                ] = worker_rewards

                dones[
                    worker_id
                ] = worker_dones

                batch_games += worker_games

                batch_pieces.extend(
                    worker_pieces
                )

                received += 1

            next_states = np.concatenate(
                new_states,
                axis=0
            )

            reward_array = np.concatenate(
                rewards,
                axis=0
            )

            done_array = np.concatenate(
                dones,
                axis=0
            )

            # ------------------------------------------------
            # Replay
            # ------------------------------------------------

            replay.add_batch(
                current_states,
                actions,
                reward_array,
                next_states,
                done_array
            )

            current_states = next_states

            global_step += total_envs

            games += batch_games

            if batch_pieces:

                episode_pieces.extend(
                    batch_pieces
                )

                best_pieces = max(
                    best_pieces,
                    max(batch_pieces)
                )

            # ------------------------------------------------
            # GPU training
            # ------------------------------------------------

            if (
                replay.size >= WARMUP_STEPS
                and global_step % 2 == 0
            ):

                for _ in range(
                    TRAIN_UPDATES
                ):

                    (
                        states_np,
                        actions_np,
                        rewards_np,
                        next_states_np,
                        dones_np
                    ) = replay.sample(
                        BATCH_SIZE
                    )

                    states = torch.from_numpy(
                        states_np
                    ).pin_memory().to(
                        DEVICE,
                        non_blocking=True
                    )

                    next_states = torch.from_numpy(
                        next_states_np
                    ).pin_memory().to(
                        DEVICE,
                        non_blocking=True
                    )

                    actions_t = torch.from_numpy(
                        actions_np
                    ).to(
                        DEVICE,
                        non_blocking=True
                    )

                    rewards_t = torch.from_numpy(
                        rewards_np
                    ).to(
                        DEVICE,
                        non_blocking=True
                    )

                    dones_t = torch.from_numpy(
                        dones_np
                    ).to(
                        DEVICE,
                        non_blocking=True
                    )

                    # ----------------------------------------
                    # Q(s,a)
                    # ----------------------------------------

                    with torch.amp.autocast(
                        device_type="cuda",
                        enabled=(DEVICE.type == "cuda")
                    ):

                        q = policy(
                            states
                        )

                        q_selected = q.gather(
                            1,
                            actions_t.unsqueeze(1)
                        ).squeeze(1)

                        # ------------------------------------
                        # Double DQN
                        # ------------------------------------

                        with torch.no_grad():

                            next_actions = policy(
                                next_states
                            ).argmax(
                                dim=1
                            )

                            next_q = target(
                                next_states
                            ).gather(
                                1,
                                next_actions.unsqueeze(1)
                            ).squeeze(1)

                            target_q = (
                                rewards_t
                                + GAMMA
                                * next_q
                                * (1.0 - dones_t)
                            )

                        loss = nn.functional.smooth_l1_loss(
                            q_selected,
                            target_q
                        )

                    optimizer.zero_grad(
                        set_to_none=True
                    )

                    if DEVICE.type == "cuda":

                        scaler.scale(
                            loss
                        ).backward()

                        scaler.unscale_(
                            optimizer
                        )

                        torch.nn.utils.clip_grad_norm_(
                            policy.parameters(),
                            10.0
                        )

                        scaler.step(
                            optimizer
                        )

                        scaler.update()

                    else:

                        loss.backward()

                        torch.nn.utils.clip_grad_norm_(
                            policy.parameters(),
                            10.0
                        )

                        optimizer.step()

            # ------------------------------------------------
            # Target network
            # ------------------------------------------------

            if (
                global_step > 0
                and global_step % TARGET_UPDATE
                < total_envs
            ):

                target.load_state_dict(
                    policy.state_dict()
                )

            # ------------------------------------------------
            # Save
            # ------------------------------------------------

            if (
                global_step > 0
                and global_step % SAVE_EVERY
                < total_envs
            ):

                checkpoint = {

                    "model": policy.state_dict(),

                    "target": target.state_dict(),

                    "optimizer": optimizer.state_dict(),

                    "step": global_step,

                    "epsilon": epsilon,

                    "state_dim": state_dim,

                    "action_dim": action_dim,

                    "best_pieces": best_pieces
                }

                torch.save(
                    checkpoint,
                    "tricky_dqn_v3.pt"
                )

                print()
                print(
                    "CHECKPOINT SAVED:",
                    global_step
                )

            # ------------------------------------------------
            # Statistics
            # ------------------------------------------------

            now = time.time()

            if now - last_print >= 5:

                elapsed = now - start_time

                delta_steps = (
                    global_step
                    - last_print_step
                )

                steps_sec = (
                    delta_steps
                    / max(
                        now - last_print,
                        0.001
                    )
                )

                if episode_pieces:

                    recent = episode_pieces[-100:]

                    avg_pieces = (
                        sum(recent)
                        / len(recent)
                    )

                else:

                    avg_pieces = 0

                gpu_mem = 0

                if DEVICE.type == "cuda":

                    gpu_mem = (
                        torch.cuda.memory_allocated()
                        / 1024**3
                    )

                print(
                    f"steps={global_step:,} | "
                    f"{steps_sec:,.0f} steps/s | "
                    f"games={games:,} | "
                    f"avg pieces={avg_pieces:.1f} | "
                    f"best={best_pieces} | "
                    f"eps={epsilon:.3f} | "
                    f"replay={replay.size:,} | "
                    f"GPU mem={gpu_mem:.2f}GB"
                )

                last_print = now
                last_print_step = global_step

    except KeyboardInterrupt:

        print()
        print("Stopping...")

    finally:

        # save before exit

        checkpoint = {

            "model": policy.state_dict(),

            "target": target.state_dict(),

            "optimizer": optimizer.state_dict(),

            "step": global_step,

            "epsilon": get_epsilon(global_step),

            "state_dim": state_dim,

            "action_dim": action_dim,

            "best_pieces": best_pieces
        }

        torch.save(
            checkpoint,
            "tricky_dqn_v3_last.pt"
        )

        print(
            "Saved:",
            "tricky_dqn_v3_last.pt"
        )

        # stop workers

        for q in command_queues:

            q.put(None)

        for p in workers:

            p.join(timeout=3)

            if p.is_alive():

                p.terminate()


# ============================================================

if __name__ == "__main__":

    mp.freeze_support()

    train()