const SudokuSolver = struct {
    const countEql = std.simd.countElementsWithValue;
    const join = std.simd.join;
    // Each cell exists in memory three times.
    // This makes it easy to compare a guess against all relevant cells simultaneously.
    rows: [9]@Vector(9, u8),
    cols: [9]@Vector(9, u8),
    boxes: [9]@Vector(9, u8),

    fn init(board: [9][9]u8) !SudokuSolver {
        var res: SudokuSolver = .{
            .rows = undefined,
            .cols = undefined,
            .boxes = undefined,
        };
        for (board, 0..) |r, row| {
            for (r, 0..) |_, col| {
                const cell = board[row][col];
                if (cell > 9) return error.InvalidValue;
                res.setCell(row, col, cell);
            }
        }

        return res;
    }

    fn possibleMovesForCell(self: SudokuSolver, row: usize, col: usize) std.StaticBitSet(9) {
        var cell_guesses: std.StaticBitSet(9) = .initEmpty();
        for (1..10) |n| {
            // A number is a possible move iff it does not already exist in the current row, col, or box.
            if (countEql(
                join(self.rows[row], join(self.cols[col], self.boxes[boxOf(row, col)])),
                @intCast(n),
            ) == 0) {
                cell_guesses.set(n - 1);
            }
        }

        return cell_guesses;
    }

    const PossibleMoves = struct {
        row: usize,
        col: usize,
        moves: std.StaticBitSet(9),
    };

    fn possibleMovesLessThan(_: void, lhs: PossibleMoves, rhs: PossibleMoves) bool {
        return lhs.moves.count() < rhs.moves.count();
    }

    fn solve(self: *SudokuSolver) bool {
        if (self.isSolved()) return true;

        var most_constrained_moves: ?PossibleMoves = null;

        for (self.rows, 0..) |r, row| {
            for (@as([9]u8, r), 0..) |cell, col| {
                if (cell == 0) {
                    const current_possible_moves = self.possibleMovesForCell(row, col);
                    switch (current_possible_moves.count()) {
                        1 => {
                            self.setCell(
                                row,
                                col,
                                @intCast(current_possible_moves.findFirstSet().? + 1),
                            );
                            return self.solve();
                        },
                        0 => return false,
                        else => |c| if (most_constrained_moves == null or c < most_constrained_moves.?.moves.count()) {
                            most_constrained_moves = .{
                                .row = row,
                                .col = col,
                                .moves = current_possible_moves,
                            };
                        },
                    }
                }
            }
        }

        if (most_constrained_moves) |pm| {
            var moves_iter = pm.moves.iterator(.{});

            while (moves_iter.next()) |current_move| {
                const backup = self.*;

                self.setCell(pm.row, pm.col, @intCast(current_move + 1));

                if (self.solve()) {
                    return true;
                }

                self.* = backup;
            }
        }

        return false;
    }

    fn isSolved(self: SudokuSolver) bool {
        for (self.rows) |row| {
            if (countEql(row, 0) > 0) return false;
        }
        return true;
    }

    fn setCell(self: *SudokuSolver, row: usize, col: usize, value: u8) void {
        self.rows[row][col] = value;
        self.cols[col][row] = value;
        self.boxes[boxOf(row, col)][cellOfInBox(row, col)] = value;
    }

    fn boxOf(row: usize, col: usize) usize {
        // Due to integer division, this is not the same as just adding row.
        return (row / 3) * 3 + col / 3;
    }

    fn cellOfInBox(row: usize, col: usize) usize {
        return (row % 3) * 3 + col % 3;
    }

    fn display(self: SudokuSolver, writer: std.io.AnyWriter) !void {
        for (self.rows, 0..) |row, i| {
            // Can't loop over row directly?
            for (@as([9]u8, row), 0..) |cell, j| {
                try writer.print("{d} ", .{cell});
                if (j % 3 == 2 and j < 8) {
                    try writer.print("| ", .{});
                }
            }
            try writer.print("\n", .{});
            if (i % 3 == 2 and i < 8) {
                try writer.print("{s}\n", .{"-" ** 21});
            }
        }
    }
};

