const std = @import("std");
const Allocator = std.mem.Allocator;

pub const GRID_SIZE = 100;
pub const NUM_ANTS = 20;

pub const TILE_EMPTY = 0;
pub const TILE_WALL = 1;
pub const TILE_FOOD = 2;

pub const Ant = struct {
    x: usize,
    y: usize,
    alive: bool,
    energy: f32,

    gene_speed: f32,
    gene_metabolism: f32,

    generation: usize,
    score: usize,
};

pub const World = struct {
    grid: [GRID_SIZE][GRID_SIZE]u8,
    pheromones: [GRID_SIZE][GRID_SIZE]f32,
    ants: [NUM_ANTS]Ant,
    food_x: usize,
    food_y: usize,

    best_ant_idx: usize,
    prng: std.Random.DefaultPrng,

    pub fn init(seed: u64) World {
        var w = World{
            .grid = undefined,
            .pheromones = undefined,
            .ants = undefined,
            .food_x = 0,
            .food_y = 0,
            .best_ant_idx = 0,
            .prng = std.Random.DefaultPrng.init(seed),
        };
        w.reset();
        return w;
    }

    pub fn reset(self: *World) void {
        const random = self.prng.random();

        for (0..GRID_SIZE) |y| {
            for (0..GRID_SIZE) |x| {
                self.pheromones[y][x] = 0.0;
                if (x == 0 or y == 0 or x == GRID_SIZE - 1 or y == GRID_SIZE - 1) {
                    self.grid[y][x] = TILE_WALL;
                } else {
                    self.grid[y][x] = TILE_EMPTY;
                }
            }
        }
        self.respawnFood();

        const center = GRID_SIZE / 2;
        for (0..NUM_ANTS) |i| {
            self.ants[i] = Ant{
                .x = center,
                .y = center,
                .alive = true,
                .energy = 100.0,
                .gene_speed = 0.8 + random.float(f32) * 0.4,
                .gene_metabolism = 1.0,
                .generation = 0,
                .score = 0,
            };
        }
    }

    fn respawnFood(self: *World) void {
        const random = self.prng.random();
        while (true) {
            const rx = random.intRangeAtMost(usize, 1, GRID_SIZE - 2);
            const ry = random.intRangeAtMost(usize, 1, GRID_SIZE - 2);
            if (self.grid[ry][rx] != TILE_WALL) {
                self.food_x = rx;
                self.food_y = ry;
                break;
            }
        }
    }

    fn getScent(self: *World, x: usize, y: usize) f32 {
        const dx = if (x > self.food_x) x - self.food_x else self.food_x - x;
        const dy = if (y > self.food_y) y - self.food_y else self.food_y - y;
        const dist = @as(f32, @floatFromInt(dx + dy));
        return 1.0 / (dist + 1.0);
    }

    fn isOccupied(self: *World, x: usize, y: usize) bool {
        for (self.ants) |ant| {
            if (ant.alive and ant.x == x and ant.y == y) return true;
        }
        return false;
    }

    pub fn spawnChild(self: *World, parent_idx: usize) Ant {
        const parent = self.ants[parent_idx];
        const random = self.prng.random();

        var new_speed = parent.gene_speed + (random.float(f32) * 0.2 - 0.1);

        if (new_speed < 0.5) new_speed = 0.5;
        if (new_speed > 3.0) new_speed = 3.0;

        const new_metabolism = new_speed * new_speed;

        return Ant{
            .x = GRID_SIZE / 2,
            .y = GRID_SIZE / 2,
            .alive = true,
            .energy = 100.0,
            .gene_speed = new_speed,
            .gene_metabolism = new_metabolism,
            .generation = parent.generation + 1,
            .score = 0,
        };
    }

    pub fn updateBestAnt(self: *World) void {
        var max_score: usize = 0;
        var best_idx: usize = 0;

        for (self.ants, 0..) |ant, i| {
            if (ant.alive and ant.score > max_score) {
                max_score = ant.score;
                best_idx = i;
            }
        }
        if (max_score > 0) {
            self.best_ant_idx = best_idx;
        }
    }

    pub fn stepAnt(self: *World, ant_idx: usize, action: usize) struct { f32, bool } {
        var ant = &self.ants[ant_idx];
        if (!ant.alive) return .{ 0.0, true };

        const random = self.prng.random();

        var moves_to_make: usize = 0;
        var chance = ant.gene_speed;

        while (chance > 0) {
            if (chance >= 1.0) {
                moves_to_make += 1;
                chance -= 1.0;
            } else {
                if (random.float(f32) < chance) moves_to_make += 1;
                chance = 0;
            }
        }

        if (moves_to_make == 0) {
            ant.energy -= 0.1;
            if (ant.energy <= 0) return .{ -10.0, true };
            return .{ 0.0, false };
        }

        var total_reward: f32 = 0.0;

        for (0..moves_to_make) |_| {
            if (ant.energy <= 0) break;

            ant.energy -= (0.2 * ant.gene_metabolism);

            const old_dist_x = if (ant.x > self.food_x) ant.x - self.food_x else self.food_x - ant.x;
            const old_dist_y = if (ant.y > self.food_y) ant.y - self.food_y else self.food_y - ant.y;
            const old_dist = old_dist_x + old_dist_y;

            var new_x = ant.x;
            var new_y = ant.y;
            if (action == 0) new_y -= 1;
            if (action == 1) new_y += 1;
            if (action == 2) new_x -= 1;
            if (action == 3) new_x += 1;

            if (self.grid[new_y][new_x] == TILE_WALL or self.isOccupied(new_x, new_y)) {
                ant.energy -= 1.0;
                total_reward -= 0.5;
                break;
            }

            ant.x = new_x;
            ant.y = new_y;
            self.pheromones[new_y][new_x] = 1.0;

            const new_dist = if (new_x > self.food_x) new_x - self.food_x else self.food_x - new_x;
            const total_new_dist = new_dist + (if (new_y > self.food_y) new_y - self.food_y else self.food_y - new_y);

            const scent = self.getScent(new_x, new_y);
            if (total_new_dist < old_dist) total_reward += 1.5 + scent;
            if (total_new_dist > old_dist) total_reward -= 1.0;

            if (new_x == self.food_x and new_y == self.food_y) {
                self.respawnFood();
                ant.energy = 100.0;
                ant.score += 1;
                total_reward += 50.0;

                self.updateBestAnt();
            }
        }

        if (ant.energy <= 0) {
            self.ants[ant_idx] = self.spawnChild(self.best_ant_idx);
            return .{ -10.0, true };
        }

        return .{ total_reward, false };
    }

    pub fn evaporatePheromones(self: *World) void {
        for (0..GRID_SIZE) |y| {
            for (0..GRID_SIZE) |x| {
                if (self.pheromones[y][x] > 0) self.pheromones[y][x] *= 0.995;
            }
        }
    }

    pub fn getAntObservation(self: *World, allocator: Allocator, ant_idx: usize) ![]f32 {
        var obs = try allocator.alloc(f32, 15);
        const ant = self.ants[ant_idx];
        var idx: usize = 0;

        const ax = @as(i32, @intCast(ant.x));
        const ay = @as(i32, @intCast(ant.y));

        var dy: i32 = -1;
        while (dy <= 1) : (dy += 1) {
            var dx: i32 = -1;
            while (dx <= 1) : (dx += 1) {
                const py = @as(usize, @intCast(ay + dy));
                const px = @as(usize, @intCast(ax + dx));
                var val: f32 = 0.0;
                if (self.grid[py][px] == TILE_WALL) {
                    val = -1.0;
                } else if (px == self.food_x and py == self.food_y) {
                    val = 1.0;
                } else if (self.isOccupied(px, py) and (dx != 0 or dy != 0)) {
                    val = -0.5;
                }
                obs[idx] = val;
                idx += 1;
            }
        }
        obs[9] = self.getScent(ant.x, ant.y);
        obs[10] = self.pheromones[ant.y][ant.x];
        obs[11] = if (self.food_y < ant.y) 1.0 else 0.0;
        obs[12] = if (self.food_y > ant.y) 1.0 else 0.0;
        obs[13] = if (self.food_x < ant.x) 1.0 else 0.0;
        obs[14] = if (self.food_x > ant.x) 1.0 else 0.0;
        return obs;
    }
};
