const std = @import("std");

const Str = []const u8;

const path = "data/day20/input.txt";
const max_rounds = 50;
const IMAGE_SIZE = 100;

const Pixels = std.DynamicBitSet;

const bitset_size = 512;
const BitSet = std.StaticBitSet(bitset_size);

const Image = struct {
    IEA: BitSet,
    lit: Pixels,
    size: usize,
    flippy: bool = false,
};

pub fn main() !void {
    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();

    var timer = try std.time.Timer.start();
    const ret = try second(arena.allocator());
    const t = timer.lap() / 1000;

    try std.testing.expectEqual(@as(usize, 19156), ret);

    std.debug.print("Day 20b result: {d} \t\ttime: {d}us\n", .{ ret, t });
}

pub fn second(allocator: ?std.mem.Allocator) !usize {
    var image = try parseInput(allocator.?);
    defer image.lit.deinit();

    var round: usize = 1;
    while (round <= max_rounds) : (round += 1) {
        const next_size = IMAGE_SIZE + round * 2; // 7, 9, 11...

        var next = try Pixels.initEmpty(allocator.?, (next_size * next_size));

        var row: usize = 0;
        while (row < next_size) : (row += 1) {
            var col: usize = 0;
            while (col < next_size) : (col += 1) {
                if (getValue(&image, row, col)) next.set(row * next_size + col);
            }
        }

        image.lit.deinit();
        image.lit = next;
        image.size = next_size;

        if (image.IEA.isSet(0) and !image.IEA.isSet(bitset_size - 1)) {
            image.flippy = !image.flippy;
        }
    }

    return image.lit.count();
}

const directions = [9][2]isize{
    // order is important!
    .{ 1, 1 },
    .{ 1, 0 },
    .{ 1, -1 },
    .{ 0, 1 },
    .{ 0, 0 },
    .{ 0, -1 },
    .{ -1, 1 },
    .{ -1, 0 },
    .{ -1, -1 },
};

fn getValue(image: *Image, row: usize, col: usize) bool {
    var index: u9 = 0;

    for (directions) |d, idx| {
        const r = @intCast(isize, row) - 1 + d[0];
        const c = @intCast(isize, col) - 1 + d[1];

        if (r < 0 or c < 0) { // upper of left
            if (image.flippy) index |= @as(u9, 1) <<| idx;
            continue;
        }
        if (r >= image.size or c >= image.size) { // bottom or right
            if (image.flippy) index |= @as(u9, 1) <<| idx;
            continue;
        }

        const rr = @intCast(usize, r);
        const cc = @intCast(usize, c);

        if (image.lit.isSet(rr * image.size + cc)) index |= @as(u9, 1) <<| idx;
    }

    std.debug.assert(index < bitset_size);
    return image.IEA.isSet(index);
}

fn parseInput(a: std.mem.Allocator) !Image {
    const input = @embedFile(path);
    var lines = std.mem.split(u8, input, "\n");

    var i: Image = undefined;

    i.IEA = BitSet.initEmpty();

    for (lines.next().?) |item, idx| {
        if (item == '#') i.IEA.set(idx);
    }
    _ = lines.next(); // drop empty line

    i.lit = try Pixels.initEmpty(a, IMAGE_SIZE * IMAGE_SIZE);

    var row: usize = 0;
    while (lines.next()) |line| : (row += 1) {
        for (line) |ch, col| {
            if (ch == '#') i.lit.set(IMAGE_SIZE * row + col);
        }
    }

    i.size = row - 1;

    return i;
}

fn printImage(image: *Image) void {
    std.debug.print("\n", .{});
    var row: usize = 0;
    while (row < image.size) : (row += 1) {
        var col: usize = 0;
        while (col < image.size) : (col += 1) {
            if (image.lit.isSet(row * image.size + col)) {
                std.debug.print("#", .{});
            } else {
                std.debug.print(".", .{});
            }
        }
        std.debug.print("\n", .{});
    }
}

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