---

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

from tricky_sim import TrickySimulator


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

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

TOTAL_STEPS = 2_000_000

BATCH_SIZE = 256
REPLAY_SIZE = 200_000

GAMMA = 0.99

LEARNING_RATE = 0.0005

EPSILON_START = 1.0
EPSILON_END = 0.05
EPSILON_DECAY = 500_000

LEARN_EVERY = 4
TARGET_UPDATE = 5_000

WARMUP = 10_000

MODEL_FILE = "tricky_dqn.pt"


# =========================================================
# DEVICE
# =========================================================

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

print("=" * 60)
print("TRICKY TOWERS DQN")
print("=" * 60)

print("CPU threads:", os.cpu_count())
print("Parallel games:", NUM_ENVS)
print("Device:", device)

if device.type == "cuda":
    print("GPU:", torch.cuda.get_device_name(0))

print("=" * 60)


# =========================================================
# ENVIRONMENTS
# =========================================================

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

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

STATE_SIZE = len(states[0])

ACTION_SIZE = 5

print("State size:", STATE_SIZE)
print("Actions:", ACTION_SIZE)


# =========================================================
# NEURAL NETWORK
# =========================================================

class DQN(nn.Module):

    def __init__(self):

        super().__init__()

        self.net = nn.Sequential(

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

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

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

            nn.Linear(128, ACTION_SIZE)
        )

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


policy_net = DQN().to(device)

target_net = DQN().to(device)

target_net.load_state_dict(
    policy_net.state_dict()
)

target_net.eval()


optimizer = optim.Adam(
    policy_net.parameters(),
    lr=LEARNING_RATE
)

loss_function = nn.SmoothL1Loss()


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

class ReplayBuffer:

    def __init__(self, capacity):

        self.capacity = capacity

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

        self.position = 0

    def add(
        self,
        state,
        action,
        reward,
        next_state,
        done
    ):

        if len(self.states) < self.capacity:

            self.states.append(state)
            self.actions.append(action)
            self.rewards.append(reward)
            self.next_states.append(next_state)
            self.dones.append(done)

        else:

            self.states[self.position] = state
            self.actions[self.position] = action
            self.rewards[self.position] = reward
            self.next_states[self.position] = next_state
            self.dones[self.position] = done

        self.position = (
            self.position + 1
        ) % self.capacity

    def sample(self, batch_size):

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

        states = np.asarray(
            [self.states[i] for i in indexes],
            dtype=np.float32
        )

        actions = np.asarray(
            [self.actions[i] for i in indexes],
            dtype=np.int64
        )

        rewards = np.asarray(
            [self.rewards[i] for i in indexes],
            dtype=np.float32
        )

        next_states = np.asarray(
            [self.next_states[i] for i in indexes],
            dtype=np.float32
        )

        dones = np.asarray(
            [self.dones[i] for i in indexes],
            dtype=np.float32
        )

        return (
            states,
            actions,
            rewards,
            next_states,
            dones
        )

    def __len__(self):
        return len(self.states)


replay = ReplayBuffer(REPLAY_SIZE)


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

def get_epsilon(step):

    fraction = min(
        1.0,
        step / EPSILON_DECAY
    )

    return (
        EPSILON_START
        + fraction
        * (
            EPSILON_END
            - EPSILON_START
        )
    )


# =========================================================
# ACTION
# =========================================================

@torch.no_grad()
def choose_actions(states, epsilon):

    actions = []

    # Все состояния одним батчем
    state_tensor = torch.tensor(
        np.asarray(states),
        dtype=torch.float32,
        device=device
    )

    q_values = policy_net(
        state_tensor
    )

    for i in range(len(states)):

        if random.random() < epsilon:

            action = random.randrange(
                ACTION_SIZE
            )

        else:

            action = int(
                torch.argmax(
                    q_values[i]
                ).item()
            )

        actions.append(action)

    return actions


# =========================================================
# TRAIN STEP
# =========================================================

