# Configuración general y dependencias
import random
import math
from dataclasses import dataclass
from collections import deque

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
from matplotlib import animation
from IPython.display import HTML

# Semillas para reproducibilidad
SEED = 42
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Dispositivo disponible: {device}")

# Entorno Snake desde cero
class SnakeEnv:
    def __init__(self, size=10, max_steps_without_food=120):
        self.size = size
        self.max_steps_without_food = max_steps_without_food
        self.reset()

    def reset(self):
        # Dirección: 0=arriba, 1=derecha, 2=abajo, 3=izquierda
        self.direction = 1
        mid = self.size // 2
        self.snake = [(mid, mid), (mid, mid - 1), (mid, mid - 2)]  # cabeza primero
        self._spawn_food()
        self.done = False
        self.score = 0
        self.steps = 0
        self.steps_since_food = 0
        self.prev_distance = self._distance_to_food(self.snake[0])
        return self._get_state()

    def _spawn_food(self):
        empty = [(r, c) for r in range(self.size) for c in range(self.size) if (r, c) not in self.snake]
        self.food = random.choice(empty)

    @staticmethod
    def _move_point(point, direction):
        r, c = point
        if direction == 0:
            return (r - 1, c)
        if direction == 1:
            return (r, c + 1)
        if direction == 2:
            return (r + 1, c)
        return (r, c - 1)

    def _distance_to_food(self, point):
        return abs(point[0] - self.food[0]) + abs(point[1] - self.food[1])

    def _is_collision(self, point):
        r, c = point
        if r < 0 or r >= self.size or c < 0 or c >= self.size:
            return True
        if point in self.snake:
            return True
        return False

    def _rotate(self, action):
        # action: 0=recto, 1=izquierda, 2=derecha
        if action == 1:
            self.direction = (self.direction - 1) % 4
        elif action == 2:
            self.direction = (self.direction + 1) % 4

    def _danger(self, relative_turn):
        # relative_turn: -1 izquierda, 0 frente, +1 derecha
        test_dir = (self.direction + relative_turn) % 4
        next_head = self._move_point(self.snake[0], test_dir)
        return 1.0 if self._is_collision(next_head) else 0.0

    def _get_state(self):
        head_r, head_c = self.snake[0]
        food_r, food_c = self.food

        state = [
            self._danger(0),   # peligro al frente
            self._danger(-1),  # peligro a la izquierda
            self._danger(+1),  # peligro a la derecha
            1.0 if self.direction == 0 else 0.0,
            1.0 if self.direction == 1 else 0.0,
            1.0 if self.direction == 2 else 0.0,
            1.0 if self.direction == 3 else 0.0,
            1.0 if food_r < head_r else 0.0,   # comida arriba
            1.0 if food_r > head_r else 0.0,   # comida abajo
            1.0 if food_c < head_c else 0.0,   # comida izquierda
            1.0 if food_c > head_c else 0.0,   # comida derecha
            self._distance_to_food(self.snake[0]) / (2 * self.size),
        ]
        return np.array(state, dtype=np.float32)

    def step(self, action):
        if self.done:
            raise RuntimeError("El episodio terminó. Llama a reset().")

        self.steps += 1
        self.steps_since_food += 1

        # 1) Actualizar dirección según acción
        self._rotate(action)

        # 2) Proponer nuevo movimiento
        new_head = self._move_point(self.snake[0], self.direction)

        # 3) Recompensas/castigos (reward shaping)
        reward = 0.0

        # Colisión => castigo fuerte y fin
        if self._is_collision(new_head):
            self.done = True
            reward = -10.0
            return self._get_state(), reward, self.done, {"score": self.score}

        # Avance normal
        self.snake.insert(0, new_head)

        # Comer comida => premio alto
        if new_head == self.food:
            self.score += 1
            reward = 10.0
            self.steps_since_food = 0
            self._spawn_food()
        else:
            # Si no come, elimina cola (movimiento normal)
            self.snake.pop()

            # Pequeña penalización por paso para evitar vagar sin objetivo
            reward -= 0.03

            # Recompensa densa por acercarse a la comida
            new_dist = self._distance_to_food(new_head)
            if new_dist < self.prev_distance:
                reward += 0.2
            else:
                reward -= 0.2
            self.prev_distance = new_dist

        # Timeout sin comer => episodio finalizado (castigo moderado)
        if self.steps_since_food >= self.max_steps_without_food:
            self.done = True
            reward -= 2.0

        return self._get_state(), reward, self.done, {"score": self.score}

