agent_RL.py

import torch
import random
import pygame
import numpy as np
import sys
from collections import deque
from game import SnakeGameAI, Direction, Point
from model import LinearQNet, QTrainer, ReplayBuffer, MultiLinearQNet
from plotme import TrainingPlot
from genetic import GeneticAlgorithm

"""
Deep Q-Learning 
Q new value = Q current   + Learn Rate * [ Reward + Discount Rate * Max Future Expected Reward - Q_current ]
Q_new(s,a) = Q_current(s,a) + ALPHA [R (s,a) + GAMMA MAX Q'(s',a') - Q_current(s,a) ]"""

#########################################################################################################################
"""
ALPHA - LEARNING RATE DQN

The learning rate determines the extent to which newly acquired information overrides old information. 
It regulates how much the Q-values are updated based on new experiences.
A higher learning rate means faster updates, but it might lead to instability or overshooting optimal values."""

#########################################################################################################################
"""
GAMMA - DISCOUNT RATE DQN

The discount factor signifies the importance of future rewards compared to immediate rewards. 
It determines how much the agent values future rewards over immediate ones. 
A higher discount factor values long-term rewards more, influencing the agent’s decision-making."""

#########################################################################################################################
"""EPSILON - EXPLORATION RATE Q_current -> Q_new using Bellman Equation

Loss = E [(rt + GAMMA * max(Q_st+1, a', theta-) - Q (st, at, theta))^2
Loss = greedy strategy ---> starts high and decreases over time

if random_number < epsilon:
   select_random_action()
else:
   select_action_with_highest_q_value()"""

# VARIABLES

MAX_MEMORY = 100_000 # Maximum memory for the agent  

param_ranges = {
    # Continuous parameters
    'learning_rate': (0.001, 0.1), # Alpha / Higher values allow faster learning, while lower values ensure more stability
    'discount_factor': (0.9, 0.999), #Gamme / Closer to 1 indicate future rewards are highly important, emphasizing long-term rewards
    'dropout_rate': (0.1, 0.5), # Higher drops out a more neurons -> prevent overfit in complex models or datasets with limited samples
    'exploration_rate': (0.1, 0.5), #Epsilon /Higher more exploration -> Possibly better actions /Lower -> More stability using learned policy
    
    # Discrete parameters
    # 'batch_size': [10, 100, 250, 500, 1000, 2000, 5000], # Number of experiences sampled from the replay buffer for training.
    # 'activation_function': ['relu', 'sigmoid', 'tanh'],
    # 'optimizer': ['adam', 'sgd', 'rmsprop'], 
    
    # Integer parameters (num_inputs, num_outputs of NN)
    # 'num_hidden_layers': [1, 2, 3, 4, 5],
    # 'neurons_per_layer': [32, 64, 128, 256, 512, 1024]

    # Other parameters
    #'MAX_MEMORY' -> capacity of replay memory
}

MUTATION_RATE = 0.1
CROSSOVER_RATE = 0.8
POPULATION_SIZE = 20
CHROMOSOME_LENGTH, NUM_GENERATIONS = 15, 5 



