diff --git a/main.py b/main.py
index 9877935..04f4d2a 100644
--- a/main.py
+++ b/main.py
@@ -1,54 +1,52 @@
+import argparse
 from src.game.tetris import Tetris
-from src.game.TetrisGameManager import *
-from src.agents.agent import Agent, play_game
+from src.game.TetrisGameManager import TetrisGameManager
 from src.agents.agent_factory import create_agent
-from src.agents.heuristic import (
-    utility
-)
-from src.agents.heuristic_trainer import train
 from src.agents.geneticAlgAgentJon import GeneticAlgAgentJM
-
-def test():
-    # algAgent = GeneticAlgAgentJM()
-    # algAgent.number_of_selection(2)
-    # print(algAgent.getBestPop())
-    
+def self_play():
+    """Start a self-playing Tetris game."""
     board = Tetris()
-    agent = create_agent("heuristic")
+    manager = TetrisGameManager(board)
+    manager.startGame()
+def demonstrate(agent_type: str):
+    """Demonstrate gameplay with a specified agent."""
+    board = Tetris()
+    agent = create_agent(agent_type)
     manager = TetrisGameManager(board)
     manager.startDemo(agent)
-    
+def train_genetic_algorithm(population_size: int = 10):
+    """Train the genetic algorithm agent."""
+    alg_agent = GeneticAlgAgentJM()
+    alg_agent.number_of_selection(population_size)
+    print(alg_agent.getBestPop())
+def main():
+    """Main entry point to handle command-line arguments."""
+    parser = argparse.ArgumentParser(description="Tetris Game with different options.")
+    subparsers = parser.add_subparsers(dest="command", help="Sub-command help")
+    # Self-play parser
+    subparsers.add_parser("play", help="Start a self-playing Tetris game.")
+    # Demonstrate parser
+    demonstrate_parser = subparsers.add_parser(
+        "agent", help="Demonstrate gameplay with a specific agent."
+    )
+    demonstrate_parser.add_argument(
+        "agent", type=str, help="Agent type for demonstration."
+    )
+    # Genetic algorithm training parser
+    train_parser = subparsers.add_parser("train", help="Train the genetic algorithm agent.")
+    train_parser.add_argument(
+        "--population_size", type=int, default=10, help="Population size for the genetic algorithm."
+    )
+    # Parse the arguments
+    args = parser.parse_args()
+    # Route commands to the appropriate functions
+    if args.command == "play":
+        self_play()
+    elif args.command == "agent":
+        demonstrate(args.agent)
+    elif args.command == "train":
+        train_genetic_algorithm(args.population_size)
+    else:
+        parser.print_help()
 if __name__ == "__main__":
-
-    # game = Tetris()
-    # agent: Agent = create_agent("heuristic")
-    # sum_rows_removed = 0
-    # for i in range(10):
-    #     end_board = play_game(agent, game, 7)
-    #     end_board.printBoard()
-    #     sum_rows_removed += end_board.rowsRemoved
-
-    # print(f"Average rows removed: {sum_rows_removed / 10}")
-
-    # possible_moves = game.getPossibleBoards()
-    # for boards in possible_moves:
-    #     print(utility(boards, 0, -1, 0, 0, 0))
-    #     boards.printBoard()
-    
-    # board = Tetris()
-    # manager = TetrisGameManager(board)
-    # agent = create_agent("heuristic")
-    
-    # # manager.startGame()
-
-    # # train()
-
-
-    # algAgent = GeneticAlgAgentJM()
-    # algAgent.number_of_selection(2)
-    # print(algAgent.getBestPop())
-    
-    test()
-    
-        
-    # cProfile.run('main()', 'restats')  
+    main()
\ No newline at end of file
diff --git a/src/agents/random_agent.py b/src/agents/random_agent.py
index 845986c..fbc7169 100644
--- a/src/agents/random_agent.py
+++ b/src/agents/random_agent.py
@@ -7,8 +7,8 @@
 class RandomAgent(Agent):
     """Random agent that selects a random move from the list of possible moves"""
 
-    def result(self, board: Tetris) -> Action:
+    def result(self, board: Tetris) -> list[Action]:
         possible_actions = get_all_actions()
-        return choice(possible_actions)
+        return [choice(possible_actions)]