# Inspección básica del entorno
env = SnakeEnv(size=10)
state = env.reset()

print("Dimensión del estado:", state.shape)
print("Ejemplo de estado inicial:", state)
print("Acciones disponibles: 0=recto, 1=izquierda, 2=derecha")

# Simulación aleatoria corta para observar recompensas
episode_rewards = []
for ep in range(5):
    s = env.reset()
    done = False
    total_r = 0
    while not done:
        a = random.randint(0, 2)
        s, r, done, info = env.step(a)
        total_r += r
    episode_rewards.append(total_r)

print("Recompensas totales (política aleatoria, 5 episodios):", np.round(episode_rewards, 2))

# Visualización rápida de recompensas aleatorias
plt.figure(figsize=(7, 4))
plt.plot(episode_rewards, marker='o')
plt.title('Recompensa total por episodio (agente aleatorio)')
plt.xlabel('Episodio')
plt.ylabel('Recompensa total')
plt.grid(True, alpha=0.3)
plt.show()

# ──────────────────────────────────────────────────
# 3.1  Replay Buffer
# ──────────────────────────────────────────────────
class ReplayBuffer:
    """Almacena transiciones (s, a, r, s', done) en una cola circular."""

    def __init__(self, capacity: int = 50_000):
        self.buffer = deque(maxlen=capacity)

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

    def sample(self, batch_size: int):
        """Muestreo uniforme aleatorio de un mini-batch."""
        batch = random.sample(self.buffer, batch_size)
        states, actions, rewards, next_states, dones = map(np.array, zip(*batch))
        return states, actions, rewards, next_states, dones

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


# ──────────────────────────────────────────────────
# 3.2  Red Neuronal Q (MLP)
# ──────────────────────────────────────────────────
class QNetwork(nn.Module):
    """
    MLP que aproxima Q(s, a) para todas las acciones simultáneamente.
    Entrada:  vector de estado (dim = state_dim)
    Salida:   Q-values para cada acción (dim = action_dim)
    """

    def __init__(self, state_dim: int, action_dim: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, 128),   # Capa de entrada → capa oculta 1
            nn.ReLU(),
            nn.Linear(128, 128),         # Capa oculta 1 → capa oculta 2
            nn.ReLU(),
            nn.Linear(128, action_dim),  # Capa oculta 2 → salida (sin activación)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


# ──────────────────────────────────────────────────
# 3.3  Configuración de hiperparámetros
# ──────────────────────────────────────────────────
@dataclass
class DQNConfig:
    # --- Descuento y optimización ---
    gamma: float = 0.99            # Factor de descuento γ
    lr: float = 1e-3               # Learning rate del optimizador Adam

    # --- Replay Buffer ---
    buffer_capacity: int = 30_000  # Capacidad máxima del buffer
    min_buffer_size: int = 1_000   # Mínimo de transiciones antes de entrenar
    batch_size: int = 128          # Tamaño del mini-batch

    # --- Target Network ---
    use_soft_update: bool = True   # True = Polyak averaging, False = hard update
    tau: float = 0.005             # Coeficiente de soft update (τ)
    hard_update_every: int = 200   # Frecuencia de hard update (solo si use_soft_update=False)

    # --- Exploración (Epsilon-Greedy) ---
    epsilon_start: float = 1.0     # ε inicial (exploración pura)
    epsilon_end: float = 0.05      # ε mínimo (siempre algo de exploración)
    epsilon_decay: float = 0.995   # Factor de decaimiento multiplicativo


