# Configuración general y dependencias
import sys
import os
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

# Asegurar que scripts/ está en el path
PROJECT_ROOT = os.path.abspath(os.path.join(os.getcwd(), "..", "..", "..", ".."))
if PROJECT_ROOT not in sys.path:
    sys.path.insert(0, PROJECT_ROOT)

from scripts.pacman.level_generator import generate_level, generate_training_levels, Level
from scripts.pacman.game import PacmanEnv, STATE_DIM, NUM_ACTIONS, ACTION_NAMES
from scripts.pacman.renderer import render_frame, record_episode, create_animation, render_level_static

# 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}")
print(f"Dimensión del estado: {STATE_DIM}")
print(f"Número de acciones: {NUM_ACTIONS} ({', '.join(ACTION_NAMES)})")

# Generar un nivel de ejemplo y visualizarlo
level = generate_level(width=15, height=15, seed=42)
env = PacmanEnv(level=level, max_steps=500)
state = env.reset()

print("=== Nivel generado ===")
level.print()
print(f"\nPac-Man spawn: {level.pacman_spawn}")
print(f"Ghost spawns:  {level.ghost_spawns}")
print(f"Dots: {level.total_dots} | Power Pellets: {level.total_pellets}")

# Visualización gráfica del nivel
render_level_static(env, cell_size=16, figsize=(7, 7))

# Inspección del espacio de estados
state = env.reset()
print(f"Dimensión del estado: {state.shape}")
print(f"Ejemplo de estado inicial:\n{np.round(state, 2)}")
print(f"\nFeatures:")
feature_names = [
    "wall_up", "wall_right", "wall_down", "wall_left",
    "ghost_danger_up", "ghost_danger_right", "ghost_danger_down", "ghost_danger_left",
    "dot_up", "dot_right", "dot_down", "dot_left",
    "dot_dist_norm",
    "ghost_up", "ghost_right", "ghost_down", "ghost_left",
    "ghost_dist_norm",
    "power_active",
    "dots_remaining_ratio",
]
for i, (name, val) in enumerate(zip(feature_names, state)):
    print(f"  [{i:2d}] {name:25s} = {val:.2f}")

print(f"\nAcciones válidas desde posición actual: {env.get_valid_actions()}")
print(f"  (0=UP, 1=RIGHT, 2=DOWN, 3=LEFT)")

# Línea base: rendimiento de un agente aleatorio
random_rewards = []
random_scores = []
random_steps = []
random_wins = []

for ep in range(50):
    s = env.reset()
    done = False
    total_r = 0
    while not done:
        valid = env.get_valid_actions()
        a = random.choice(valid) if valid else 0
        s, r, done, info = env.step(a)
        total_r += r
    random_rewards.append(total_r)
    random_scores.append(info["score"])
    random_steps.append(info["steps"])
    random_wins.append(info["won"])

print("=== Línea base: agente aleatorio (50 episodios) ===")
print(f"Reward medio:  {np.mean(random_rewards):.1f} ± {np.std(random_rewards):.1f}")
print(f"Score medio:   {np.mean(random_scores):.1f} ± {np.std(random_scores):.1f}")
print(f"Pasos medios:  {np.mean(random_steps):.0f}")
print(f"Victorias:     {sum(random_wins)} / 50")

fig, axes = plt.subplots(1, 3, figsize=(14, 4))
axes[0].hist(random_rewards, bins=15, color="#7c5cff", alpha=0.7, edgecolor="white")
axes[0].set_title("Distribución de recompensas (aleatorio)")
axes[0].set_xlabel("Recompensa total")
axes[1].hist(random_scores, bins=15, color="#22c55e", alpha=0.7, edgecolor="white")
axes[1].set_title("Distribución de puntuación (aleatorio)")
axes[1].set_xlabel("Score")
axes[2].hist(random_steps, bins=15, color="#f59e0b", alpha=0.7, edgecolor="white")
axes[2].set_title("Distribución de pasos (aleatorio)")
axes[2].set_xlabel("Pasos hasta muerte/timeout")
plt.tight_layout()
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):
        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 más profunda)
