const std = @import("std");

const PATH = "input/day17.txt";

const Str = []const u8;

const Coord3 = @Vector(3, isize);
const CubeHash3 = std.AutoHashMap(Coord3, void);

// Buggy with i2, see: https://github.com/ziglang/zig/issues/11611
const Dir3 = @Vector(3, isize);

const Grid3 = struct {
    rows: isize,
    cols: isize,
    cubes: CubeHash3,

    const directions: [3 * 3 * 3 - 1]Dir3 = blk: {
        var ret: [3 * 3 * 3 - 1]Dir3 = undefined;

        var idx: usize = 0;
        for ([_]isize{ -1, 0, 1 }) |x| {
            for ([_]isize{ -1, 0, 1 }) |y| {
                for ([_]isize{ -1, 0, 1 }) |z| {
                    if (x == 0 and y == 0 and z == 0) continue;
                    ret[idx] = Dir3{ x, y, z };
                    idx += 1;
                }
            }
        }

        break :blk ret;
    };

    fn countNeighbor(self: @This(), pos: Coord3) u5 {
        var sum: u5 = 0;
        for (directions) |d| {
            const diffpos = pos + d;
            if (diffpos[0] < 0 or diffpos[1] < 0 or diffpos[2] < 0) continue;

            if (self.cubes.contains(diffpos)) sum += 1;
        }

        std.debug.assert(sum < 3 * 3 * 3);

        return sum;
    }
};

const Coord4 = @Vector(4, isize);
const CubeHash4 = std.AutoHashMap(Coord4, void);

// Buggy with i2, see: https://github.com/ziglang/zig/issues/11611
const Dir4 = @Vector(4, isize);

const Grid4 = struct {
    rows: isize,
    cols: isize,
    cubes: CubeHash4,

    const directions: [3 * 3 * 3 * 3 - 1]Dir4 = blk: {
        var ret: [3 * 3 * 3 * 3 - 1]Dir4 = undefined;

        var idx: usize = 0;
        for ([_]isize{ -1, 0, 1 }) |x| {
            for ([_]isize{ -1, 0, 1 }) |y| {
                for ([_]isize{ -1, 0, 1 }) |z| {
                    for ([_]isize{ -1, 0, 1 }) |w| {
                        if (x == 0 and y == 0 and z == 0 and w == 0) continue;
                        ret[idx] = Dir4{ x, y, z, w };
                        idx += 1;
                    }
                }
            }
        }

        break :blk ret;
    };

    fn countNeighbor(self: @This(), pos: Coord4) u6 {
        var sum: u6 = 0;
        for (directions) |d| {
            const diffpos = pos + d;
            if (diffpos[0] < 0 or diffpos[1] < 0 or diffpos[2] < 0 or diffpos[3] < 0) continue;

            if (self.cubes.contains(diffpos)) sum += 1;
        }

        std.debug.assert(sum < 3 * 3 * 3 * 3);

        return sum;
    }
};

pub fn first(allocator: ?std.mem.Allocator) anyerror!usize {
    var grid = try parseInput3(allocator.?, @embedFile(PATH));
    // var grid = try parseInput3(allocator.?, test_input);
    defer grid.cubes.deinit();

    var cycle: usize = 1;
    while (cycle <= 6) : (cycle += 1) {
        var next_grid = Grid3{
            .rows = grid.rows + 2,
            .cols = grid.cols + 2,
            .cubes = CubeHash3.init(allocator.?),
        };

        var row: isize = 0;
        while (row < grid.rows + 2) : (row += 1) {
            var col: isize = 0;
            while (col < grid.cols + 2) : (col += 1) {
                var depth: isize = 0;
                while (depth <= cycle * 2) : (depth += 1) {
                    const nbr = grid.countNeighbor(.{ row - 1, col - 1, depth - 1 });

                    // if active
                    if (grid.cubes.contains(.{ row - 1, col - 1, depth - 1 })) {
                        if (nbr == 2 or nbr == 3) try next_grid.cubes.put(.{ row, col, depth }, {});
                    } else {
                        if (nbr == 3) try next_grid.cubes.put(.{ row, col, depth }, {});
                    }
                }
            }
        }

        grid.cubes.deinit();
        grid = next_grid;
    }

    return grid.cubes.count();
}

pub fn second(allocator: ?std.mem.Allocator) anyerror!usize {
    var grid = try parseInput4(allocator.?, @embedFile(PATH));
    // var grid = try parseInput4(allocator.?, test_input);
    defer grid.cubes.deinit();

    var cycle: usize = 1;
    while (cycle <= 6) : (cycle += 1) {
        var next_grid = Grid4{
            .rows = grid.rows + 2,
            .cols = grid.cols + 2,
            .cubes = CubeHash4.init(allocator.?),
        };

        var row: isize = 0;
        while (row < grid.rows + 2) : (row += 1) {
            var col: isize = 0;
            while (col < grid.cols + 2) : (col += 1) {
                var depth: isize = 0;
                while (depth <= cycle * 2) : (depth += 1) {
                    var hyper: isize = 0;
                    while (hyper <= cycle * 2) : (hyper += 1) {
                        const nbr = grid.countNeighbor(.{ row - 1, col - 1, depth - 1, hyper - 1 });

                        // if active
                        if (grid.cubes.contains(.{ row - 1, col - 1, depth - 1, hyper - 1 })) {
                            if (nbr == 2 or nbr == 3) try next_grid.cubes.put(.{ row, col, depth, hyper }, {});
                        } else {
                            if (nbr == 3) try next_grid.cubes.put(.{ row, col, depth, hyper }, {});
                        }
                    }
                }
            }
        }

        grid.cubes.deinit();
        grid = next_grid;
    }

    return grid.cubes.count();
}

fn parseInput3(allocator: std.mem.Allocator, input: Str) !Grid3 {
    var lines = std.mem.tokenize(u8, input, "\n");

    var grid = Grid3{
        .rows = 0,
        .cols = 0,
        .cubes = CubeHash3.init(allocator),
    };

    var row: isize = 0;
    while (lines.next()) |line| : (row += 1) {
        grid.cols = @intCast(isize, line.len);
        for (line) |ch, col| {
            if (ch == '#') try grid.cubes.put(.{ row, @intCast(isize, col), 0 }, {});
        }
    }

    grid.rows = row;

    return grid;
}

fn parseInput4(allocator: std.mem.Allocator, input: Str) !Grid4 {
    var lines = std.mem.tokenize(u8, input, "\n");

    var grid = Grid4{
        .rows = 0,
        .cols = 0,
        .cubes = CubeHash4.init(allocator),
    };

    var row: isize = 0;
    while (lines.next()) |line| : (row += 1) {
        grid.cols = @intCast(isize, line.len);
        for (line) |ch, col| {
            if (ch == '#') try grid.cubes.put(.{ row, @intCast(isize, col), 0, 0 }, {});
        }
    }

    grid.rows = row;

    return grid;
}

test "day17a" {
    try std.testing.expectEqual(@as(usize, 223), try first(std.testing.allocator));
}

test "day17b" {
    try std.testing.expectEqual(@as(usize, 1884), try second(std.testing.allocator));
}

const test_input =
    \\.#.
    \\..#
    \\###
;