# ──────────────────────────────────────────────────
# 3.4  Agente DQN
# ──────────────────────────────────────────────────
class DQNAgent:
    def __init__(self, state_dim: int, action_dim: int, config: DQNConfig = DQNConfig()):
        self.cfg = config
        self.action_dim = action_dim

        # Red online (la que aprende) y red objetivo (proporciona TD targets estables)
        self.q_net = QNetwork(state_dim, action_dim).to(device)
        self.target_net = QNetwork(state_dim, action_dim).to(device)
        self.target_net.load_state_dict(self.q_net.state_dict())  # Iniciar iguales
        self.target_net.eval()  # La target network nunca se entrena directamente

        self.optimizer = optim.Adam(self.q_net.parameters(), lr=self.cfg.lr)
        self.criterion = nn.SmoothL1Loss()  # Huber Loss

        self.replay = ReplayBuffer(self.cfg.buffer_capacity)
        self.epsilon = self.cfg.epsilon_start
        self.learn_steps = 0

    # ---- Selección de acción: epsilon-greedy ----
    def act(self, state: np.ndarray, explore: bool = True) -> int:
        """
        Con probabilidad ε → acción aleatoria (exploración).
        Con probabilidad 1-ε → argmax Q (explotación).
        En modo evaluación (explore=False), siempre explota.
        """
        if explore and random.random() < self.epsilon:
            return random.randint(0, self.action_dim - 1)

        state_t = torch.tensor(state, dtype=torch.float32, device=device).unsqueeze(0)
        with torch.no_grad():
            q_values = self.q_net(state_t)
        return int(torch.argmax(q_values, dim=1).item())

    # ---- Almacenar transición en el buffer ----
    def remember(self, s, a, r, ns, d):
        self.replay.push(s, a, r, ns, d)

    # ---- Paso de entrenamiento ----
    def train_step(self):
        """
        1. Muestrear mini-batch del replay buffer.
        2. Calcular Q(s,a) con la red online.
        3. Calcular TD target con la red objetivo.
        4. Backpropagation con Huber Loss.
        5. Actualizar target network (soft o hard).
        """
        if len(self.replay) < self.cfg.min_buffer_size:
            return None  # No hay suficientes datos aún

        # (1) Muestreo aleatorio del buffer
        states, actions, rewards, next_states, dones = self.replay.sample(self.cfg.batch_size)

        # Convertir a tensores de PyTorch
        states_t = torch.tensor(states, dtype=torch.float32, device=device)
        actions_t = torch.tensor(actions, dtype=torch.long, device=device).unsqueeze(1)
        rewards_t = torch.tensor(rewards, dtype=torch.float32, device=device).unsqueeze(1)
        next_states_t = torch.tensor(next_states, dtype=torch.float32, device=device)
        dones_t = torch.tensor(dones, dtype=torch.float32, device=device).unsqueeze(1)

        # (2) Q(s, a) actual — solo para la acción tomada
        #     gather(1, actions_t) selecciona el Q-value de la acción que se ejecutó
        q_values = self.q_net(states_t).gather(1, actions_t)

        # (3) TD target: y = r + γ(1 - done) * max_a' Q_target(s', a')
        with torch.no_grad():
            max_next_q = self.target_net(next_states_t).max(dim=1, keepdim=True)[0]
            td_target = rewards_t + self.cfg.gamma * (1 - dones_t) * max_next_q

        # (4) Calcular pérdida y backpropagation
        loss = self.criterion(q_values, td_target)
        self.optimizer.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(self.q_net.parameters(), max_norm=10.0)  # Gradient clipping
        self.optimizer.step()

        # (5) Actualizar target network
        self.learn_steps += 1
        if self.cfg.use_soft_update:
            # Soft update (Polyak averaging): θ̄ ← τθ + (1-τ)θ̄
            for param, target_param in zip(self.q_net.parameters(), self.target_net.parameters()):
                target_param.data.copy_(
                    self.cfg.tau * param.data + (1.0 - self.cfg.tau) * target_param.data
                )
        else:
            # Hard update: copiar pesos cada N pasos
            if self.learn_steps % self.cfg.hard_update_every == 0:
                self.target_net.load_state_dict(self.q_net.state_dict())

        return float(loss.item())

    # ---- Decaimiento de epsilon ----
    def decay_epsilon(self):
        """Decaimiento exponencial: ε ← max(ε_min, ε * δ)"""
        self.epsilon = max(self.cfg.epsilon_end, self.epsilon * self.cfg.epsilon_decay)

# Función de evaluación (sin exploración, modo greedy puro)
def evaluate_agent(agent, env_size=10, episodes=10):
    """
    Ejecuta el agente en modo puramente greedy (ε=0) durante varios episodios.
    Retorna la recompensa media y el score medio.
    Es el equivalente a "evaluar en validación" en aprendizaje supervisado.
    """
    eval_env = SnakeEnv(size=env_size)
    rewards, scores = [], []

    for _ in range(episodes):
        s = eval_env.reset()
        done = False
        total_r = 0.0
        while not done:
            a = agent.act(s, explore=False)  # explore=False → siempre argmax Q
            s, r, done, info = eval_env.step(a)
            total_r += r
        rewards.append(total_r)
        scores.append(info['score'])

    return float(np.mean(rewards)), float(np.mean(scores))

