const std = @import("std");

const Input = []Machine;
const Machine = struct {
    a: @Vector(2, u32),
    b: @Vector(2, u32),
    prize: @Vector(2, u32),
};

const test_input =
    \\Button A: X+94, Y+34
    \\Button B: X+22, Y+67
    \\Prize: X=8400, Y=5400
    \\
    \\Button A: X+26, Y+66
    \\Button B: X+67, Y+21
    \\Prize: X=12748, Y=12176
    \\
    \\Button A: X+17, Y+86
    \\Button B: X+84, Y+37
    \\Prize: X=7870, Y=6450
    \\
    \\Button A: X+69, Y+23
    \\Button B: X+27, Y+71
    \\Prize: X=18641, Y=10279
    \\
;

pub fn parse(allocator: std.mem.Allocator, data: []const u8) !Input {
    var it = std.mem.splitSequence(u8, data, "\n\n");
    var machines = std.ArrayList(Machine).init(allocator);
    while (it.next()) |m| {
        var itt = std.mem.splitSequence(u8, m, ": ");
        _ = itt.next() orelse return error.InvalidInput;
        const a = itt.next() orelse return error.InvalidInput;
        const b = itt.next() orelse return error.InvalidInput;
        const prize = itt.next() orelse return error.InvalidInput;
        if (itt.next() != null) return error.InvalidInput;
        var machine: Machine = undefined;
        inline for (.{ a, b, prize }, .{ "a", "b", "prize" }) |str, name| {
            const comma = std.mem.indexOf(u8, str, ", ") orelse return error.InvalidInput;
            const end = std.mem.indexOfScalar(u8, str, '\n') orelse str.len;
            const x = try std.fmt.parseInt(u32, str[2..comma], 10);
            const y = try std.fmt.parseInt(u32, str[comma + 4 .. end], 10);
            @field(machine, name) = .{ x, y };
        }
        try machines.append(machine);
    }
    return machines.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(480, output);
}

pub fn part1(_: std.mem.Allocator, input: Input) !u32 {
    var total_tokens: u32 = 0;
    for (input) |machine| {
        var fewest: u32 = std.math.maxInt(u32);
        for (0..101) |a_presses| for (0..101) |b_presses| {
            const tokens: u32 = @intCast(a_presses * 3 + b_presses);
            if (tokens >= fewest) continue;
            const a_presses_vec: @Vector(2, u32) = @splat(@intCast(a_presses));
            const b_presses_vec: @Vector(2, u32) = @splat(@intCast(b_presses));
            if (@reduce(
                .And,
                machine.a * a_presses_vec + machine.b * b_presses_vec == machine.prize,
            )) {
                // std.debug.print("can win {} with A={}, B={}, tokens={}\n", .{ machine, a_presses, b_presses, tokens });
                fewest = tokens;
            }
        };
        if (fewest != std.math.maxInt(@TypeOf(fewest)))
            total_tokens += fewest;
    }
    return total_tokens;
}

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(875318608908, output);
}

test "solveDiophantine" {
    inline for (.{
        .{ .a = 69, .b = 42, .c = 100071, .solutions = 104 },
        .{ .a = 2, .b = 4, .c = 7, .solutions = 0 },
    }) |t| {
        var it = solveDiophantine(t.a, t.b, t.c);
        var solutions: usize = 0;
        while (it.next()) |sol| {
            // std.debug.print("{} * {} + {} * {} = {}\n", .{ t.a, sol[0], t.b, sol[1], t.c });
            try std.testing.expectEqual(t.c, t.a * sol[0] + t.b * sol[1]);
            try std.testing.expect(sol[0] >= 0);
            try std.testing.expect(sol[1] >= 0);
            solutions += 1;
        }
        try std.testing.expectEqual(t.solutions, solutions);
    }
}

const DiophantineIterator = struct {
    curr: @Vector(2, i64),
    diff: @Vector(2, i64),

    pub fn next(self: *@This()) ?@Vector(2, i64) {
        if (self.curr[1] < 0) return null;
        defer self.curr += self.diff;
        return self.curr;
    }
};

