---

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


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

EPISODES = 50000

BATCH_SIZE = 128
MEMORY_SIZE = 100000

GAMMA = 0.99

LEARNING_RATE = 0.001

EPSILON_START = 1.0
EPSILON_END = 0.05
EPSILON_DECAY = 30000

TARGET_UPDATE = 1000

MODEL_FILE = "tricky_model.pt"


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

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

print("Устройство:", device)


# =========================
# СОСТОЯНИЕ
# =========================

def get_features(sim):

    heights = sim.get_column_heights()

    max_height = max(
        sim.get_tower_height(),
        1
    )

    # Нормализуем высоты
    heights = heights / 20.0

    # Текущая фигура
    piece = sim.piece.flatten()

    # Дополним до 12 значений
    piece_data = np.zeros(12)

    n = min(
        len(piece),
        12
    )

    piece_data[:n] = piece[:n]

    # Положение фигуры
    piece_x = sim.piece_x / 10.0
    piece_y = sim.piece_y / 20.0

    # Дополнительные характеристики
    tower_height = sim.get_tower_height() / 20.0
    holes = sim.count_holes() / 100.0
    stability = sim.get_stability()

    state = np.concatenate([
        heights,
        piece_data,
        [piece_x],
        [piece_y],
        [tower_height],
        [holes],
        [stability]
    ])

    return state.astype(np.float32)


# =========================
# NETWORK
# =========================

class DQN(nn.Module):

    def __init__(self, input_size):

        super().__init__()

        self.net = nn.Sequential(

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

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

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

            # 5 действий:
            # nothing / left / right / rotate / down
            nn.Linear(64, 5)
        )

    def forward(self, x):

        return self.net(x)


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

class ReplayBuffer:

    def __init__(self, capacity):

        self.capacity = capacity
        self.buffer = []
        self.position = 0

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

        data = (
            state,
            action,
            reward,
            next_state,
            done
        )

        if len(self.buffer) < self.capacity:
            self.buffer.append(data)

        else:
            self.buffer[self.position] = data

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

    def sample(self, batch_size):

        batch = random.sample(
            self.buffer,
            batch_size
        )

        states, actions, rewards, next_states, dones = zip(
            *batch
        )

        return (
            np.array(states),
            np.array(actions),
            np.array(rewards),
            np.array(next_states),
            np.array(dones)
        )

    def __len__(self):

        return len(self.buffer)


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

def get_epsilon(step):

    if step >= EPSILON_DECAY:
        return EPSILON_END

    fraction = (
        step / EPSILON_DECAY
    )

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


# =========================
# INIT
# =========================

sim = TrickySimulator()

state = get_features(sim)

input_size = len(state)

print(
    "Размер состояния:",
    input_size
)

policy_net = DQN(
    input_size
).to(device)

target_net = DQN(
    input_size
).to(device)

target_net.load_state_dict(
    policy_net.state_dict()
)

target_net.eval()

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

memory = ReplayBuffer(
    MEMORY_SIZE
)


# =========================
# LOAD OLD MODEL
# =========================

if os.path.exists(MODEL_FILE):

    print(
        "Найдена старая модель."
    )

    try:

        policy_net.load_state_dict(
            torch.load(
                MODEL_FILE,
                map_location=device
            )
        )

        target_net.load_state_dict(
            policy_net.state_dict()
        )

        print(
            "Модель загружена."
        )

    except Exception as e:

        print(
            "Не удалось загрузить модель:",
            e
        )


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

def train_step():

    if len(memory) < BATCH_SIZE:
        return None

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

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

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

    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.unsqueeze(1)
    ).squeeze(1)

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

        next_q = target_net(
            next_states
        ).max(
            1
        )[0]

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

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

    optimizer.zero_grad()

    loss.backward()

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

    optimizer.step()

    return loss.item()


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

global_step = 0

best_score = 0

print()
print("НАЧАЛО ОБУЧЕНИЯ")
print()


for episode in range(1, EPISODES + 1):

    sim.reset()

    state = get_features(sim)

    episode_reward = 0

    loss_sum = 0
    loss_count = 0

    while not sim.game_over:

        epsilon = get_epsilon(
            global_step
        )

        # Exploration
        if random.random() < epsilon:

            action = random.randrange(5)

        # Exploitation
        else:

            with torch.no_grad():

                state_tensor = torch.tensor(
                    state,
                    dtype=torch.float32,
                    device=device
                ).unsqueeze(0)

                q_values = policy_net(
                    state_tensor
                )

                action = int(
                    torch.argmax(
                        q_values
                    ).item()
                )

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

        next_state = get_features(
            sim
        )

        memory.push(
            state,
            action,
            reward,
            next_state,
            done
        )

        state = next_state

        episode_reward += reward

        loss = train_step()

        if loss is not None:

            loss_sum += loss
            loss_count += 1

        global_step += 1

        # Target network
        if global_step % TARGET_UPDATE == 0:

            target_net.load_state_dict(
                policy_net.state_dict()
            )

    # =====================
    # SAVE BEST
    # =====================

    if sim.score > best_score:

        best_score = sim.score

        torch.save(
            policy_net.state_dict(),
            MODEL_FILE
        )

    # =====================
    # LOG
    # =====================

    if episode % 100 == 0:

        avg_loss = (
            loss_sum / loss_count
            if loss_count
            else 0
        )

        print(
            f"Игра {episode:6d} | "
            f"score {sim.score:5d} | "
            f"pieces {sim.pieces_placed:3d} | "
            f"height {sim.get_tower_height():2d} | "
            f"reward {episode_reward:8.2f} | "
            f"epsilon {epsilon:.3f} | "
            f"loss {avg_loss:.5f}"
        )


print()
print("ОБУЧЕНИЕ ЗАКОНЧЕНО")

torch.save(
    policy_net.state_dict(),
    MODEL_FILE
)

print(
    "Модель сохранена:",
    MODEL_FILE
)

print(
    "Лучший score:",
    best_score
)