const std = @import("std");

const Vec = @Vector(2, i16);

const Input = struct {
    s: Vec,
    robots: []Robot,
};

const Robot = struct {
    pos: Vec,
    vel: Vec,
};

const test_input =
    \\p=0,4 v=3,-3
    \\p=6,3 v=-1,-3
    \\p=10,3 v=-1,2
    \\p=2,0 v=2,-1
    \\p=0,0 v=1,3
    \\p=3,0 v=-2,-2
    \\p=7,6 v=-1,-3
    \\p=3,0 v=-1,-2
    \\p=9,3 v=2,3
    \\p=7,3 v=-1,2
    \\p=2,4 v=2,-3
    \\p=9,5 v=-3,-3
    \\
;

pub fn parse(allocator: std.mem.Allocator, data: []const u8) !Input {
    const size: struct { w: i8, h: i8 } = if (data.ptr == test_input.ptr) .{ .w = 11, .h = 7 } else .{ .w = 101, .h = 103 };
    var lines = std.mem.splitScalar(u8, data, '\n');
    var robots = std.ArrayList(Robot).init(allocator);
    while (lines.next()) |line| {
        if (line.len == 0) break;
        var it = std.mem.tokenizeAny(u8, line, "pv=, ");
        const px = it.next() orelse return error.InvalidInput;
        const py = it.next() orelse return error.InvalidInput;
        const vx = it.next() orelse return error.InvalidInput;
        const vy = it.next() orelse return error.InvalidInput;
        // std.debug.print("{s} {s} {s} {s}\n", .{ px, py, vx, vy });

        try robots.append(.{
            .pos = .{
                try std.fmt.parseInt(i8, px, 10),
                try std.fmt.parseInt(i8, py, 10),
            },
            .vel = .{
                try std.fmt.parseInt(i8, vx, 10),
                try std.fmt.parseInt(i8, vy, 10),
            },
        });
    }
    return .{ .s = .{ size.w, size.h }, .robots = robots.items };
}

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));
    try std.testing.expectEqual(12, output);
}

pub fn part1(allocator: std.mem.Allocator, input: Input) !u32 {
    const robots = try allocator.alloc(Robot, input.robots.len);
    @memcpy(robots, input.robots);
    for (0..100) |_| {
        for (robots) |*robot| {
            robot.pos = @mod(robot.pos + robot.vel, input.s);
        }
    }
    var quadrants = @Vector(4, u32){ 0, 0, 0, 0 };
    for (robots) |robot| {
        if (@reduce(.Or, robot.pos == input.s / @as(Vec, @splat(2))))
            continue;

        const idx = @reduce(.Add, @select(u2, robot.pos < input.s / @as(Vec, @splat(2)), .{ 0, 0 }, .{ 2, 1 }));
        quadrants[idx] += 1;
    }
    return @reduce(.Mul, quadrants);
}

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});
    try std.testing.expectEqual(2, output);
}

pub fn part2(allocator: std.mem.Allocator, input: Input) !u32 {
    const robots = try allocator.alloc(Robot, input.robots.len);
    @memcpy(robots, input.robots);
    var i: usize = 0;
    const map = try allocator.alloc(u8, @intCast(@reduce(.Mul, input.s)));
    return while (true) : (i += 1) {
        @memset(map, 0);
        for (robots) |*robot| {
            robot.pos = @mod(robot.pos + robot.vel, input.s);
            map[@intCast(robot.pos[1] * input.s[0] + robot.pos[0])] +|= 1;
        }
        var buf: [(101 + 1) * 103]u8 = undefined;
        var fbs = std.io.fixedBufferStream(&buf);
        var out = fbs.writer();
        for (0..@intCast(input.s[1])) |y| {
            for (0..@intCast(input.s[0])) |x| {
                try out.print("{c}", .{switch (map[y * @as(u32, @intCast(input.s[0])) + x]) {
                    0 => '.',
                    1...9 => |c| c + '0',
                    10...10 + 26 - 1 => |c| c - 10 + 'a',
                    10 + 26...10 + 26 + 26 - 1 => |c| c - 10 + 'A',
                    else => '+',
                }});
            }
            try out.print("\n", .{});
        }
        std.debug.print("{s}", .{fbs.getWritten()});
        try std.io.getStdIn().reader().skipUntilDelimiterOrEof('\n');
        if (false) break i;
    };
}