fn solveDiophantine(a: i64, b: i64, c: i64) DiophantineIterator {
    // std.debug.print("solveDiophantine({}, {}, {})\n", .{ a, b, c });
    const Vec = @Vector(2, i64);
    const cast = struct {
        fn f(x: u64) i64 {
            return @intCast(x);
        }
    }.f;

    const d = std.math.gcd(@abs(a), @abs(b));
    if (@rem(c, cast(d)) != 0) return .{ .curr = .{ undefined, -1 }, .diff = undefined };
    const u = @divExact(@abs(a), d);
    const v = @divExact(@abs(b), d);
    var x = u;
    var y = v;
    var sub_x: Vec = .{ 1, 0 };
    var sub_y: Vec = .{ 0, 1 };
    if (x < y) {
        std.mem.swap(u64, &x, &y);
        std.mem.swap(Vec, &sub_x, &sub_y);
    }
    while (y > 1) {
        const rem = @rem(x, y);
        const quot = @divTrunc(x, y);
        const sub: Vec = sub_x - @as(Vec, @splat(cast(quot))) * sub_y;
        // std.debug.print("{} = {} * {} + {}. {} = {} * {} + {} * {}\n", .{ x, y, quot, rem, rem, @divExact(a, d), sub[0], @divExact(b, d), sub[1] });

        x = y;
        y = rem;
        sub_x = sub_y;
        sub_y = sub;
    }
    std.debug.assert(y == 1);

    const one = sub_y;
    var solution = one * @as(Vec, @splat(cast(@divExact(@abs(c), d))));
    const k: i64 = -@divFloor(solution[0], cast(v));
    solution += .{ k * cast(v), -k * cast(u) };
    const s = std.math.sign;
    const sign_adjust: Vec = .{ s(a) * s(c), s(b) * s(c) };
    // std.debug.print("{} + k{}\n", .{ solution * sign_adjust, Vec{ cast(v), -cast(u) } * sign_adjust });
    return .{ .curr = solution * sign_adjust, .diff = Vec{ cast(v), -cast(u) } * sign_adjust };
}

pub fn part2(allocator: std.mem.Allocator, input: Input) !u64 {
    var total_tokens = std.atomic.Value(u64).init(0);
    var tpool: std.Thread.Pool = undefined;
    try tpool.init(.{ .allocator = allocator });
    defer tpool.deinit();
    var wg = std.Thread.WaitGroup{};
    for (input) |machine| {
        tpool.spawnWg(&wg, checkMachine, .{ machine, &total_tokens });
    }
    wg.wait();
    return total_tokens.load(.monotonic);
}

fn checkMachine(machine: Machine, total_tokens: *std.atomic.Value(u64)) void {
    // std.debug.print("machine[{} of {}] = {}\n", .{ i, input.len, machine });
    var fewest: u64 = std.math.maxInt(u64);

    var comb_sols = solveDiophantine(
        machine.a[0] + machine.a[1],
        machine.b[0] + machine.b[1],
        @as(i64, @intCast(machine.prize[0] + machine.prize[1])) + 20000000000000,
    );

    while (comb_sols.next()) |sol| {
        const solves_x = machine.a[0] * sol[0] + machine.b[0] * sol[1] == @as(i64, @intCast(machine.prize[0])) + 10000000000000;
        if (!solves_x) continue;
        const solves_y = machine.a[1] * sol[0] + machine.b[1] * sol[1] == @as(i64, @intCast(machine.prize[1])) + 10000000000000;
        std.debug.assert(solves_y);

        std.debug.print("found solution {}\n", .{sol});

        const tokens = @reduce(.Add, sol * @Vector(2, i64){ 3, 1 });
        fewest = @min(fewest, @as(u64, @intCast(tokens)));
    }

    if (fewest != std.math.maxInt(@TypeOf(fewest))) {
        _ = total_tokens.fetchAdd(fewest, .monotonic); // total_tokens += fewest;
    }
    // return total_tokens;
}