# Bucle de entrenamiento principal
env = SnakeEnv(size=10, max_steps_without_food=120)
state_dim = env.reset().shape[0]
action_dim = 3

agent = DQNAgent(state_dim, action_dim)

num_episodes = 400
eval_every = 25

history = {
    'train_reward': [],
    'train_score': [],
    'mean_loss': [],
    'epsilon': [],
    'eval_reward': [],
    'eval_score': [],
    'eval_episode': [],
}

for ep in range(1, num_episodes + 1):
    s = env.reset()
    done = False
    total_reward = 0.0
    losses = []

    while not done:
        a = agent.act(s, explore=True)
        ns, r, done, info = env.step(a)

        agent.remember(s, a, r, ns, done)
        loss = agent.train_step()
        if loss is not None:
            losses.append(loss)

        s = ns
        total_reward += r

    agent.decay_epsilon()

    history['train_reward'].append(total_reward)
    history['train_score'].append(info['score'])
    history['mean_loss'].append(np.mean(losses) if losses else np.nan)
    history['epsilon'].append(agent.epsilon)

    if ep % eval_every == 0:
        eval_r, eval_s = evaluate_agent(agent, env_size=10, episodes=15)
        history['eval_reward'].append(eval_r)
        history['eval_score'].append(eval_s)
        history['eval_episode'].append(ep)

    if ep % 25 == 0:
        last50 = history['train_score'][-50:]
        print(
            f"Ep {ep:3d} | "
            f"score medio(últ.50)={np.mean(last50):.2f} | "
            f"epsilon={agent.epsilon:.3f}"
        )

# Visualización del schedule de decaimiento de epsilon
eps = 1.0
eps_min = 0.05
eps_decay = 0.995
epsilons = []
for _ in range(600):
    epsilons.append(eps)
    eps = max(eps_min, eps * eps_decay)

plt.figure(figsize=(8, 4))
plt.plot(epsilons, color='purple', linewidth=2)
plt.axhline(y=eps_min, color='red', linestyle='--', alpha=0.5, label=f'ε_min = {eps_min}')
plt.fill_between(range(len(epsilons)), epsilons, alpha=0.15, color='purple')
plt.title('Schedule de Epsilon: de Exploración a Explotación')
plt.xlabel('Episodio')
plt.ylabel('ε (probabilidad de exploración)')
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

print(f"ε en episodio 50:  {epsilons[50]:.3f}")
print(f"ε en episodio 200: {epsilons[200]:.3f}")
print(f"ε en episodio 350: {epsilons[350]:.3f}")
print(f"ε en episodio 500: {epsilons[500]:.3f}")

# Utilidad para media móvil

def moving_average(x, window=20):
    x = np.asarray(x, dtype=np.float32)
    if len(x) < window:
        return x
    c = np.convolve(x, np.ones(window) / window, mode='valid')
    pad = np.full(window - 1, np.nan)
    return np.concatenate([pad, c])

plt.figure(figsize=(14, 10))

# (1) Reward train vs eval
ax1 = plt.subplot(2, 2, 1)
ax1.plot(history['train_reward'], alpha=0.35, label='Train reward (episodio)')
ax1.plot(moving_average(history['train_reward'], 20), linewidth=2, label='Train reward MA(20)')
ax1.plot(history['eval_episode'], history['eval_reward'], marker='o', label='Eval reward (greedy)')
ax1.set_title('Recompensa: entrenamiento vs evaluación')
ax1.set_xlabel('Episodio')
ax1.set_ylabel('Reward')
ax1.grid(alpha=0.3)
ax1.legend()

# (2) Score train vs eval
ax2 = plt.subplot(2, 2, 2)
ax2.plot(history['train_score'], alpha=0.35, label='Train score (episodio)')
ax2.plot(moving_average(history['train_score'], 20), linewidth=2, label='Train score MA(20)')
ax2.plot(history['eval_episode'], history['eval_score'], marker='o', label='Eval score (greedy)')
ax2.set_title('Score (comida): entrenamiento vs evaluación')
ax2.set_xlabel('Episodio')
ax2.set_ylabel('Comidas')
ax2.grid(alpha=0.3)
ax2.legend()