# ──────────────────────────────────────────────────
class QNetwork(nn.Module):
    """
    MLP de 3 capas ocultas para aproximar Q(s, a) en Pac-Man.
    Entrada:  vector de estado (dim = 20)
    Salida:   Q-values para cada acción (dim = 4)
    """

    def __init__(self, state_dim: int, action_dim: int,
                 hidden_sizes: tuple = (256, 128, 64)):
        super().__init__()
        layers = []
        prev = state_dim
        for h in hidden_sizes:
            layers.append(nn.Linear(prev, h))
            layers.append(nn.ReLU())
            prev = h
        layers.append(nn.Linear(prev, action_dim))
        self.net = nn.Sequential(*layers)

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


# ──────────────────────────────────────────────────
# 3.3  Configuración de hiperparámetros
# ──────────────────────────────────────────────────
@dataclass
class DQNConfig:
    # --- Red ---
    hidden_sizes: tuple = (256, 128, 64)

    # --- Descuento y optimización ---
    gamma: float = 0.99
    lr: float = 5e-4

    # --- Replay Buffer ---
    buffer_capacity: int = 50_000
    min_buffer_size: int = 2_000
    batch_size: int = 128

    # --- Target Network ---
    use_soft_update: bool = True
    tau: float = 0.005
    hard_update_every: int = 200

    # --- Exploración (Epsilon-Greedy) ---
    epsilon_start: float = 1.0
    epsilon_end: float = 0.05
    epsilon_decay: float = 0.998  # Más lento que Snake


# ──────────────────────────────────────────────────
# 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

        self.q_net = QNetwork(state_dim, action_dim, config.hidden_sizes).to(device)
        self.target_net = QNetwork(state_dim, action_dim, config.hidden_sizes).to(device)
        self.target_net.load_state_dict(self.q_net.state_dict())
        self.target_net.eval()

        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

    def act(self, state: np.ndarray, explore: bool = True) -> int:
        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())

    def remember(self, s, a, r, ns, d):
        self.replay.push(s, a, r, ns, d)

    def train_step(self):
        if len(self.replay) < self.cfg.min_buffer_size:
            return None

        states, actions, rewards, next_states, dones = self.replay.sample(self.cfg.batch_size)

        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)

        q_values = self.q_net(states_t).gather(1, actions_t)

        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

        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)
        self.optimizer.step()

        self.learn_steps += 1
        if self.cfg.use_soft_update:
            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:
            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())

    def decay_epsilon(self):
        self.epsilon = max(self.cfg.epsilon_end, self.epsilon * self.cfg.epsilon_decay)

print("✅ Componentes DQN definidos")
print(f"   Red: {STATE_DIM} → {DQNConfig.hidden_sizes} → {NUM_ACTIONS}")
total_params = sum(p.numel() for p in QNetwork(STATE_DIM, NUM_ACTIONS).parameters())
print(f"   Parámetros totales: {total_params:,}")

# Función de evaluación (sin exploración, modo greedy puro)
def evaluate_agent(agent, levels, episodes_per_level=2, max_steps=500):
    """
    Ejecuta el agente greedy en todos los niveles de entrenamiento.
    """
    rewards, scores, wins = [], [], []

    for lv in levels:
        eval_env = PacmanEnv(level=lv, max_steps=max_steps)
        for _ in range(episodes_per_level):
            s = eval_env.reset()
            done = False
            total_r = 0.0
            while not done:
                a = agent.act(s, explore=False)
                s, r, done, info = eval_env.step(a)
                total_r += r
            rewards.append(total_r)
            scores.append(info["score"])
            wins.append(info["won"])

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


# Generar niveles de entrenamiento (variados)
training_levels = generate_training_levels(n=5, width=15, height=15, base_seed=42)
print(f"Niveles de entrenamiento: {len(training_levels)}")
for i, lv in enumerate(training_levels):
    print(f"  Nivel {i}: {lv.total_dots} dots, {lv.total_pellets} pellets")

# Crear entorno con rotación de niveles
env = PacmanEnv(levels=training_levels, max_steps=500)