def train_step():

    if len(replay) < max(
        BATCH_SIZE,
        WARMUP
    ):
        return None

    (
        states,
        actions,
        rewards,
        next_states,
        dones
    ) = replay.sample(BATCH_SIZE)

    states = torch.tensor(
        states,
        dtype=torch.float32,
        device=device
    )

    actions = torch.tensor(
        actions,
        dtype=torch.long,
        device=device
    ).unsqueeze(1)

    rewards = torch.tensor(
        rewards,
        dtype=torch.float32,
        device=device
    )

    next_states = torch.tensor(
        next_states,
        dtype=torch.float32,
        device=device
    )

    dones = torch.tensor(
        dones,
        dtype=torch.float32,
        device=device
    )

    # Q(s,a)
    current_q = policy_net(
        states
    ).gather(
        1,
        actions
    ).squeeze(1)

    # max Q(s',a')
    with torch.no_grad():

        next_q = target_net(
            next_states
        ).max(
            1
        ).values

        target_q = (
            rewards
            + GAMMA
            * next_q
            * (1.0 - dones)
        )

    loss = loss_function(
        current_q,
        target_q
    )

    optimizer.zero_grad()

    loss.backward()

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

    optimizer.step()

    return float(loss.item())


# =========================================================
# STATISTICS
# =========================================================

total_games = 0
total_pieces = 0

best_pieces = 0

recent_games = []

losses = []


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

print()
print("Training started...")
print()

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

    epsilon = get_epsilon(step)

    # -----------------------------------------------------
    # Choose action for all 32 games
    # -----------------------------------------------------

    actions = choose_actions(
        states,
        epsilon
    )

    # -----------------------------------------------------
    # Step environments
    # -----------------------------------------------------

    for i, env in enumerate(envs):

        state = states[i]

        action = actions[i]

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

        replay.add(
            state,
            action,
            reward,
            next_state,
            done
        )

        states[i] = next_state

        if done:

            pieces = env.pieces

            total_games += 1
            total_pieces += pieces

            best_pieces = max(
                best_pieces,
                pieces
            )

            recent_games.append(
                pieces
            )

            if len(recent_games) > 100:
                recent_games.pop(0)

            states[i] = env.reset()

    # -----------------------------------------------------
    # Learn
    # -----------------------------------------------------

    if (
        step >= WARMUP
        and step % LEARN_EVERY == 0
    ):

        loss = train_step()

        if loss is not None:
            losses.append(loss)

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

    if step % TARGET_UPDATE == 0:

        target_net.load_state_dict(
            policy_net.state_dict()
        )

    # -----------------------------------------------------
    # Console
    # -----------------------------------------------------

    if step % 1_000 == 0:

        if recent_games:

            avg_recent = (
                sum(recent_games)
                / len(recent_games)
            )

        else:

            avg_recent = 0

        if losses:

            avg_loss = (
                sum(losses[-100:])
                / min(
                    100,
                    len(losses)
                )
            )

        else:

            avg_loss = 0

        print(
            f"step={step:,} | "
            f"games={total_games:,} | "
            f"avg100={avg_recent:.1f} | "
            f"best={best_pieces} | "
            f"epsilon={epsilon:.3f} | "
            f"loss={avg_loss:.4f} | "
            f"buffer={len(replay):,}",
            flush=True
        )

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

    if step % 50_000 == 0:

        torch.save(
            {
                "model": policy_net.state_dict(),
                "target": target_net.state_dict(),
                "optimizer": optimizer.state_dict(),
                "step": step,
                "best_pieces": best_pieces
            },
            MODEL_FILE
        )

        print(
            f"Model saved: {MODEL_FILE}",
            flush=True
        )


# =========================================================
# FINAL SAVE
# =========================================================

torch.save(
    {
        "model": policy_net.state_dict(),
        "target": target_net.state_dict(),
        "optimizer": optimizer.state_dict(),
        "step": TOTAL_STEPS,
        "best_pieces": best_pieces
    },
    MODEL_FILE
)

print()
print("=" * 60)
print("TRAINING FINISHED")
print("Best:", best_pieces)
print("Model:", MODEL_FILE)
print("=" * 60)_