import pygame
import numpy as np
from config import MazeConfig
from maze import render_static

class Visualizer:
    def __init__(self, screen: pygame.Surface, grid: np.ndarray, cfg: MazeConfig):
        self.screen = screen
        self.grid = grid
        self.cfg = cfg
        self.maze_w = cfg.maze_cols * cfg.cell_size
        self.maze_h = cfg.maze_rows * cfg.cell_size
        self.hud_x = self.maze_w + 10
        # Pre-render static maze
        self.maze_surface = render_static(grid, cfg)
        # Trail surface with per-pixel alpha
        self.trail_surf = pygame.Surface((self.maze_w, self.maze_h), pygame.SRCALPHA)
        self.trail_surf.fill((0, 0, 0, 0))
        # Pre-allocate fade surface — reused every frame (avoids per-frame Surface creation)
        self._fade_surf = pygame.Surface((self.maze_w, self.maze_h), pygame.SRCALPHA)
        subtract = max(1, int(255 * (1.0 - cfg.trail_decay)))
        self._fade_surf.fill((0, 0, 0, subtract))
        # Pre-render static gradient bar for HUD
        self._gradient_bar = self._make_gradient_bar(170)
        # Font
        pygame.font.init()
        self.font = pygame.font.SysFont("monospace", cfg.hud_font_size)
        self.font_sm = pygame.font.SysFont("monospace", cfg.hud_font_size - 4)
        self.clock = pygame.time.Clock()
        self._goal_flash = 0
        self._first_goal = False

    def _make_gradient_bar(self, width: int) -> pygame.Surface:
        surf = pygame.Surface((width, 10))
        for i in range(width):
            t = i / width
            surf.set_at((i, 0), (int(255 * (1 - t)), int(255 * t), 30))
            for y in range(1, 10):
                surf.set_at((i, y), (int(255 * (1 - t)), int(255 * t), 30))
        return surf

    def _rank_color(self, rank: int, total: int) -> tuple[int, int, int]:
        t = rank / max(1, total - 1)
        return (int(255 * t), int(255 * (1.0 - t)), 30)

    def _fade_trails(self):
        self.trail_surf.blit(self._fade_surf, (0, 0), special_flags=pygame.BLEND_RGBA_SUB)

    def _draw_hud(self, gen: int, best_fitness: float, step: int, alive: int, reached: int):
        # Dark sidebar background
        sidebar_rect = pygame.Rect(self.maze_w, 0, 200, self.maze_h)
        pygame.draw.rect(self.screen, (10, 10, 16), sidebar_rect)

        lines = [
            ("GEN", f"{gen:>4}"),
            ("STEP", f"{step:>4}/{self.cfg.max_steps}"),
            ("ALIVE", f"{alive:>4}/{self.cfg.population}"),
            ("GOAL", f"{reached:>4}"),
            ("BEST", f"{best_fitness:>8.1f}"),
        ]
        y = 20
        accent = (0, 220, 200)
        for label, value in lines:
            lbl_surf = self.font_sm.render(label, True, (100, 100, 120))
            val_surf = self.font.render(value, True, accent)
            self.screen.blit(lbl_surf, (self.hud_x, y))
            self.screen.blit(val_surf, (self.hud_x, y + 16))
            y += 52

        # Legend — pre-rendered gradient bar (no per-frame line drawing)
        y += 20
        legend_label = self.font_sm.render("RANK", True, (80, 80, 100))
        self.screen.blit(legend_label, (self.hud_x, y))
        y += 18
        self.screen.blit(self._gradient_bar, (self.hud_x, y))
        best_lbl = self.font_sm.render("best", True, (0, 255, 30))
        wrst_lbl = self.font_sm.render("worst", True, (255, 0, 30))
        self.screen.blit(best_lbl, (self.hud_x, y + 14))
        self.screen.blit(wrst_lbl, (self.hud_x + 120, y + 14))

    def render(self, agents, gen: int, step: int):
        # Sort by fitness for ranking colors
        sorted_agents = sorted(agents, key=lambda a: a.fitness)

        # Check first goal
        reached = sum(1 for a in agents if a.reached_goal)
        if reached > 0 and not self._first_goal:
            self._first_goal = True
            self._goal_flash = 45

        # Blit static maze
        self.screen.blit(self.maze_surface, (0, 0))

        # Fade and draw trails
        self._fade_trails()
        for rank, agent in enumerate(sorted_agents):
            if agent.trail:
                color = self._rank_color(rank, len(sorted_agents))
                px, py = agent.trail[-1]
                pygame.draw.circle(self.trail_surf, (*color, 130), (px, py), 3)
        self.screen.blit(self.trail_surf, (0, 0))

        # Draw agents on top
        for rank, agent in enumerate(sorted_agents):
            color = self._rank_color(rank, len(sorted_agents))
            px = agent.col * self.cfg.cell_size + self.cfg.cell_size // 2
            py = agent.row * self.cfg.cell_size + self.cfg.cell_size // 2
            radius = 5 if agent.reached_goal else 3
            pygame.draw.circle(self.screen, color, (px, py), radius)
            if agent.reached_goal:
                pygame.draw.circle(self.screen, (255, 255, 255), (px, py), radius + 2, 1)

        # Goal flash overlay
        if self._goal_flash > 0:
            gx = (self.cfg.maze_cols - 1) * self.cfg.cell_size + self.cfg.cell_size // 2
            gy = (self.cfg.maze_rows - 1) * self.cfg.cell_size + self.cfg.cell_size // 2
            alpha = int(200 * self._goal_flash / 45)
            flash_surf = pygame.Surface((self.cfg.cell_size, self.cfg.cell_size), pygame.SRCALPHA)
            flash_surf.fill((255, 220, 0, alpha))
            self.screen.blit(flash_surf, (gx - self.cfg.cell_size//2, gy - self.cfg.cell_size//2))
            self._goal_flash -= 1

        # HUD
        best_fit = max((a.fitness for a in agents), default=0.0)
        alive = sum(1 for a in agents if not a.reached_goal)
        self._draw_hud(gen, best_fit, step, alive, reached)

        pygame.display.flip()

    def reset_trails(self):
        self.trail_surf.fill((0, 0, 0, 0))
        self._first_goal = False

    def new_maze(self, grid, cfg):
        self.grid = grid
        self.cfg = cfg
        self.maze_surface = render_static(grid, cfg)
        self.reset_trails()