class QLearningAgent:

    def __init__(self, parameters=None):
        parameters = parameters or {}
        self.n_games = 0 # Number of games played
        self.base_epsilon = float(parameters.get('exploration_rate', 0.3)) # Parameter for exploration-exploitation trade-off
        self.epsilon = self.base_epsilon
        self.gamma = float(parameters.get('discount_factor', 0.9)) # Discount factor for future rewards
        self.dropout_rate = float(parameters.get('dropout_rate', 0.2))
        self.lr = float(parameters.get('learning_rate', 0.001))
        self.memory = deque(maxlen=MAX_MEMORY) # Replay memory for storing experiences
        input_size, hidden_size, output_size, num_hidden_layers = 11, 256, 3, 1
        #num_hidden_layers = parameters.get('num_hidden_layers', 1)
        #activation_function = parameters.get('activation_function', 'relu')
        self.model = LinearQNet(input_size, hidden_size, output_size, self.dropout_rate, num_hidden_layers, activation_function = 'relu')
        
        #optimizer = parameters.get('optimizer','adam')
        self.trainer = QTrainer(self.model, lr = self.lr, gamma = self.gamma, optimizer_name = 'adam') 
        #self.batch_size = parameters.get('batch_size', batch_size)
        self.batch_size = 1000 # Learning rate for the model
        self.replay_buffer = ReplayBuffer(capacity = self.batch_size)
        self.target_model = LinearQNet(input_size, hidden_size, output_size, self.dropout_rate, num_hidden_layers, activation_function = 'relu')
        self.target_model.load_state_dict(self.model.state_dict())  # Sync initial weights

    def update_hyperparameters(self, parameters):
        self.base_epsilon = float(parameters.get('exploration_rate', self.base_epsilon))
        self.lr = float(parameters.get('learning_rate', self.lr))
        self.gamma = float(parameters.get('discount_factor', self.gamma))
        self.trainer = QTrainer(self.model, lr=self.lr, gamma=self.gamma, optimizer_name='adam')

    def get_state(self, game):
        # Function to obtain the state representation based on the game state
        # Generates information about dangers, direction, and food location
        head = game.snake[0]
        point_l = Point(head.x - 20, head.y)
        point_r = Point(head.x + 20, head.y)
        point_u = Point(head.x, head.y - 20)
        point_d = Point(head.x, head.y + 20)
        
        dir_l = game.direction == Direction.LEFT
        dir_r = game.direction == Direction.RIGHT
        dir_u = game.direction == Direction.UP
        dir_d = game.direction == Direction.DOWN

        state = [
            # Danger straight
            (dir_r and game.is_collision(point_r)) or 
            (dir_l and game.is_collision(point_l)) or 
            (dir_u and game.is_collision(point_u)) or 
            (dir_d and game.is_collision(point_d)),

            # Danger right
            (dir_u and game.is_collision(point_r)) or 
            (dir_d and game.is_collision(point_l)) or 
            (dir_l and game.is_collision(point_u)) or 
            (dir_r and game.is_collision(point_d)),

            # Danger left
            (dir_d and game.is_collision(point_r)) or 
            (dir_u and game.is_collision(point_l)) or 
            (dir_r and game.is_collision(point_u)) or 
            (dir_l and game.is_collision(point_d)),
            
            # Move direction
            dir_l,
            dir_r,
            dir_u,
            dir_d,
            
            # Food location 
            game.food.x < game.head.x,  # food left
            game.food.x > game.head.x,  # food right
            game.food.y < game.head.y,  # food up
            game.food.y > game.head.y  # food down
            ]

        return np.array(state, dtype=int)

    """Store experience (state, action, reward, next_state, done) in memory // MAX_MEMORY"""
    def remember(self, state, action, reward, next_state, done):
        self.memory.append((state, action, reward, next_state, done))


    """Sample from memory and perform a training step using QTrainer"""
    def train_long_memory(self):
        if not self.memory:
            return
        if len(self.memory) > self.batch_size:
            mini_sample = random.sample(self.memory, self.batch_size) # List of tuples
        else:
            mini_sample = self.memory

        states, actions, rewards, next_states, dones = zip(*mini_sample)
        self.trainer.train_step(states, actions, rewards, next_states, dones, ReplayBuffer, self.batch_size)

    """Perform single training step using a single experience tuple (short-term memory)"""
    def train_short_memory(self, state, action, reward, next_state, done):
        self.trainer.train_step(state, action, reward, next_state, done, ReplayBuffer, self.batch_size)

    """Select actions based on an epsilon-greedy strategy, balancing exploration and exploitation in the agent's decision-making process.
    - 1. Epsilon Decay; 2. Action Selection; 3. Outcome"""
    def get_action(self, state):
        # Select actions based on an epsilon-greedy strategy
        self.epsilon = max(5, int(self.base_epsilon * 200) - self.n_games)
        final_move = [0,0,0]
        if random.randint(0, 200) < self.epsilon:
            move = random.randint(0, 2)
            final_move[move] = 1
        else:
            state0 = torch.tensor(state, dtype=torch.float)
            prediction = self.model(state0)
            move = torch.argmax(prediction).item()
            final_move[move] = 1

        return final_move

def create_random_parameters(param_ranges): # For initialization of a population
    return {param: random.uniform(*ranges) if isinstance(ranges, tuple) else random.choice(ranges)
            for param, ranges in param_ranges.items()}