pub fn main() !void {
    var board: [9][9]u8 = undefined;

    const stdout = std.io.getStdOut().writer().any();
    const stdin = std.io.getStdIn().reader().any();

    for (0..9) |row| {
        for (0..9) |col| {
            const byte = try stdin.readByte();
            board[row][col] = try std.fmt.charToDigit(byte, 10);
        }
    }

    // Unreachable because we validate while parsing.
    var solver = SudokuSolver.init(board) catch unreachable;
    try solver.display(stdout);
    std.debug.print("\n", .{});

    if (solver.solve()) {
        try solver.display(stdout);
    } else {
        std.debug.print("Unable to solve board\n", .{});
    }
}

const std = @import("std");

test "boxOf" {
    const expectEql = std.testing.expectEqual;
    try expectEql(0, SudokuSolver.boxOf(0, 0));
    try expectEql(1, SudokuSolver.boxOf(0, 3));
    try expectEql(2, SudokuSolver.boxOf(2, 8));
    try expectEql(4, SudokuSolver.boxOf(4, 5));
}

test "cellOfInBox" {
    const expectEql = std.testing.expectEqual;
    try expectEql(0, SudokuSolver.cellOfInBox(0, 0));
    try expectEql(1, SudokuSolver.cellOfInBox(0, 1));
    try expectEql(2, SudokuSolver.cellOfInBox(0, 2));
    try expectEql(3, SudokuSolver.cellOfInBox(1, 0));

    try expectEql(0, SudokuSolver.cellOfInBox(3, 3));
    try expectEql(0, SudokuSolver.cellOfInBox(6, 6));
}

test "get available moves for easy board" {
    // Taken from https://ziggit.dev/t/i-wrote-a-simple-sudoku-solver/9924/14?u=zambyte
    if (false) {
        const board: [9][9]u8 = .{
            .{ 4, 0, 3, 5, 0, 0, 0, 9, 6 },
            .{ 5, 0, 7, 8, 0, 0, 3, 0, 0 },
            .{ 0, 8, 0, 0, 0, 0, 0, 5, 4 },
            .{ 0, 0, 0, 7, 0, 6, 0, 2, 3 },
            .{ 0, 0, 0, 0, 2, 0, 0, 0, 0 },
            .{ 9, 5, 0, 3, 0, 4, 0, 0, 0 },
            .{ 7, 9, 0, 0, 0, 0, 0, 6, 0 },
            .{ 0, 0, 4, 0, 0, 7, 5, 0, 8 },
            .{ 8, 3, 0, 0, 0, 5, 4, 0, 9 },
        };

        const solver: SudokuSolver = SudokuSolver.init(.{ .board = board }) catch unreachable;

        for (board, 0..) |row, row_i| {
            for (row, 0..) |cell, col_i| {
                if (cell == 0) {
                    const possible_moves = solver.possibleMovesForCell(row_i, col_i);
                    if (possible_moves.count() == 1) {
                        std.debug.print("row: {}, col: {}, value: {}\n", .{ row_i, col_i, possible_moves.findFirstSet().? });
                    }
                }
            }
        }
    }
    return error.SkipZigTest;
}

test "Solve easy board" {
    // Taken from https://ziggit.dev/t/i-wrote-a-simple-sudoku-solver/9924/14?u=zambyte
    const board: [9][9]u8 = .{
        .{ 4, 0, 3, 5, 0, 0, 0, 9, 6 },
        .{ 5, 0, 7, 8, 0, 0, 3, 0, 0 },
        .{ 0, 8, 0, 0, 0, 0, 0, 5, 4 },
        .{ 0, 0, 0, 7, 0, 6, 0, 2, 3 },
        .{ 0, 0, 0, 0, 2, 0, 0, 0, 0 },
        .{ 9, 5, 0, 3, 0, 4, 0, 0, 0 },
        .{ 7, 9, 0, 0, 0, 0, 0, 6, 0 },
        .{ 0, 0, 4, 0, 0, 7, 5, 0, 8 },
        .{ 8, 3, 0, 0, 0, 5, 4, 0, 9 },
    };

    var solver: SudokuSolver = SudokuSolver.init(board) catch unreachable;

    const stderr = std.io.getStdErr().writer().any();
    std.debug.print("input:\n", .{});
    try solver.display(stderr);
    std.debug.print("output:\n", .{});

    _ = solver.solve();
    try solver.display(stderr);
}