import numpy as np
from maze import get_walls, is_passable, NORTH, SOUTH, EAST, WEST
from config import MazeConfig

class NeuralNet:
    WEIGHT_COUNT = 8*12 + 12 + 12*4 + 4  # = 160

    def __init__(self, weights=None):
        if weights is None:
            weights = np.random.randn(self.WEIGHT_COUNT) * 0.5
        self.weights = weights.astype(np.float32)

    def forward(self, x: np.ndarray) -> int:
        # x shape: (8,) float32
        # Slices: W1=weights[0:96].reshape(8,12), b1=weights[96:108]
        #         W2=weights[108:156].reshape(12,4), b2=weights[156:160]
        # Forward: h = ReLU(x @ W1 + b1), logits = h @ W2 + b2
        # Return argmax(logits) — deterministic, no sampling
        W1 = self.weights[0:96].reshape(8, 12)
        b1 = self.weights[96:108]
        W2 = self.weights[108:156].reshape(12, 4)
        b2 = self.weights[156:160]
        h = np.maximum(0.0, x @ W1 + b1)
        logits = h @ W2 + b2
        return int(np.argmax(logits))


class Agent:
    def __init__(self, net: NeuralNet, start_row: int, start_col: int):
        self.net = net
        self.row = start_row
        self.col = start_col
        self.alive = True
        self.reached_goal = False
        self.steps = 0
        self.fitness = 0.0
        self.trail: list[tuple[int,int]] = []  # pixel coords
        self.dist_traveled_toward = 0.0

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

    def get_inputs(self, grid, goal_row, goal_col, cfg) -> np.ndarray:
        wN, wS, wE, wW = get_walls(grid, self.row, self.col)
        dx = (goal_col - self.col) / cfg.maze_cols
        dy = (goal_row - self.row) / cfg.maze_rows
        px = self.col / cfg.maze_cols
        py = self.row / cfg.maze_rows
        return np.array([wN, wS, wE, wW, dx, dy, px, py], dtype=np.float32)

    def step(self, grid, goal_row, goal_col, cfg):
        if self.reached_goal:
            return
        x = self.get_inputs(grid, goal_row, goal_col, cfg)
        action = self.net.forward(x)
        dr = [-1, 1, 0, 0][action]
        dc = [0, 0, 1, -1][action]
        prev_dist = abs(goal_row - self.row) + abs(goal_col - self.col)
        if is_passable(grid, self.row, self.col, action):
            self.row += dr
            self.col += dc
            new_dist = abs(goal_row - self.row) + abs(goal_col - self.col)
            self.dist_traveled_toward += max(0.0, prev_dist - new_dist)
        self.steps += 1
        self.trail.append(self._pixel_center(cfg))
        if self.row == goal_row and self.col == goal_col:
            self.reached_goal = True

    def compute_fitness(self, goal_row, goal_col, cfg):
        dist = abs(goal_row - self.row) + abs(goal_col - self.col)
        self.fitness = (
            self.dist_traveled_toward
            - dist * 0.5
            + (cfg.goal_bonus if self.reached_goal else 0.0)
            - self.steps * cfg.step_penalty
        )
        return self.fitness