def _build_hamiltonian_cycle(width, height, block_size):
    cols = width // block_size
    rows = height // block_size
    if cols < 2 or rows < 2 or (cols * rows) % 2 != 0:
        raise ValueError("Hamiltonian cycle requires at least a 2x2 even-sized grid.")

    cells = [(0, 0)]
    for x in range(1, cols):
        cells.append((x, 0))

    for y in range(1, rows):
        if y % 2 == 1:
            cells.append((cols - 1, y))
            for x in range(cols - 2, 0, -1):
                cells.append((x, y))
        else:
            cells.append((1, y))
            for x in range(2, cols):
                cells.append((x, y))

    for y in range(rows - 1, 0, -1):
        cells.append((0, y))

    if len(cells) != cols * rows or len(set(cells)) != len(cells):
        raise ValueError("Failed to build a valid Hamiltonian cycle.")

    for i in range(len(cells)):
        x1, y1 = cells[i]
        x2, y2 = cells[(i + 1) % len(cells)]
        if abs(x1 - x2) + abs(y1 - y2) != 1:
            raise ValueError("Invalid Hamiltonian cycle adjacency.")

    return [Point(x * block_size, y * block_size) for x, y in cells]

def _target_direction(current_point, next_point):
    if next_point.x > current_point.x:
        return Direction.RIGHT
    if next_point.x < current_point.x:
        return Direction.LEFT
    if next_point.y > current_point.y:
        return Direction.DOWN
    return Direction.UP

def _direction_to_action(current_direction, wanted_direction):
    clockwise_directions = [Direction.RIGHT, Direction.DOWN, Direction.LEFT, Direction.UP]
    current_index = clockwise_directions.index(current_direction)
    wanted_index = clockwise_directions.index(wanted_direction)

    if wanted_index == current_index:
        return [1, 0, 0]
    if wanted_index == (current_index + 1) % 4:
        return [0, 1, 0]
    if wanted_index == (current_index - 1) % 4:
        return [0, 0, 1]
    return None

def train_hamiltonian(max_games=None):
    plotter = TrainingPlot()
    plot_scores = []
    plot_mean_scores = []
    total_score = 0
    record = 0
    n_games = 0
    game = SnakeGameAI()

    block_size = abs(game.snake[0].x - game.snake[1].x) if len(game.snake) > 1 else 20
    cycle = _build_hamiltonian_cycle(game.width, game.height, block_size)
    next_lookup = {cycle[i]: cycle[(i + 1) % len(cycle)] for i in range(len(cycle))}
    prev_lookup = {cycle[(i + 1) % len(cycle)]: cycle[i] for i in range(len(cycle))}

    while True:
        head = game.snake[0]
        next_candidates = [next_lookup.get(head), prev_lookup.get(head)]
        final_move = [1, 0, 0]

        for candidate in next_candidates:
            if candidate is None:
                continue
            wanted_direction = _target_direction(head, candidate)
            action = _direction_to_action(game.direction, wanted_direction)
            if action is not None:
                final_move = action
                break

        _, done, score, _, _ = game.play_step(final_move)

        if done:
            game._init_game()
            n_games += 1

            if score > record:
                record = score

            print('Hamiltonian Game', n_games, 'Score', score, 'Record:', record)
            plot_scores.append(score)
            total_score += score
            mean_score = total_score / n_games
            plot_mean_scores.append(mean_score)
            plotter.update(plot_scores, plot_mean_scores)

            if max_games is not None and n_games >= max_games:
                break

def _hamiltonian_move(game, next_lookup, prev_lookup):
    head = game.snake[0]
    next_candidates = [next_lookup.get(head), prev_lookup.get(head)]

    for candidate in next_candidates:
        if candidate is None:
            continue
        wanted_direction = _target_direction(head, candidate)
        action = _direction_to_action(game.direction, wanted_direction)
        if action is not None:
            return action

    return [1, 0, 0]

def _draw_competition_panel(surface, game, origin_x, origin_y, title, score, record, n_games, mean_score, title_font, stat_font):
    panel_rect = pygame.Rect(origin_x, origin_y, game.width, game.height)
    pygame.draw.rect(surface, (32, 32, 32), panel_rect)
    pygame.draw.rect(surface, (90, 90, 90), panel_rect, 2)

    for pt in game.snake:
        snake_rect = pygame.Rect(origin_x + pt.x, origin_y + pt.y, 20, 20)
        pygame.draw.rect(surface, (0, 0, 255), snake_rect)
        pygame.draw.rect(surface, (0, 100, 255), snake_rect.inflate(-8, -8))

    food_rect = pygame.Rect(origin_x + game.food.x, origin_y + game.food.y, 20, 20)
    pygame.draw.rect(surface, (200, 0, 0), food_rect)

    title_text = title_font.render(title, True, (255, 255, 255))
    stat_text = stat_font.render(
        f"G:{n_games} S:{score} R:{record} M:{mean_score:.2f}",
        True,
        (220, 220, 220),
    )
    surface.blit(title_text, (origin_x, origin_y - 48))
    surface.blit(stat_text, (origin_x, origin_y - 24))

