import argparse
import sys
import os
from config import MazeConfig
from maze import generate, bfs_distances, find_short_seed
from agent import Agent, Genome
from genetic import next_generation


def handle_events() -> bool:
    """Returns False if user wants to quit."""
    import pygame
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            return False
        if event.type == pygame.KEYDOWN and event.key == pygame.K_ESCAPE:
            return False
    return True


def run_generation_headless(agents, grid, goal_row, goal_col, cfg):
    """Run one generation without any rendering."""
    for step in range(cfg.max_steps):
        for agent in agents:
            agent.step(grid, goal_row, goal_col, cfg)


def run_generation_visual(agents, grid, goal_row, goal_col, cfg, viz, clock, fps, gen, steps_per_frame) -> bool:
    """Run one generation with rendering. Returns False if user quit."""
    step = 0
    while step < cfg.max_steps:
        # Advance simulation N steps per render frame
        for _ in range(steps_per_frame):
            if step >= cfg.max_steps:
                break
            for agent in agents:
                agent.step(grid, goal_row, goal_col, cfg)
            step += 1
        if not handle_events():
            return False
        viz.render(agents, gen, step)
        clock.tick(fps)
    return True


def main():
    parser = argparse.ArgumentParser(description="Maze ML — Genetic Algorithm Visualizer")
    parser.add_argument("--fast", action="store_true", help="Train without rendering (headless)")
    parser.add_argument("--generations", type=int, default=200)
    parser.add_argument("--maze-size", type=int, default=25, help="Maze cols and rows (square)")
    parser.add_argument("--seed", type=int, default=None)
    parser.add_argument("--speed", type=int, default=60, help="FPS cap for visual mode")
    parser.add_argument("--steps-per-frame", type=int, default=3, help="Simulation steps rendered per frame (higher = faster)")
    args = parser.parse_args()

    cfg = MazeConfig(maze_cols=args.maze_size, maze_rows=args.maze_size)
    goal_row = cfg.maze_rows - 1
    goal_col = cfg.maze_cols - 1

    if args.seed is None:
        print("Auto-selecting a maze with a short solution path...")
        chosen_seed, grid, bfs = find_short_seed(cfg.maze_rows, cfg.maze_cols)
        print(f"Using seed={chosen_seed}, BFS path={int(bfs[0,0])} steps")
    else:
        grid = generate(cfg.maze_rows, cfg.maze_cols, seed=args.seed)
        bfs = bfs_distances(grid, goal_row, goal_col)
        path_len = int(bfs[0, 0])
        print(f"Maze BFS path: {path_len} steps")
        if path_len > cfg.maze_rows * cfg.maze_cols // 3:
            print(f"  [hint: long path ({path_len} steps) -- try omitting --seed for auto-selection]")

    def make_agents(genomes):
        agents = [Agent(g, 0, 0) for g in genomes]
        for a in agents:
            a._bfs = bfs
        return agents

    # Initialize population
    nets = [Genome(cfg=cfg) for _ in range(cfg.population)]

    print(f"Starting training: {args.generations} generations, population={cfg.population}")

    if args.fast:
        # --- Headless fast mode: train without display ---
        print("Fast mode: rendering disabled during training.")

        for gen in range(args.generations):
            agents = make_agents(nets)
            run_generation_headless(agents, grid, goal_row, goal_col, cfg)

            for agent in agents:
                agent.compute_fitness(goal_row, goal_col, cfg)

            best_fit = max(a.fitness for a in agents)
            reached = sum(1 for a in agents if a.reached_goal)
            print(f"Gen {gen+1:>4}/{args.generations} | Best fitness: {best_fit:>10.2f} | Reached goal: {reached}/{cfg.population}")

            nets = next_generation(agents, cfg)

        # Init pygame now for the visual replay
        print("\nTraining complete. Opening replay window...")
        import pygame
        from visualizer import Visualizer

        pygame.init()
        win_w = cfg.maze_cols * cfg.cell_size + 200
        win_h = cfg.maze_rows * cfg.cell_size
        screen = pygame.display.set_mode((win_w, win_h))
        pygame.display.set_caption(cfg.window_title + " — Replay")
        clock = pygame.time.Clock()
        viz = Visualizer(screen, grid, cfg)

        replay_agents = make_agents(nets)

        ok = run_generation_visual(replay_agents, grid, goal_row, goal_col, cfg, viz, clock, args.speed, args.generations, args.steps_per_frame)
        if ok:
            print("Done. Press ESC or close window to exit.")
            running = True
            while running:
                running = handle_events()
                clock.tick(30)
        pygame.quit()
        sys.exit(0)

    else:
        # --- Visual mode: initialise pygame ---
        import pygame
        from visualizer import Visualizer

        pygame.init()
        win_w = cfg.maze_cols * cfg.cell_size + 200
        win_h = cfg.maze_rows * cfg.cell_size
        screen = pygame.display.set_mode((win_w, win_h))
        pygame.display.set_caption(cfg.window_title)
        clock = pygame.time.Clock()

        viz = Visualizer(screen, grid, cfg)

        for gen in range(args.generations):
            agents = make_agents(nets)

            ok = run_generation_visual(agents, grid, goal_row, goal_col, cfg, viz, clock, args.speed, gen, args.steps_per_frame)
            if not ok:
                print("Quit by user.")
                pygame.quit()
                sys.exit(0)

            for agent in agents:
                agent.compute_fitness(goal_row, goal_col, cfg)

            best_fit = max(a.fitness for a in agents)
            reached = sum(1 for a in agents if a.reached_goal)
            print(f"Gen {gen+1:>4}/{args.generations} | Best fitness: {best_fit:>10.2f} | Reached goal: {reached}/{cfg.population}")

            nets = next_generation(agents, cfg)
            viz.reset_trails()

        # Final visual replay with best generation
        print("\nTraining complete. Showing final generation replay...")
        replay_agents = make_agents(nets)

        viz.reset_trails()
        run_generation_visual(replay_agents, grid, goal_row, goal_col, cfg, viz, clock, args.speed, args.generations, args.steps_per_frame)

        print("Done. Press ESC or close window to exit.")
        running = True
        while running:
            running = handle_events()
            clock.tick(30)

        pygame.quit()


if __name__ == "__main__":
    main()