# Bucle de entrenamiento principal
state_dim = STATE_DIM
action_dim = NUM_ACTIONS

agent = DQNAgent(state_dim, action_dim)

num_episodes = 4000
eval_every = 25

history = {
    'train_reward': [],
    'train_score': [],
    'train_wins': [],
    'mean_loss': [],
    'epsilon': [],
    'eval_reward': [],
    'eval_score': [],
    'eval_wins': [],
    '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['train_wins'].append(info['won'])
    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, eval_w = evaluate_agent(agent, training_levels)
        history['eval_reward'].append(eval_r)
        history['eval_score'].append(eval_s)
        history['eval_wins'].append(eval_w)
        history['eval_episode'].append(ep)

    if ep % 50 == 0:
        last50 = history['train_score'][-50:]
        wins50 = sum(history['train_wins'][-50:])
        print(
            f"Ep {ep:4d} | "
            f"score medio(últ.50)={np.mean(last50):.1f} | "
            f"wins(últ.50)={wins50} | "
            f"epsilon={agent.epsilon:.4f}"
        )

def moving_average(x, window=30):
    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])

fig, axes = plt.subplots(2, 3, figsize=(16, 10))

# (1) Reward
ax = axes[0, 0]
ax.plot(history['train_reward'], alpha=0.25, label='Train reward')
ax.plot(moving_average(history['train_reward'], 30), linewidth=2, label='MA(30)')
ax.plot(history['eval_episode'], history['eval_reward'], marker='o', markersize=4, label='Eval reward')
ax.set_title('Recompensa por episodio')
ax.set_xlabel('Episodio')
ax.legend(fontsize=8)
ax.grid(alpha=0.3)

# (2) Score
ax = axes[0, 1]
ax.plot(history['train_score'], alpha=0.25, label='Train score')
ax.plot(moving_average(history['train_score'], 30), linewidth=2, label='MA(30)')
ax.plot(history['eval_episode'], history['eval_score'], marker='o', markersize=4, label='Eval score')
ax.set_title('Score (comida) por episodio')
ax.set_xlabel('Episodio')
ax.legend(fontsize=8)
ax.grid(alpha=0.3)

# (3) Wins acumulados
ax = axes[0, 2]
cumulative_wins = np.cumsum(history['train_wins'])
ax.plot(cumulative_wins, color='green', linewidth=2)
ax.set_title('Victorias acumuladas (entrenamiento)')
ax.set_xlabel('Episodio')
ax.set_ylabel('Wins totales')
ax.grid(alpha=0.3)

# (4) Loss
ax = axes[1, 0]
loss_clean = np.nan_to_num(history['mean_loss'], nan=0.0)
ax.plot(loss_clean, alpha=0.3, label='Loss media')
ax.plot(moving_average(loss_clean, 30), linewidth=2, label='MA(30)')
ax.set_title('Pérdida de Bellman (Huber)')
ax.set_xlabel('Episodio')
ax.legend(fontsize=8)
ax.grid(alpha=0.3)

# (5) Epsilon
ax = axes[1, 1]
ax.plot(history['epsilon'], color='purple', linewidth=2)
ax.set_title('Decaimiento de epsilon')
ax.set_xlabel('Episodio')
ax.grid(alpha=0.3)

# (6) Eval wins
ax = axes[1, 2]
if history['eval_wins']:
    ax.bar(history['eval_episode'], history['eval_wins'], width=15, color='#22c55e', alpha=0.7)
ax.set_title('Victorias en evaluación')
ax.set_xlabel('Episodio')
ax.set_ylabel('Wins (de 10 eval eps)')
ax.grid(alpha=0.3)

plt.tight_layout()
plt.show()

# Evaluación final exhaustiva
final_eval_r, final_eval_s, final_eval_w = evaluate_agent(
    agent, training_levels, episodes_per_level=10, max_steps=500
)
total_eval_eps = len(training_levels) * 10
print(f"=== Evaluación final ({total_eval_eps} episodios, {len(training_levels)} niveles) ===")
print(f"Recompensa media: {final_eval_r:.1f}")
print(f"Score medio:      {final_eval_s:.1f}")
print(f"Victorias:        {final_eval_w} / {total_eval_eps}")

