import numpy as np
from maze import get_walls, is_passable, bfs_distances
from config import MazeConfig


class Genome:
    """Direct action-sequence genome: each gene is a move (0=N 1=S 2=E 3=W).
    The GA evolves the sequence directly — no NN weights to get in the way.
    This is the approach used in "AI learns maze" viral videos.
    """
    N_ACTIONS = 4

    def __init__(self, genes: np.ndarray | None = None, cfg: MazeConfig | None = None):
        steps = cfg.max_steps if cfg is not None else 300
        if genes is None:
            genes = np.random.randint(0, self.N_ACTIONS, size=steps, dtype=np.uint8)
        self.genes = genes

    @property
    def weights(self) -> np.ndarray:
        """Alias so genetic.py can treat Genome like NeuralNet."""
        return self.genes.astype(np.float32)

    @classmethod
    def from_weights(cls, w: np.ndarray) -> "Genome":
        return cls(genes=np.clip(np.round(w), 0, 3).astype(np.uint8))


class Agent:
    def __init__(self, genome: Genome, start_row: int, start_col: int):
        self.net = genome          # kept as .net so visualizer/main.py don't need changes
        self.genome = genome
        self.row = start_row
        self.col = start_col
        self.reached_goal = False
        self.steps = 0
        self.fitness = 0.0
        self.trail: list[tuple[int, int]] = []
        self._visited: set[tuple[int, int]] = {(start_row, start_col)}
        self._visit_counts: dict[tuple[int, int], int] = {(start_row, start_col): 1}
        self.wall_hits = 0
        self._min_bfs: int = 10000   # closest BFS approach to goal this episode
        self._best_step: int = 0     # genome index when min_bfs was achieved
        self._bfs: np.ndarray | None = None  # set externally before stepping

    def _pixel_center(self, cfg: MazeConfig) -> tuple[int, int]:
        return (self.col * cfg.cell_size + cfg.cell_size // 2,
                self.row * cfg.cell_size + cfg.cell_size // 2)

    def step(self, grid, goal_row: int, goal_col: int, cfg: MazeConfig):
        if self.reached_goal or self.steps >= len(self.genome.genes):
            return
        action = int(self.genome.genes[self.steps])
        dr = [-1, 1, 0, 0][action]
        dc = [0,  0, 1, -1][action]

        if is_passable(grid, self.row, self.col, action):
            self.row += dr
            self.col += dc
        else:
            self.wall_hits += 1

        self.steps += 1
        cell = (self.row, self.col)
        self._visited.add(cell)
        self._visit_counts[cell] = self._visit_counts.get(cell, 0) + 1
        self.trail.append(self._pixel_center(cfg))

        if self._bfs is not None:
            bfs_d = int(self._bfs[self.row, self.col])
            if bfs_d >= 0 and bfs_d < self._min_bfs:
                self._min_bfs = bfs_d
                self._best_step = self.steps  # record when we achieved best approach

        if self.row == goal_row and self.col == goal_col:
            self.reached_goal = True

    def compute_fitness(self, goal_row: int, goal_col: int, cfg: MazeConfig) -> float:
        # BFS distance from start (0,0) to goal = max_bfs
        # best BFS approach: the closest the agent ever got (through actual maze corridors)
        max_bfs = int(self._bfs[0, 0]) if self._bfs is not None else (goal_row + goal_col)
        best_bfs = self._min_bfs if self._min_bfs < 10000 else max_bfs
        unique_cells = len(self._visited)
        total_visits = sum(self._visit_counts.values())
        revisit_penalty = total_visits - unique_cells

        self.fitness = (
            (max_bfs - best_bfs) * 12.0             # real maze proximity (primary)
            + unique_cells * 1.0                     # exploration
            - revisit_penalty * 0.1                  # penalize looping
            - self.wall_hits * 0.05
            + (cfg.goal_bonus if self.reached_goal else 0.0)
            + ((cfg.max_steps - self.steps) * 5.0 if self.reached_goal else 0.0)
        )
        return self.fitness


# Keep NeuralNet as alias so any external code importing it still works
NeuralNet = Genome