def _next_direction(current_direction, action):
    clockwise_directions = [Direction.RIGHT, Direction.DOWN, Direction.LEFT, Direction.UP]
    current_index = clockwise_directions.index(current_direction)
    if action == [0, 1, 0]:
        new_index = (current_index + 1) % 4
    elif action == [0, 0, 1]:
        new_index = (current_index - 1) % 4
    else:
        new_index = current_index
    return clockwise_directions[new_index]

def _step_point(point, direction, block_size):
    if direction == Direction.RIGHT:
        return Point(point.x + block_size, point.y)
    if direction == Direction.LEFT:
        return Point(point.x - block_size, point.y)
    if direction == Direction.DOWN:
        return Point(point.x, point.y + block_size)
    return Point(point.x, point.y - block_size)

class _BattleView:
    def __init__(self, arena, snake_id):
        self._arena = arena
        self._snake_id = snake_id
        self.snake = arena.snakes[snake_id]
        self.direction = arena.directions[snake_id]
        self.food = arena.food
        self.head = self.snake[0]

    def is_collision(self, pt=None):
        if pt is None:
            pt = self.head
        return self._arena.is_collision(self._snake_id, pt)

class SnakeBattleArena:
    def __init__(self, width=640, height=640, block_size=20, speed=220):
        self.width = width
        self.height = height
        self.block_size = block_size
        self.speed = speed
        self.display = pygame.display.set_mode((self.width, self.height + 70))
        pygame.display.set_caption('Snake Battle: GA vs RL vs Hamiltonian')
        self.clock = pygame.time.Clock()
        self.title_font = pygame.font.SysFont('arial', 24)
        self.info_font = pygame.font.SysFont('arial', 18)
        self.snake_colors = {
            'ga': ((0, 190, 255), (0, 120, 210)),
            'rl': ((0, 255, 130), (0, 170, 90)),
            'ham': ((255, 170, 0), (210, 120, 0)),
        }
        self.wins = {'ga': 0, 'rl': 0, 'ham': 0}
        self.round = 0
        self._reset_round()

    def _reset_round(self):
        cx = self.width // 2
        cy = self.height // 2
        b = self.block_size

        self.snakes = {
            'ga': [Point(b * 5, b * 5), Point(b * 4, b * 5), Point(b * 3, b * 5)],
            'rl': [Point(self.width - b * 6, b * 5), Point(self.width - b * 5, b * 5), Point(self.width - b * 4, b * 5)],
            'ham': [Point(cx, self.height - b * 6), Point(cx, self.height - b * 5), Point(cx, self.height - b * 4)],
        }
        self.directions = {'ga': Direction.RIGHT, 'rl': Direction.LEFT, 'ham': Direction.UP}
        self.alive = {'ga': True, 'rl': True, 'ham': True}
        self.scores = {'ga': 0, 'rl': 0, 'ham': 0}
        self.frame_iteration = 0
        self.round += 1
        self._place_food()

    def _occupied_points(self):
        occupied = set()
        for snake in self.snakes.values():
            occupied.update(snake)
        return occupied

    def _place_food(self):
        occupied = self._occupied_points()
        while True:
            x = random.randint(0, (self.width - self.block_size) // self.block_size) * self.block_size
            y = random.randint(0, (self.height - self.block_size) // self.block_size) * self.block_size
            candidate = Point(x, y)
            if candidate not in occupied:
                self.food = candidate
                return

    def is_collision(self, snake_id, pt):
        if pt.x < 0 or pt.y < 0 or pt.x > self.width - self.block_size or pt.y > self.height - self.block_size:
            return True
        for other_id, snake in self.snakes.items():
            body = snake if other_id != snake_id else snake[1:]
            if pt in body:
                return True
        return False

    def get_view(self, snake_id):
        return _BattleView(self, snake_id)

    def step(self, actions):
        self.frame_iteration += 1
        rewards = {'ga': 0, 'rl': 0, 'ham': 0}
        new_heads = {}
        new_directions = {}

        for snake_id in self.snakes.keys():
            if not self.alive[snake_id]:
                continue
            new_direction = _next_direction(self.directions[snake_id], actions.get(snake_id, [1, 0, 0]))
            new_head = _step_point(self.snakes[snake_id][0], new_direction, self.block_size)
            new_directions[snake_id] = new_direction
            new_heads[snake_id] = new_head

        dead = set()
        seen_heads = {}
        for snake_id, new_head in new_heads.items():
            if self.is_collision(snake_id, new_head):
                dead.add(snake_id)
            if new_head in seen_heads:
                dead.add(snake_id)
                dead.add(seen_heads[new_head])
            seen_heads[new_head] = snake_id

        ate_food = False
        for snake_id in self.snakes.keys():
            if not self.alive[snake_id]:
                continue
            if snake_id in dead:
                self.alive[snake_id] = False
                rewards[snake_id] = -10
                continue

            self.directions[snake_id] = new_directions[snake_id]
            self.snakes[snake_id].insert(0, new_heads[snake_id])
            if new_heads[snake_id] == self.food:
                self.scores[snake_id] += 1
                rewards[snake_id] = 10
                ate_food = True
            else:
                self.snakes[snake_id].pop()
                rewards[snake_id] = 0

        if ate_food:
            self._place_food()

        alive_ids = [sid for sid, alive in self.alive.items() if alive]
        round_done = len(alive_ids) <= 1 or self.frame_iteration > 200
        if round_done and len(alive_ids) == 1:
            self.wins[alive_ids[0]] += 1

        return rewards, round_done

    def render(self):
        self.display.fill((24, 24, 24))
        board_origin_y = 70
        pygame.draw.rect(self.display, (34, 34, 34), pygame.Rect(0, board_origin_y, self.width, self.height))

        food_rect = pygame.Rect(self.food.x, board_origin_y + self.food.y, self.block_size, self.block_size)
        pygame.draw.rect(self.display, (200, 0, 0), food_rect)

        for snake_id, snake in self.snakes.items():
            outer, inner = self.snake_colors[snake_id]
            for pt in snake:
                rect = pygame.Rect(pt.x, board_origin_y + pt.y, self.block_size, self.block_size)
                pygame.draw.rect(self.display, outer, rect)
                pygame.draw.rect(self.display, inner, rect.inflate(-8, -8))

        header = self.title_font.render(
            f"Round {self.round} | GA W:{self.wins['ga']} S:{self.scores['ga']}  RL W:{self.wins['rl']} S:{self.scores['rl']}  HAM W:{self.wins['ham']} S:{self.scores['ham']}",
            True, (255, 255, 255)
        )
        info = self.info_font.render("Close window to stop. Mode: same-board battle", True, (200, 200, 200))
        self.display.blit(header, (8, 10))
        self.display.blit(info, (8, 40))
        pygame.display.flip()
        self.clock.tick(self.speed)

def train_battle(max_rounds=None):
    arena = SnakeBattleArena()
    ga_genetic = GeneticAlgorithm(
        POPULATION_SIZE=POPULATION_SIZE,
        param_ranges=param_ranges,
        MUTATION_RATE=MUTATION_RATE,
        NUM_GENERATIONS=NUM_GENERATIONS
    )
    ga_agent = QLearningAgent(ga_genetic.get_current_parameters())
    rl_agent = QLearningAgent()
    ga_record = 0
    ga_metrics = []

    block_size = arena.block_size
    cycle = _build_hamiltonian_cycle(arena.width, arena.height, block_size)
    next_lookup = {cycle[i]: cycle[(i + 1) % len(cycle)] for i in range(len(cycle))}
    prev_lookup = {cycle[(i + 1) % len(cycle)]: cycle[i] for i in range(len(cycle))}

    completed_rounds = 0
    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                return

        ga_view_old = arena.get_view('ga')
        rl_view_old = arena.get_view('rl')
        ga_state_old = ga_agent.get_state(ga_view_old)
        rl_state_old = rl_agent.get_state(rl_view_old)

        ga_action = ga_agent.get_action(ga_state_old) if arena.alive['ga'] else [1, 0, 0]
        rl_action = rl_agent.get_action(rl_state_old) if arena.alive['rl'] else [1, 0, 0]
        ham_action = _hamiltonian_move(arena.get_view('ham'), next_lookup, prev_lookup) if arena.alive['ham'] else [1, 0, 0]
        rewards, round_done = arena.step({'ga': ga_action, 'rl': rl_action, 'ham': ham_action})

        ga_view_new = arena.get_view('ga')
        rl_view_new = arena.get_view('rl')
        ga_state_new = ga_agent.get_state(ga_view_new)
        rl_state_new = rl_agent.get_state(rl_view_new)
        ga_done = (not arena.alive['ga']) or round_done
        rl_done = (not arena.alive['rl']) or round_done

        ga_agent.train_short_memory(ga_state_old, ga_action, rewards['ga'], ga_state_new, ga_done)
        ga_agent.remember(ga_state_old, ga_action, rewards['ga'], ga_state_new, ga_done)
        rl_agent.train_short_memory(rl_state_old, rl_action, rewards['rl'], rl_state_new, rl_done)
        rl_agent.remember(rl_state_old, rl_action, rewards['rl'], rl_state_new, rl_done)

        if round_done:
            completed_rounds += 1
            ga_agent.n_games += 1
            rl_agent.n_games += 1
            ga_agent.train_long_memory()
            rl_agent.train_long_memory()

            ga_score = arena.scores['ga']
            ga_record = max(ga_record, ga_score)
            ga_metrics.append({
                'score': ga_score,
                'record': ga_record,
                'steps': arena.frame_iteration,
                'collisions': 0,
                'same_positions': 0
            })
            _, best_params, _ = ga_genetic.genetic(
                NUM_GENERATIONS,
                score=ga_score,
                record=ga_record,
                steps=arena.frame_iteration,
                collisions=0,
                same_positions_counter=0,
                game_metrics_list=ga_metrics
            )
            ga_agent.update_hyperparameters(ga_genetic.get_current_parameters())
            if completed_rounds % 20 == 0 and isinstance(best_params, dict):
                ga_agent.update_hyperparameters(best_params)

            if max_rounds is not None and completed_rounds >= max_rounds:
                break
            arena._reset_round()

        arena.render()

def train_competition(max_games=None):
    panel_size = 320
    margin = 12
    top_space = 64
    window_width = panel_size * 3 + margin * 4
    window_height = panel_size + top_space + margin

    screen = pygame.display.set_mode((window_width, window_height))
    pygame.display.set_caption('Snake: GA vs RL vs Hamiltonian')
    clock = pygame.time.Clock()
    title_font = pygame.font.SysFont('arial', 22)
    stat_font = pygame.font.SysFont('arial', 17)

    ga_game = SnakeGameAI(width=panel_size, height=panel_size, render=False)
    rl_game = SnakeGameAI(width=panel_size, height=panel_size, render=False)
    ham_game = SnakeGameAI(width=panel_size, height=panel_size, render=False)

    ga_genetic = GeneticAlgorithm(
        POPULATION_SIZE=POPULATION_SIZE,
        param_ranges=param_ranges,
        MUTATION_RATE=MUTATION_RATE,
        NUM_GENERATIONS=NUM_GENERATIONS
    )
    ga_agent = QLearningAgent(ga_genetic.get_current_parameters())
    rl_agent = QLearningAgent()

    ga_visited_positions = set()
    ga_same_positions_counter = 0
    ga_game_metrics = []
    ga_record = 0
    ga_total_score = 0

    rl_record = 0
    rl_total_score = 0

    ham_record = 0
    ham_total_score = 0
    ham_games = 0

    block_size = abs(ham_game.snake[0].x - ham_game.snake[1].x) if len(ham_game.snake) > 1 else 20
    cycle = _build_hamiltonian_cycle(ham_game.width, ham_game.height, block_size)
    next_lookup = {cycle[i]: cycle[(i + 1) % len(cycle)] for i in range(len(cycle))}
    prev_lookup = {cycle[(i + 1) % len(cycle)]: cycle[i] for i in range(len(cycle))}

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                return

        ga_state_old = ga_agent.get_state(ga_game)
        ga_move = ga_agent.get_action(ga_state_old)
        ga_reward, ga_done, ga_score, ga_collisions, ga_steps = ga_game.play_step(ga_move)
        ga_state_new = ga_agent.get_state(ga_game)
        ga_agent.train_short_memory(ga_state_old, ga_move, ga_reward, ga_state_new, ga_done)
        ga_agent.remember(ga_state_old, ga_move, ga_reward, ga_state_new, ga_done)

        ga_position = (ga_game.snake[0].x, ga_game.snake[0].y)
        if ga_position in ga_visited_positions:
            ga_same_positions_counter += 1
        ga_visited_positions.add(ga_position)

        if ga_done:
            ga_game._init_game()
            ga_agent.n_games += 1
            ga_agent.train_long_memory()
            ga_record = max(ga_record, ga_score)
            ga_total_score += ga_score
            ga_game_metrics.append({
                'score': ga_score,
                'record': ga_record,
                'steps': ga_steps,
                'collisions': ga_collisions,
                'same_positions': ga_same_positions_counter
            })
            _, best_params, _ = ga_genetic.genetic(
                NUM_GENERATIONS,
                score=ga_score,
                record=ga_record,
                steps=ga_steps,
                collisions=ga_collisions,
                same_positions_counter=ga_same_positions_counter,
                game_metrics_list=ga_game_metrics
            )
            ga_agent.update_hyperparameters(ga_genetic.get_current_parameters())
            if ga_agent.n_games % 20 == 0 and isinstance(best_params, dict):
                ga_agent.update_hyperparameters(best_params)
            ga_same_positions_counter = 0
            ga_visited_positions.clear()

        rl_state_old = rl_agent.get_state(rl_game)
        rl_move = rl_agent.get_action(rl_state_old)
        rl_reward, rl_done, rl_score, _, _ = rl_game.play_step(rl_move)
        rl_state_new = rl_agent.get_state(rl_game)
        rl_agent.train_short_memory(rl_state_old, rl_move, rl_reward, rl_state_new, rl_done)
        rl_agent.remember(rl_state_old, rl_move, rl_reward, rl_state_new, rl_done)

        if rl_done:
            rl_game._init_game()
            rl_agent.n_games += 1
            rl_agent.train_long_memory()
            rl_record = max(rl_record, rl_score)
            rl_total_score += rl_score

        ham_move = _hamiltonian_move(ham_game, next_lookup, prev_lookup)
        _, ham_done, ham_score, _, _ = ham_game.play_step(ham_move)
        if ham_done:
            ham_game._init_game()
            ham_games += 1
            ham_record = max(ham_record, ham_score)
            ham_total_score += ham_score

        ga_mean = ga_total_score / ga_agent.n_games if ga_agent.n_games else 0.0
        rl_mean = rl_total_score / rl_agent.n_games if rl_agent.n_games else 0.0
        ham_mean = ham_total_score / ham_games if ham_games else 0.0

        screen.fill((16, 16, 16))
        _draw_competition_panel(
            screen, ga_game, margin, top_space, 'GA-Optimized RL',
            ga_game.score, ga_record, ga_agent.n_games, ga_mean, title_font, stat_font
        )
        _draw_competition_panel(
            screen, rl_game, margin * 2 + panel_size, top_space, 'Standard RL',
            rl_game.score, rl_record, rl_agent.n_games, rl_mean, title_font, stat_font
        )
        _draw_competition_panel(
            screen, ham_game, margin * 3 + panel_size * 2, top_space, 'Hamiltonian',
            ham_game.score, ham_record, ham_games, ham_mean, title_font, stat_font
        )
        pygame.display.flip()
        clock.tick(30)

        if max_games is not None and (
            ga_agent.n_games >= max_games and
            rl_agent.n_games >= max_games and
            ham_games >= max_games
        ):
            break

def train(max_games=None):
    plotter = TrainingPlot() # To store game scores for plotting
    plot_scores = [] # To store scores for plotting  
    plot_mean_scores = []  # To store mean scores for plotting 
    total_score = 0
    record = 0
    visited_positions = set() # Unique Values
    same_positions_counter = 0

    # Initialize the agent with random parameters from param_range
    genetic = GeneticAlgorithm(
                            POPULATION_SIZE = POPULATION_SIZE,
                            param_ranges = param_ranges, 
                            MUTATION_RATE = MUTATION_RATE,
                            NUM_GENERATIONS = NUM_GENERATIONS
                            )
    agent = QLearningAgent(genetic.get_current_parameters()) 
    game = SnakeGameAI(speed=1080) # Slightly faster GA training loop
    
    game_metrics_list = []  # List to store game metrics (score, record, steps, collisions, same positions)
    
    while True: 
        # Capture the state before taking an action
        state_old = agent.get_state(game) 
        # Determine the next move/action using the RL agent
        final_move = agent.get_action(state_old)
        # Execute the selected move and observe the game's response
        reward, done, score, collisions, steps = game.play_step(final_move)
        # Update the agent's internal score
        game.score = score
        # Capture the new state after the action
        state_new = agent.get_state(game)
        # Train the agent using this experience
        agent.train_short_memory(state_old, final_move, reward, state_new, done)
        agent.remember(state_old, final_move, reward, state_new, done)

        # Get the current position of the snake's head
        current_position = (game.snake[0].x, game.snake[0].y)

        if current_position in visited_positions:
            same_positions_counter += 1

        # Add the current position to the set of visited positions
        visited_positions.add(current_position)
            
        # # Check if the snake's head position has been visited before
        # same_positions = len(visited_positions) != len(set(visited_positions))


        if done: # He is dead

            # Initialize the game for the next iteration
            game._init_game()
            agent.n_games += 1
            agent.train_long_memory()

            # Update the highest record if the current score surpasses
            if score > record:
                record = score
                agent.model.save()

            print('Game', agent.n_games, 'Score', score, 'Record:', record)

            # Update plot data for visualization
            plot_scores.append(score)
            total_score += score
            mean_score = total_score / agent.n_games
            plot_mean_scores.append(mean_score)
            if agent.n_games <= 10 or agent.n_games % 2 == 0:
                plotter.update(plot_scores, plot_mean_scores)

            # Store game metrics in a dictionary
            game_metrics = {
                'score': score,
                'record': record,
                'steps': steps,
                'collisions': collisions,
                'same_positions': same_positions_counter
            }

            game_metrics_list.append(game_metrics)

            # Evaluate one GA candidate per game and rotate candidates continuously.
            _, best_params, _ = genetic.genetic(NUM_GENERATIONS, score = score, record = record, steps = steps, 
                                                    collisions = collisions, same_positions_counter = same_positions_counter,  
                                                    game_metrics_list = game_metrics_list)
            agent.update_hyperparameters(genetic.get_current_parameters())
            if agent.n_games % 20 == 0 and isinstance(best_params, dict):
                agent.update_hyperparameters(best_params)

            same_positions_counter = 0
            steps = 0

            if max_games is not None and agent.n_games >= max_games:
                break

def train_RL(max_games=None):
    plotter = TrainingPlot() # To store game scores for plotting
    plot_scores = [] # To store scores for plotting  
    plot_mean_scores = []  # To store mean scores for plotting 
    total_score = 0
    record = 0
    visited_positions = set() # Unique Values
    same_positions_counter = 0
    agent = QLearningAgent() 
    game = SnakeGameAI() # Initialize the game environment
    game_metrics_list = []  # List to store game metrics (score, record, steps, collisions, same positions)

    while True:
        # get old state
        state_old = agent.get_state(game)

        # get move
        final_move = agent.get_action(state_old)

        # perform move and get new state
        reward, done, score, collisions, steps = game.play_step(final_move)
        state_new = agent.get_state(game)

        # train short memory
        agent.train_short_memory(state_old, final_move, reward, state_new, done)

        # remember
        agent.remember(state_old, final_move, reward, state_new, done)

        # Get the current position of the snake's head
        current_position = (game.snake[0].x, game.snake[0].y)

        if current_position in visited_positions:
            same_positions_counter += 1

        # Add the current position to the set of visited positions
        visited_positions.add(current_position)

        if done:
            # train long memory, plot result
            game._init_game()
            agent.n_games += 1
            agent.train_long_memory()

            if score > record:
                record = score
                agent.model.save()

            print('Game', agent.n_games, 'Score', score, 'Record:', record)

            # Update plot data for visualization
            plot_scores.append(score)
            total_score += score
            mean_score = total_score / agent.n_games
            plot_mean_scores.append(mean_score)
            plotter.update(plot_scores, plot_mean_scores)

            # Store game metrics in a dictionary
            game_metrics = {
                'score': score,
                'record': record,
                'steps': steps,
                'collisions': collisions,
                'same_positions': same_positions_counter
            }

            game_metrics_list.append(game_metrics)

            if max_games is not None and agent.n_games >= max_games:
                break

def train_standard_rl(max_games=None):
    train_RL(max_games=max_games)

def train_ga_optimized(max_games=None):
    train(max_games=max_games)

def train_hamiltonian_cycle(max_games=None):
    train_hamiltonian(max_games=max_games)

def train_competition_mode(max_games=None):
    train_competition(max_games=max_games)

def train_battle_mode(max_rounds=None):
    train_battle(max_rounds=max_rounds)

if __name__ == "__main__": 
    mode = sys.argv[1].lower() if len(sys.argv) > 1 else "ga"
    if mode in ("rl", "standard", "baseline"):
        train_standard_rl()
    elif mode in ("ham", "hamiltonian", "cycle"):
        train_hamiltonian_cycle()
    elif mode in ("compare", "competition", "all3"):
        train_competition_mode()
    elif mode in ("battle", "arena", "fight", "vs"):
        train_battle_mode()
    else:
        train_ga_optimized()
    train_RL()