Traceback (most recent call last):
  File "C:\Users\hp\Downloads\snk.py", line 753, in <module>
    main()
  File "C:\Users\hp\Downloads\snk.py", line 750, in main
    trainer.train()
  File "C:\Users\hp\Downloads\snk.py", line 544, in train
    self.agent.replay()
  File "C:\Users\hp\Downloads\snk.py", line 438, in replay
    current_q = self.policy_net(states).gather(1, actions)
                ^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\hp\AppData\Local\Programs\Python\Python312\Lib\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\hp\AppData\Local\Programs\Python\Python312\Lib\site-packages\torch\nn\modules\module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\hp\Downloads\snk.py", line 368, in forward
    x = F.relu(self.fc1(x))
               ^^^^^^^^^^^
  File "C:\Users\hp\AppData\Local\Programs\Python\Python312\Lib\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\hp\AppData\Local\Programs\Python\Python312\Lib\site-packages\torch\nn\modules\module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\hp\AppData\Local\Programs\Python\Python312\Lib\site-packages\torch\nn\modules\linear.py", line 116, in forward
    return F.linear(input, self.weight, self.bias)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: mat1 and mat2 shapes cannot be multiplied (64x11 and 19x256)

```
import pygame
import sys
import random
import heapq
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import matplotlib.pyplot as plt
from collections import deque, namedtuple
import numpy as np
import os
import math

# Initialize Pygame
pygame.init()

# ---------------------- Constants and Configurations ----------------------

# Game Constants
CELL_SIZE = 20
GRID_WIDTH = 30
GRID_HEIGHT = 20
WINDOW_WIDTH = CELL_SIZE * GRID_WIDTH + 300  # Extra width for UI
WINDOW_HEIGHT = CELL_SIZE * GRID_HEIGHT
FPS = 10

# Boost Constants
BOOST_DURATION = 50  # Number of frames boost is active
BOOST_COOLDOWN = 100  # Number of frames before boost can be activated again
BOOST_SPEED_MULTIPLIER = 2  # Multiplier for FPS when boost is active

# Q-Learning Constants
ACTIONS = [
    ("UP", False),
    ("DOWN", False),
    ("LEFT", False),
    ("RIGHT", False),
    ("UP", True),
    ("DOWN", True),
    ("LEFT", True),
    ("RIGHT", True)
]

ACTION_INDEX = {action: idx for idx, action in enumerate(ACTIONS)}
NUM_ACTIONS = len(ACTIONS)

# Reward Constants
REWARD_FOOD = 100
REWARD_STEP = 1
REWARD_BOOST_EFFECTIVE = 5
REWARD_CLOSE_DISTANCE = 2  # Bonus per unit distance reduced
PENALTY_FAR_DISTANCE = -1  # Penalty per unit distance increased
PENALTY_COLLISION = -100
PENALTY_BOOST_UNUSED = -2

# Neural Network Parameters
STATE_SIZE = 19  # Expanded state representation
HIDDEN_SIZE = 256
LEARNING_RATE = 0.0005
BATCH_SIZE = 64
GAMMA = 0.99
EPSILON_START = 1.0
EPSILON_END = 0.05
EPSILON_DECAY = 0.995
MEMORY_SIZE = 50000
TARGET_UPDATE = 1000  # Update target network every 1000 steps

# Prioritized Experience Replay Constants
ALPHA = 0.6
BETA_START = 0.4
BETA_FRAMES = 100000

# Visualization Constants
PLOT_WINDOW = 100  # Number of episodes to display in the plot

# Colors
WHITE = (255, 255, 255)
GREY = (40, 40, 40)
BLACK = (0, 0, 0)
GREEN = (0, 255, 0)
DARK_GREEN = (0, 155, 0)
RED = (255, 0, 0)
BLUE = (0, 0, 255)
YELLOW = (255, 255, 0)
ORANGE = (255, 165, 0)

# Set up display
screen = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT))
pygame.display.set_caption("Advanced Snake AI with Double DQN and Prioritized Experience Replay")
clock = pygame.time.Clock()

# Initialize Font
pygame.font.init()
FONT_SIZE = 20
FONT = pygame.font.SysFont('Calibri', FONT_SIZE, bold=True)

# ---------------------- Helper Functions ----------------------

def render_text(text, color, position, background=None):
    """
    Renders text on the screen at the given position.
    """
    text_surface = FONT.render(text, True, color, background)
    screen.blit(text_surface, position)

def manhattan_distance(a, b):
    """Calculate Manhattan distance between two points a and b."""
    return abs(a[0] - b[0]) + abs(a[1] - b[1])

def generate_hamiltonian_cycle():
    """Generates a simple Hamiltonian cycle for the grid."""
    cycle = []
    for y in range(GRID_HEIGHT):
        row = list(range(GRID_WIDTH))
        if y % 2 == 1:
            row.reverse()
        for x in row:
            cycle.append((x, y))
    return cycle

def a_star_search(start, goal, snake_positions, obstacles):
    """
    Performs A* search from start to goal avoiding snake_positions and obstacles.
    Returns the path as a list of positions from start to goal, inclusive.
    If no path found, returns an empty list.
    """
    def heuristic(a, b):
        return manhattan_distance(a, b)  # Manhattan distance

    open_set = []
    heapq.heappush(open_set, (heuristic(start, goal), 0, start, [start]))
    closed_set = set()

    while open_set:
        estimated_total, cost, current, path = heapq.heappop(open_set)

        if current == goal:
            return path

        if current in closed_set:
            continue
        closed_set.add(current)

        for dx, dy in [(-1,0), (1,0), (0,-1), (0,1)]:
            neighbor = (current[0] + dx, current[1] + dy)
            if (0 <= neighbor[0] < GRID_WIDTH and
                0 <= neighbor[1] < GRID_HEIGHT and
                neighbor not in snake_positions and
                neighbor not in obstacles and
                neighbor not in closed_set):
                new_cost = cost + 1
                estimated = new_cost + heuristic(neighbor, goal)
                heapq.heappush(open_set, (estimated, new_cost, neighbor, path + [neighbor]))
    return []

# ---------------------- Classes ----------------------

class Snake:
    def __init__(self):
        self.reset()

    def reset(self):
        self.positions = deque([(GRID_WIDTH // 2, GRID_HEIGHT // 2)])
        self.direction = (0, -1)  # Initially moving upwards
        self.grow = False
        self.hamiltonian = generate_hamiltonian_cycle()
        self.cycle_index = self.hamiltonian.index(self.head()) if self.head() in self.hamiltonian else 0
        # Boost attributes
        self.boost_active = False
        self.boost_duration = 0
        self.boost_cooldown = 0
        # Score
        self.score = 0

    def head(self):
        return self.positions[0]

    def move(self, direction):
        if direction is not None:
            self.direction = direction
        new_x = self.head()[0] + self.direction[0]
        new_y = self.head()[1] + self.direction[1]
        new_position = (new_x, new_y)

        # Check wall collision or self-collision
        if (0 <= new_x < GRID_WIDTH and 0 <= new_y < GRID_HEIGHT and new_position not in self.positions):
            self.positions.appendleft(new_position)
            if not self.grow:
                self.positions.pop()
            else:
                self.grow = False
            return True  # Successful move
        else:
            return False  # Collision occurred

    def grow_snake(self):
        self.grow = True
        self.score += 1

    def activate_boost(self):
        if self.boost_cooldown == 0 and not self.boost_active:
            self.boost_active = True
            self.boost_duration = BOOST_DURATION
            self.boost_cooldown = BOOST_COOLDOWN
            return True
        return False

    def update_boost(self):
        if self.boost_active:
            self.boost_duration -= 1
            if self.boost_duration <= 0:
                self.boost_active = False
        if self.boost_cooldown > 0:
            self.boost_cooldown -= 1

class Food:
    def __init__(self, snake_positions, obstacles):
        self.obstacles = obstacles
        self.position = self.random_position(snake_positions)

    def random_position(self, snake_positions):
        while True:
            pos = (random.randint(0, GRID_WIDTH - 1), random.randint(0, GRID_HEIGHT - 1))
            if pos not in snake_positions and pos not in self.obstacles:
                return pos

    def respawn(self, snake_positions):
        self.position = self.random_position(snake_positions)

class Obstacle:
    def __init__(self, num_obstacles=10):
        self.positions = set()
        self.num_obstacles = num_obstacles
        self.generate_obstacles()

    def generate_obstacles(self):
        while len(self.positions) < self.num_obstacles:
            pos = (random.randint(0, GRID_WIDTH - 1), random.randint(0, GRID_HEIGHT - 1))
            self.positions.add(pos)

    def reset(self):
        self.positions.clear()
        self.generate_obstacles()

# Replay Memory with Prioritized Experience Replay
Transition = namedtuple('Transition',
                        ('state', 'action', 'reward', 'next_state', 'done'))

class SumTree:
    """
    SumTree data structure for Prioritized Experience Replay
    """
    write = 0

    def __init__(self, capacity):
        self.capacity = capacity
        self.tree = np.zeros(2 * capacity - 1)
        self.data = np.empty(capacity, dtype=object)
        self.n_entries = 0

    def _propagate(self, idx, change):
        parent = (idx - 1) // 2
        self.tree[parent] += change

        if parent != 0:
            self._propagate(parent, change)

    def update(self, idx, p):
        change = p - self.tree[idx]
        self.tree[idx] = p
        self._propagate(idx, change)

    def add(self, p, data):
        idx = self.write + self.capacity - 1

        self.data[self.write] = data
        self.update(idx, p)

        self.write += 1
        if self.write >= self.capacity:
            self.write = 0

        if self.n_entries < self.capacity:
            self.n_entries += 1

    def _retrieve(self, idx, s):
        left = 2 * idx + 1
        right = left + 1

        if left >= len(self.tree):
            return idx

        if s <= self.tree[left]:
            return self._retrieve(left, s)
        else:
            return self._retrieve(right, s - self.tree[left])

    def get(self, s):
        idx = self._retrieve(0, s)
        dataIdx = idx - self.capacity + 1

        return (idx, self.tree[idx], self.data[dataIdx])

    @property
    def total_p(self):
        return self.tree[0]

class PrioritizedReplayMemory:
    def __init__(self, capacity, alpha=ALPHA):
        self.tree = SumTree(capacity)
        self.capacity = capacity
        self.alpha = alpha
        self.epsilon = 1e-6

    def push(self, *args):
        max_p = np.max(self.tree.tree[-self.tree.capacity:])
        if max_p == 0:
            max_p = 1.0
        self.tree.add(max_p, Transition(*args))

    def sample(self, batch_size, beta):
        batch = []
        idxs = []
        segment = self.tree.total_p / batch_size
        priorities = []

        for i in range(batch_size):
            a = segment * i
            b = segment * (i + 1)
            s = random.uniform(a, b)
            idx, p, data = self.tree.get(s)
            batch.append(data)
            idxs.append(idx)
            priorities.append(p)

        sampling_probabilities = priorities / self.tree.total_p
        is_weight = np.power(self.tree.n_entries * sampling_probabilities, -beta)
        is_weight /= is_weight.max()

        return batch, idxs, is_weight

    def update_priorities(self, idxs, priorities):
        for idx, p in zip(idxs, priorities):
            p = (abs(p) + self.epsilon) ** self.alpha
            self.tree.update(idx, p)

    def __len__(self):
        return self.tree.n_entries

class QNetwork(nn.Module):
    def __init__(self, input_size=19, hidden_size=256, output_size=NUM_ACTIONS):
        super(QNetwork, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc3 = nn.Linear(hidden_size, output_size)
        self._init_weights()

    def _init_weights(self):
        nn.init.kaiming_uniform_(self.fc1.weight, nonlinearity='relu')
        nn.init.constant_(self.fc1.bias, 0)
        nn.init.kaiming_uniform_(self.fc2.weight, nonlinearity='relu')
        nn.init.constant_(self.fc2.bias, 0)
        nn.init.xavier_uniform_(self.fc3.weight)
        nn.init.constant_(self.fc3.bias, 0)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

class QLearningAgent:
    def __init__(self, actions, state_size=19, hidden_size=256, action_size=NUM_ACTIONS,
                 learning_rate=LEARNING_RATE, gamma=GAMMA, epsilon=EPSILON_START,
                 epsilon_decay=EPSILON_DECAY, epsilon_min=EPSILON_END,
                 memory_size=MEMORY_SIZE, batch_size=BATCH_SIZE, target_update=TARGET_UPDATE):
        self.actions = actions
        self.state_size = state_size
        self.action_size = action_size
        self.gamma = gamma
        self.epsilon = epsilon
        self.epsilon_decay = epsilon_decay
        self.epsilon_min = epsilon_min
        self.batch_size = batch_size
        self.target_update = target_update
        self.learn_step_counter = 0

        self.memory = PrioritizedReplayMemory(memory_size)

        self.policy_net = QNetwork(input_size=state_size, hidden_size=hidden_size, output_size=action_size)
        self.target_net = QNetwork(input_size=state_size, hidden_size=hidden_size, output_size=action_size)
        self.target_net.load_state_dict(self.policy_net.state_dict())
        self.target_net.eval()

        self.optimizer = optim.Adam(self.policy_net.parameters(), lr=learning_rate)
        self.loss_fn = nn.MSELoss()

        self.beta = BETA_START
        self.beta_increment = (1.0 - BETA_START) / BETA_FRAMES

        # For Double DQN
        self.update_target_network()

    def get_action(self, state, preferred_action=None):
        if random.random() < self.epsilon:
            return random.choice(self.actions)
        else:
            with torch.no_grad():
                q_values = self.policy_net(state)
                action_idx = torch.argmax(q_values).item()
                return self.actions[action_idx]

    def remember(self, state, action, reward, next_state, done):
        # Calculate priority
        state_np = state.numpy()
        next_state_np = next_state.numpy()
        transition = Transition(state, action, reward, next_state, done)
        self.memory.push(*transition)

    def replay(self):
        if len(self.memory) < self.batch_size:
            return

        batch, idxs, is_weights = self.memory.sample(self.batch_size, self.beta)
        self.beta = min(1.0, self.beta + self.beta_increment)

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

        states = torch.stack(states)
        actions = torch.tensor(actions, dtype=torch.long).unsqueeze(1)
        rewards = torch.tensor(rewards, dtype=torch.float32).unsqueeze(1)
        next_states = torch.stack(next_states)
        dones = torch.tensor(dones, dtype=torch.float32).unsqueeze(1)
        is_weights = torch.tensor(is_weights, dtype=torch.float32).unsqueeze(1)

        # Current Q values
        current_q = self.policy_net(states).gather(1, actions)

        # For Double DQN
        with torch.no_grad():
            next_actions = self.policy_net(next_states).argmax(1).unsqueeze(1)
            next_q = self.target_net(next_states).gather(1, next_actions)
            target_q = rewards + (1 - dones) * self.gamma * next_q

        # Compute TD error
        td_error = target_q - current_q
        loss = (is_weights * td_error.pow(2)).mean()

        # Optimize the model
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        # Update priorities
        priorities = td_error.abs().detach().numpy() + 1e-6
        self.memory.update_priorities(idxs, priorities.flatten())

    def update_epsilon(self):
        if self.epsilon > self.epsilon_min:
            self.epsilon *= self.epsilon_decay

    def update_target_network(self):
        self.target_net.load_state_dict(self.policy_net.state_dict())

    def save_model(self, path):
        torch.save(self.policy_net.state_dict(), path)

    def load_model(self, path):
        if os.path.exists(path):
            self.policy_net.load_state_dict(torch.load(path))
            self.update_target_network()
            print(f"Model loaded from {path}")

class Trainer:
    def __init__(self, agent, num_episodes=10000, max_steps=1000, render=False):
        self.agent = agent
        self.num_episodes = num_episodes
        self.max_steps = max_steps
        self.render = render
        self.plot_scores = []
        self.plot_ma_scores = []
        self.ma_window = PLOT_WINDOW
        self.total_steps = 0
        self.model_save_path = "snake_dqn.pth"

    def train(self):
        scores = deque(maxlen=100)
        ma_scores = deque(maxlen=self.ma_window)
        high_score = 0

        for episode in range(1, self.num_episodes + 1):
            snake = Snake()
            obstacles = Obstacle()
            food = Food(snake.positions, obstacles.positions)
            done = False
            step = 0

            # Initialize state
            state = get_state(snake, food, obstacles)
            
            while not done and step < self.max_steps:
                step += 1
                self.total_steps += 1

                # Agent selects action
                action = self.agent.get_action(state)
                direction, use_boost = parse_action(action)

                # Activate boost if selected
                if use_boost:
                    snake.activate_boost()

                # Take action
                moved = snake.move(direction)
                reward = REWARD_STEP

                # Check for collision
                if not moved:
                    reward = PENALTY_COLLISION
                    done = True

                # Check for food collision
                if snake.head() == food.position:
                    snake.grow_snake()
                    reward = REWARD_FOOD
                    food.respawn(snake.positions)

                # Update boost status
                snake.update_boost()

                # Calculate distance feature
                dist_current = manhattan_distance(snake.head(), food.position)
                # You can add more sophisticated distance calculations or features here
                reward += (dist_current * REWARD_CLOSE_DISTANCE)

                # Observe new state
                next_state = get_state(snake, food, obstacles)

                # Remember the experience
                self.agent.remember(state, ACTION_INDEX.get(action, 0), reward, next_state, done)

                # Replay experience
                self.agent.replay()

                # Update state
                state = next_state

                # Update target network
                if self.total_steps % self.agent.target_update == 0:
                    self.agent.update_target_network()

                if self.render:
                    self.render_game(snake, food, obstacles, episode, step, high_score)

            # End of episode
            scores.append(snake.score)
            ma_score = np.mean(scores)
            ma_scores.append(ma_score)
            self.plot_scores.append(snake.score)
            self.plot_ma_scores.append(ma_score)

            if ma_score > high_score:
                high_score = ma_score
                self.agent.save_model(self.model_save_path)

            self.agent.update_epsilon()

            # Logging
            if episode % 100 == 0:
                print(f"Episode {episode}, Score: {snake.score}, MA Score: {ma_score:.2f}, Epsilon: {self.agent.epsilon:.2f}")

            # Plotting
            if episode % 100 == 0:
                self.plot_training(scores, ma_scores)

        # Save final model
        self.agent.save_model(self.model_save_path)
        print("Training completed.")

    def render_game(self, snake, food, obstacles, episode, step, high_score):
        screen.fill(BLACK)
        draw_grid()
        draw_obstacles(obstacles)
        draw_snake(snake)
        draw_food(food)

        # Draw boost indicator
        draw_boost_indicator(snake)

        # Render Scores
        render_text(f"Episode: {episode}", WHITE, (GRID_WIDTH * CELL_SIZE + 10, 10))
        render_text(f"Step: {step}", WHITE, (GRID_WIDTH * CELL_SIZE + 10, 40))
        render_text(f"Score: {snake.score}", GREEN, (GRID_WIDTH * CELL_SIZE + 10, 70))
        render_text(f"High MA Score: {high_score}", YELLOW, (GRID_WIDTH * CELL_SIZE + 10, 100))
        render_text(f"Epsilon: {self.agent.epsilon:.2f}", ORANGE, (GRID_WIDTH * CELL_SIZE + 10, 130))

        pygame.display.flip()
        clock.tick(FPS * (BOOST_SPEED_MULTIPLIER if snake.boost_active else 1))

    def plot_training(self, scores, ma_scores):
        plt.figure(figsize=(12,5))
        plt.title("Training Progress")
        plt.plot(self.plot_scores, label='Score per Episode')
        plt.plot(self.plot_ma_scores, label=f'Mean Score over {self.ma_window} Episodes')
        plt.xlabel('Episode')
        plt.ylabel('Score')
        plt.legend()
        plt.show()

# ---------------------- Drawing Functions ----------------------

def draw_grid():
    for x in range(0, GRID_WIDTH * CELL_SIZE, CELL_SIZE):
        pygame.draw.line(screen, GREY, (x, 0), (x, GRID_HEIGHT * CELL_SIZE))
    for y in range(0, GRID_HEIGHT * CELL_SIZE, CELL_SIZE):
        pygame.draw.line(screen, GREY, (0, y), (GRID_WIDTH * CELL_SIZE, y))

def draw_snake(snake):
    for pos in snake.positions:
        rect = pygame.Rect(pos[0] * CELL_SIZE, pos[1] * CELL_SIZE, CELL_SIZE, CELL_SIZE)
        pygame.draw.rect(screen, DARK_GREEN, rect)
        pygame.draw.rect(screen, GREEN, rect, 1)

def draw_food(food):
    rect = pygame.Rect(food.position[0] * CELL_SIZE, food.position[1] * CELL_SIZE, CELL_SIZE, CELL_SIZE)
    pygame.draw.rect(screen, RED, rect)

def draw_obstacles(obstacle):
    for pos in obstacle.positions:
        rect = pygame.Rect(pos[0] * CELL_SIZE, pos[1] * CELL_SIZE, CELL_SIZE, CELL_SIZE)
        pygame.draw.rect(screen, GREY, rect)

def draw_boost_indicator(snake):
    if snake.boost_active:
        head = snake.head()
        rect = pygame.Rect(head[0] * CELL_SIZE, head[1] * CELL_SIZE, CELL_SIZE, CELL_SIZE)
        pygame.draw.rect(screen, YELLOW, rect, 3)  # Highlight the head with a yellow border

def parse_action(action):
    direction, use_boost = action
    dir_map = {
        "UP": (0, -1),
        "DOWN": (0, 1),
        "LEFT": (-1, 0),
        "RIGHT": (1, 0)
    }
    return dir_map.get(direction, (0, -1)), use_boost

def get_state(snake, food, obstacles):
    """
    Enhanced state representation including:
    - Current direction (4 one-hot)
    - Food location (2 relative)
    - Danger straight, right, left (3)
    - Boost available (1)
    - Distance to food (1)
    - Number of obstacles near (optional)
    Total: 11 features
    """
    head_x, head_y = snake.head()
    food_x, food_y = food.position

    # Relative position of food
    food_direction = (
        1 if food_x > head_x else -1 if food_x < head_x else 0,
        1 if food_y > head_y else -1 if food_y < head_y else 0
    )

    # Directions: Danger straight, right, left
    directions = {
        'UP': (0, -1),
        'RIGHT': (1, 0),
        'DOWN': (0, 1),
        'LEFT': (-1, 0)
    }
    # Current direction
    current_dir = snake.direction
    dirs = ['UP', 'RIGHT', 'DOWN', 'LEFT']
    idx = dirs.index(direction_to_string(current_dir))

    danger_straight = is_danger(snake, current_dir, obstacles)
    danger_right = is_danger(snake, dirs[(idx + 1) % 4], obstacles)
    danger_left = is_danger(snake, dirs[(idx - 1) % 4], obstacles)

    # Boost availability
    boost_available = 1 if snake.boost_cooldown == 0 else 0

    # Distance to food
    distance = manhattan_distance(snake.head(), food.position) / (GRID_WIDTH + GRID_HEIGHT)

    state = [
        # Direction (one-hot)
        1 if current_dir == directions['UP'] else 0,
        1 if current_dir == directions['RIGHT'] else 0,
        1 if current_dir == directions['DOWN'] else 0,
        1 if current_dir == directions['LEFT'] else 0,
        # Food direction
        food_direction[0],
        food_direction[1],
        # Danger
        float(danger_straight),
        float(danger_right),
        float(danger_left),
        # Boost
        float(boost_available),
        # Distance to food
    ] + [float(distance)]

    state_tensor = torch.tensor(state, dtype=torch.float32)
    return state_tensor

def is_danger(snake, direction, obstacles):
    """Check if moving in the given direction will result in danger."""
    dir_map = {
        'UP': (0, -1),
        'DOWN': (0, 1),
        'LEFT': (-1, 0),
        'RIGHT': (1, 0)
    }
    dx, dy = direction if isinstance(direction, tuple) else dir_map.get(direction, (0, -1))
    new_head = (snake.head()[0] + dx, snake.head()[1] + dy)
    return (new_head in snake.positions) or \
           (new_head in obstacles.positions) or \
           not (0 <= new_head[0] < GRID_WIDTH and 0 <= new_head[1] < GRID_HEIGHT)

def direction_to_string(direction_tuple):
    """Convert direction tuple to string representation."""
    dir_map = {
        (0, -1): 'UP',
        (0, 1): 'DOWN',
        (-1, 0): 'LEFT',
        (1, 0): 'RIGHT'
    }
    return dir_map.get(direction_tuple, 'UP')

# ---------------------- Main Function ----------------------

def main():
    # Initialize agent
    agent = QLearningAgent(actions=ACTIONS)
    # Load existing model if available
    model_path = "snake_dqn.pth"
    agent.load_model(model_path)

    # Initialize trainer
    trainer = Trainer(agent, num_episodes=5000, render=False)

    # Start training
    trainer.train()

if __name__ == "__main__":
    main()