# (3) Loss vs episodio
ax3 = plt.subplot(2, 2, 3)
ax3.plot(history['mean_loss'], label='Loss media por episodio')
ax3.plot(moving_average(np.nan_to_num(history['mean_loss'], nan=0.0), 20), linewidth=2, label='Loss MA(20)')
ax3.set_title('Pérdida de Bellman (Huber)')
ax3.set_xlabel('Episodio')
ax3.set_ylabel('Loss')
ax3.grid(alpha=0.3)
ax3.legend()

# (4) Epsilon
ax4 = plt.subplot(2, 2, 4)
ax4.plot(history['epsilon'], color='purple')
ax4.set_title('Decaimiento de epsilon (exploración)')
ax4.set_xlabel('Episodio')
ax4.set_ylabel('Epsilon')
ax4.grid(alpha=0.3)

plt.tight_layout()
plt.show()

# Métricas resumen de evaluación final
final_eval_reward, final_eval_score = evaluate_agent(agent, env_size=10, episodes=50)
print(f"Recompensa media en evaluación (50 episodios): {final_eval_reward:.2f}")
print(f"Score medio en evaluación (50 episodios):      {final_eval_score:.2f}")

# Tests simples (educativos)

test_env = SnakeEnv(size=8)
s0 = test_env.reset()
assert s0.shape == (12,), "La dimensión del estado debe ser 12"

test_agent = DQNAgent(state_dim=12, action_dim=3)

# Test forward pass
tensor_state = torch.tensor(s0, dtype=torch.float32, device=device).unsqueeze(0)
qvals = test_agent.q_net(tensor_state)
assert qvals.shape == (1, 3), "La red debe devolver Q para 3 acciones"

# Llenar buffer con transiciones aleatorias para habilitar train_step
while len(test_agent.replay) < test_agent.cfg.min_buffer_size:
    s = test_env.reset()
    done = False
    while not done and len(test_agent.replay) < test_agent.cfg.min_buffer_size:
        a = random.randint(0, 2)
        ns, r, done, _ = test_env.step(a)
        test_agent.remember(s, a, r, ns, done)
        s = ns

loss_val = test_agent.train_step()
assert loss_val is not None and np.isfinite(loss_val), "train_step debe devolver una loss válida"

print("✅ Tests de sanidad superados.")

# Render de un episodio como frames RGB

def render_grid(env):
    grid = np.zeros((env.size, env.size, 3), dtype=np.uint8)

    # Fondo
    grid[:] = np.array([20, 20, 20], dtype=np.uint8)

    # Comida (rojo)
    fr, fc = env.food
    grid[fr, fc] = np.array([220, 50, 50], dtype=np.uint8)

    # Cuerpo serpiente (verde)
    for i, (r, c) in enumerate(env.snake):
        if i == 0:
            grid[r, c] = np.array([50, 220, 120], dtype=np.uint8)  # cabeza
        else:
            grid[r, c] = np.array([40, 140, 80], dtype=np.uint8)

    return grid


def rollout_frames(agent, env_size=10, max_steps=250):
    demo_env = SnakeEnv(size=env_size)
    s = demo_env.reset()
    done = False

    frames = [render_grid(demo_env)]
    total_reward = 0.0

    steps = 0
    while not done and steps < max_steps:
        a = agent.act(s, explore=False)
        s, r, done, info = demo_env.step(a)
        total_reward += r
        frames.append(render_grid(demo_env))
        steps += 1

    return frames, info['score'], total_reward

frames, demo_score, demo_reward = rollout_frames(agent, env_size=10, max_steps=300)
print(f"Demo terminada | score={demo_score} | reward_total={demo_reward:.2f} | frames={len(frames)}")

# Animación inline en notebook
fig, ax = plt.subplots(figsize=(5, 5))
ax.axis('off')
img = ax.imshow(frames[0], interpolation='nearest')

def animate(i):
    img.set_data(frames[i])
    return [img]

ani = animation.FuncAnimation(fig, animate, frames=len(frames), interval=120, blit=True)
plt.close(fig)
HTML(ani.to_jshtml())