# Tests de sanidad
test_lv = generate_level(width=15, height=15, seed=999)
test_env = PacmanEnv(level=test_lv, max_steps=200)
s0 = test_env.reset()
assert s0.shape == (STATE_DIM,), f"Estado debe ser {STATE_DIM}D, es {s0.shape}"

test_agent = DQNAgent(state_dim=STATE_DIM, action_dim=NUM_ACTIONS)

# 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, NUM_ACTIONS), f"Red debe devolver Q para {NUM_ACTIONS} acciones"

# Test train_step con buffer lleno
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, NUM_ACTIONS - 1)
        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 loss finita"

print("✅ Tests de sanidad superados.")

# Grabar un episodio demo con el agente entrenado
demo_level = training_levels[0]
demo_env = PacmanEnv(level=demo_level, max_steps=500)

def agent_action(state):
    return agent.act(state, explore=False)

frames, demo_info = record_episode(demo_env, agent_action, max_steps=500, cell_size=12)
print(f"Demo: score={demo_info['score']} | pasos={demo_info['steps']} | won={demo_info['won']} | frames={len(frames)}")

# Animación inline del episodio
ani = create_animation(frames, interval=100)
HTML(ani.to_jshtml())

import json

def export_pacman_dqn_to_json(model: QNetwork, filepath: str,
                               state_dim: int = 20, action_dim: int = 4):
    """
    Exporta los pesos de la QNetwork a JSON compatible con la demo JS.
    Transpone weight de PyTorch (out, in) a JS (in, out).
    """
    layers_export = []
    layer_modules = [m for m in model.net if isinstance(m, nn.Linear)]

    for i, linear in enumerate(layer_modules):
        w = linear.weight.detach().cpu().numpy().T  # (in, out)
        b = linear.bias.detach().cpu().numpy()       # (out,)

        is_last = (i == len(layer_modules) - 1)
        layers_export.append({
            "type": "linear" if is_last else "relu",
            "weight": w.flatten().tolist(),
            "weight_rows": w.shape[0],
            "weight_cols": w.shape[1],
            "bias": b.tolist(),
            "bias_cols": len(b),
        })

    export_data = {
        "architecture": {
            "input_size": state_dim,
            "hidden_sizes": [l.out_features for l in layer_modules[:-1]],
            "output_size": action_dim,
        },
        "state_features": [
            "wall_up", "wall_right", "wall_down", "wall_left",
            "ghost_danger_up", "ghost_danger_right", "ghost_danger_down", "ghost_danger_left",
            "dot_up", "dot_right", "dot_down", "dot_left",
            "dot_dist_norm",
            "ghost_up", "ghost_right", "ghost_down", "ghost_left",
            "ghost_dist_norm",
            "power_active",
            "dots_remaining_ratio",
        ],
        "action_mapping": {
            "0": "up", "1": "right", "2": "down", "3": "left"
        },
        "layers": layers_export,
    }

    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    with open(filepath, "w") as f:
        json.dump(export_data, f)

    size_kb = os.path.getsize(filepath) / 1024
    print(f"✅ Modelo exportado a: {filepath}")
    print(f"   Arquitectura: {state_dim} → {[l.out_features for l in layer_modules[:-1]]} → {action_dim}")
    print(f"   Capas exportadas: {len(layers_export)}")
    print(f"   Tamaño: {size_kb:.1f} KB")

# Exportar
project_root = os.path.abspath(os.path.join(os.getcwd(), "..", "..", "..", ".."))
export_filepath = os.path.join(project_root, "static", "data", "pacman-drl", "pretrained_dqn.json")

export_pacman_dqn_to_json(agent.q_net, export_filepath)

# También exportar un nivel de ejemplo para la demo JS
level_filepath = os.path.join(project_root, "static", "data", "pacman-drl", "sample_level.json")
with open(level_filepath, "w") as f:
    json.dump(training_levels[0].to_dict(), f)
print(f"✅ Nivel de ejemplo exportado a: {level_filepath}")