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 },
        .{ .a = 2, .b = 4, .c = 6, .solutions = 2 },
    }) |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: u64, b: u64, c: u64) DiophantineIterator {
    const Vec = @Vector(2, i64);
    const cast = struct {
        fn f(x: u64) i64 {
            return @intCast(x);
        }
    }.f;

    const d = std.math.gcd(a, b);
    if (c % d != 0) return .{ .curr = .{ undefined, -1 }, .diff = undefined };
    const u = @divExact(a, d);
    const v = @divExact(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 = x % y;
        const quot = 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(c, d))));
    const k: i64 = -@divFloor(solution[0], cast(v));
    solution += .{ k * cast(v), -k * cast(u) };
    return .{ .curr = solution, .diff = .{ cast(v), -cast(u) } };
}

pub fn part2(_: std.mem.Allocator, input: Input) !u64 {
    var total_tokens: u64 = 0;
    for (input) |machine| {
        std.debug.print("machine = {}\n", .{machine});
        var fewest: u64 = std.math.maxInt(u64);
        var x_sols = solveDiophantine(machine.a[0], machine.b[0], @as(u64, @intCast(machine.prize[0])) + 10000000000000);
        const y_sols = solveDiophantine(machine.a[1], machine.b[1], @as(u64, @intCast(machine.prize[1])) + 10000000000000);
        while (x_sols.next()) |sol_x| {
            const ks = (sol_x - y_sols.curr) / y_sols.diff;
            if (@reduce(.And, @rem(sol_x - y_sols.curr, y_sols.diff) == @Vector(2, i64){ 0, 0 }) and ks[0] == ks[1]) {
                const tokens = @reduce(.Add, sol_x * @Vector(2, i64){ 3, 1 });
                fewest = @min(fewest, @as(u64, @intCast(tokens)));
                std.debug.print("    possible solution {} for {} tokens\n", .{ sol_x, tokens });
            }
        }
        if (fewest != std.math.maxInt(@TypeOf(fewest)))
            total_tokens += fewest;
    }
    return total_tokens;
}