const std = @import("std");

const Pos = @Vector(2, isize);

const Input = struct {
    width: usize,
    height: usize,
    antenna_locations: std.AutoHashMap(u8, std.ArrayList(Pos)),
};

pub fn parse(allocator: std.mem.Allocator, data: []const u8) !Input {
    var lines = std.mem.split(u8, data, "\n");
    var width: usize = undefined;
    var y: usize = 0;
    var antenna_locations = std.AutoHashMap(u8, std.ArrayList(Pos)).init(allocator);
    while (lines.next()) |line| : (y += 1) {
        if (line.len == 0) break;
        for (0.., line) |x, c| {
            switch (c) {
                '0'...'9', 'a'...'z', 'A'...'Z' => {
                    const res = try antenna_locations.getOrPut(c);
                    if (!res.found_existing)
                        res.value_ptr.* = std.ArrayList(Pos).init(allocator);
                    try res.value_ptr.*.append(.{ @intCast(x), @intCast(y) });
                },
                else => {},
            }
        }
        width = line.len;
    }
    return .{
        .width = width,
        .height = y,
        .antenna_locations = antenna_locations,
    };
}

const test_input =
    \\............
    \\........0...
    \\.....0......
    \\.......0....
    \\....0.......
    \\......A.....
    \\............
    \\............
    \\........A...
    \\.........A..
    \\............
    \\............
    \\
;

test "part1" {
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();

    const output = try part1(arena.allocator(), try parse(arena.allocator(), test_input));
    std.debug.print("got {}\n", .{output});
    std.debug.assert(output == 14);
}

pub fn part1(allocator: std.mem.Allocator, input: Input) !u32 {
    var it = input.antenna_locations.valueIterator();
    var antinodes = std.AutoHashMap(Pos, void).init(allocator);
    while (it.next()) |locations| {
        for (locations.items) |a| for (locations.items) |b| {
            if (@reduce(.And, a == b)) continue;
            const pos = b + b - a;
            if (pos[0] < 0 or input.width <= pos[0] or
                pos[1] < 0 or input.height <= pos[1]) continue;
            try antinodes.put(pos, {});
        };
    }
    return antinodes.count();
}

test "part2" {
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();

    const output = try part2(arena.allocator(), try parse(arena.allocator(), test_input));
    std.debug.print("got {}\n", .{output});
    std.debug.assert(output == 34);
}

pub fn part2(allocator: std.mem.Allocator, input: Input) !u32 {
    var it = input.antenna_locations.valueIterator();
    var antinodes = std.AutoHashMap(Pos, void).init(allocator);
    while (it.next()) |locations| {
        for (locations.items) |a| for (locations.items) |b| {
            if (@reduce(.And, a == b)) continue;
            const delta = b - a;
            var pos = b;
            while (0 <= pos[0] and pos[0] < input.width and
                0 <= pos[1] and pos[1] < input.height) : (pos += delta)
            {
                try antinodes.put(pos, {});
            }
        };
    }
    return antinodes.count();
}