Day 15: Warehouse Woes

/* Advent of Code 2024
 * Day 15: Warehouse Woes
 */

#import "Basic";
#import "File";
#import "String";

Tile :: enum {
    EMPTY;
    WALL;
    BOX;
    BOX_LEFT;
    BOX_RIGHT;
}

Vector2_Int :: struct {
    x, y: int;
}

array_free_2d :: (array: [..][..] $T) {
    for array {
        array_free(it);
    }

    array_free(array);
}

is_box :: inline (tile: Tile) -> bool {
    return tile == .BOX || tile == .BOX_LEFT || tile == .BOX_RIGHT;
}

// Moves a box in a given direction. Does not perform any validity checks; that
// is all handled by push_box(). The inputs to this procedure are assumed to be
// valid.
move_box :: inline (
    map: [..][..] Tile,
    position: Vector2_Int,
    direction: Vector2_Int
) {
    tile := map[position.y][position.x];
    map[position.y][position.x] = .EMPTY;
    map[position.y + direction.y][position.x + direction.x] = tile;
}

// Attempts to push a box in a given direction, along with all boxes in its way.
// Returns whether the push was successful.
//
// If should_push is false, the procedure will return whether the push is
// possible, but will not actually move the boxes. This is necessary for part 2,
// where we have to check whether multiple pushes are possible before pushing
// some boxes (and thus don't want a possible push to still physically happen if
// the other push isn't possible).
push_box :: (
    map: [..][..] Tile,
    position: Vector2_Int,
    direction: Vector2_Int,
    should_push := true
) -> bool {
    tile := map[position.y][position.x];
    if direction.x == 0 && (tile == .BOX_LEFT || tile == .BOX_RIGHT) {
        other_side := ifx tile == .BOX_LEFT then 1 else -1;
        next_tile_1 := map[position.y + direction.y][position.x];
        next_tile_2 := map[position.y + direction.y][position.x + other_side];

        aligned_with_next_box := (
            tile == .BOX_LEFT && next_tile_1 == .BOX_LEFT
            || tile == .BOX_RIGHT && next_tile_2 == .BOX_LEFT
        );

        if
            aligned_with_next_box && push_box(map, .{
                position.x + ifx tile == .BOX_RIGHT then -1 else 0,
                position.y + direction.y
            }, direction, false)
            || (
                next_tile_1 == .EMPTY || is_box(next_tile_1) && push_box(map, .{
                    position.x,
                    position.y + direction.y
                }, direction, false)
            ) && (
                next_tile_2 == .EMPTY || is_box(next_tile_2) && push_box(map, .{
                    position.x + other_side,
                    position.y + direction.y
                }, direction, false)
            )
        {
            if should_push {
                // I'm not a huge fan of calling push_box() a second time here,
                // but without writing a separate can_push_box() procedure for
                // above (which would heavily complicate things), I don't think
                // there's a way around it. Either way, there are no performance
                // issues in the current version of the program, so it luckily
                // doesn't seem like much of a problem.
                if aligned_with_next_box {
                    push_box(map, .{
                        position.x + ifx tile == .BOX_RIGHT then -1 else 0,
                        position.y + direction.y
                    }, direction);
                } else {
                    if next_tile_1 != .EMPTY {
                        push_box(map, .{
                            position.x,
                            position.y + direction.y
                        }, direction);
                    }

                    if next_tile_2 != .EMPTY {
                        push_box(map, .{
                            position.x + other_side,
                            position.y + direction.y
                        }, direction);
                    }
                }

                move_box(map, position, direction);
                move_box(map, .{position.x + other_side, position.y}, direction);
            }

            return true;
        }
    } else {
        next_tile := map[position.y + direction.y][position.x + direction.x];
        if next_tile == .EMPTY || is_box(next_tile) && push_box(map, .{
            position.x + direction.x,
            position.y + direction.y
        }, direction, should_push) {
            if should_push {
                move_box(map, position, direction);
            }

            return true;
        }
    }

    return false;
}

// Performs a series of instructions on a robot in a warehouse, then calculates
// the sum of the GPS coordinates of the boxes in the warehouse after the robot
// has finished moving.
//
// We don't have to bother with bounds checking, because the warehouse is
// surrounded on all sides with walls.
get_gps_sum :: (
    map: [][..] Tile,
    robot: Vector2_Int,
    instructions: [] Vector2_Int
) -> int {
    // The map and robot parameters are immutable (since we don't want to modify
    // them outside this function), so since we need to be able to modify them
    // in this function, we thus need to make copies of them.
    mutable_map: [..][..] Tile;
    defer array_free_2d(mutable_map);

    for map {
        row: [..] Tile;
        array_copy(*row, it);
        array_add(*mutable_map, row);
    }

    mutable_robot := Vector2_Int.{robot.x, robot.y};

    for instructions {
        new_position := Vector2_Int.{
            mutable_robot.x + it.x,
            mutable_robot.y + it.y
        };

        new_tile := mutable_map[new_position.y][new_position.x];
        if new_tile == .EMPTY || is_box(new_tile) && push_box(
            mutable_map,
            new_position,
            it
        ) {
            mutable_robot = new_position;
        }
    }

    gps_sum := 0;
    for row, y: mutable_map {
        for tile, x: row {
            if tile == .BOX || tile == .BOX_LEFT {
                gps_sum += x + y * 100;
            }
        }
    }

    return gps_sum;
}

main :: () {
    input, success := read_entire_file("inputs/day_15.txt");
    assert(success);

    groups := split(input, "\n\n");
    defer array_free(groups);

    map_lines := split(groups[0], "\n");
    defer array_free(map_lines);

    map: [..][..] Tile;
    defer array_free_2d(map);

    robot: Vector2_Int;
    for line, y: map_lines {
        row: [..] Tile;
        for char, x: line {
            if char == {
                case #char "@"; robot = .{x, y}; #through;
                case #char "."; array_add(*row, .EMPTY);
                case #char "#"; array_add(*row, .WALL);
                case #char "O"; array_add(*row, .BOX);
            }
        }

        array_add(*map, row);
    }

    instructions: [..] Vector2_Int;
    defer array_free(instructions);

    for groups[1] {
        if it == {
            case #char "<"; array_add(*instructions, .{-1, 0});
            case #char ">"; array_add(*instructions, .{1, 0});
            case #char "^"; array_add(*instructions, .{0, -1});
            case #char "v"; array_add(*instructions, .{0, 1});
        }
    }

    print("Part 1: %\n", get_gps_sum(map, robot, instructions));

    wide_map: [..][..] Tile;
    defer array_free_2d(wide_map);

    for map {
        row: [..] Tile;
        for it {
            if it == .BOX {
                array_add(*row, .BOX_LEFT);
                array_add(*row, .BOX_RIGHT);
            } else {
                array_add(*row, it);
                array_add(*row, it);
            }
        }

        array_add(*wide_map, row);
    }

    robot.x *= 2;
    print("Part 2: %\n", get_gps_sum(wide_map, robot, instructions));
}

// Attempts to push a box in a given direction, along with all boxes in its way.
// Returns whether the push was successful.
get_gps_sum :: (
    map: [][..] Tile,
    width: int,
    height: int,
    robot: Vector2_Int,
    instructions: [] Vector2_Int
) -> int {
    // The map and robot parameters are not pointers (since we don't want to
    // modify them outside this function), so they are immutable. Since we need
    // to be able to modify them in this function, we thus need to make copies
    // of them.
    mutable_map: [..][..] Tile;
    defer array_free_2d(mutable_map);

    for map {
        row: [..] Tile;
        array_copy(*row, it);
        array_add(*mutable_map, row);
    }

    mutable_robot := Vector2_Int.{robot.x, robot.y};

    for instructions {
        new_robot := Vector2_Int.{
            mutable_robot.x + it.x,
            mutable_robot.y + it.y
        };

        new_tile := mutable_map[new_robot.y][new_robot.x];
        if new_tile == .EMPTY {
            mutable_robot = new_robot;
        } else if
            new_tile == .BOX || new_tile == .BOX_LEFT || new_tile == .BOX_RIGHT
        {
            position := Vector2_Int.{new_robot.x + it.x, new_robot.y + it.y};
            while mutable_map[position.y][position.x] == .BOX {
                position.x += it.x;
                position.y += it.y;
            }

            if mutable_map[position.y][position.x] == .EMPTY {
                while true {
                    position.x -= it.x;
                    position.y -= it.y;

                    if
                        position.x == mutable_robot.x
                        && position.y == mutable_robot.y
                    {
                        break;
                    }

                    mutable_map[position.y][position.x] = .EMPTY;
                    mutable_map[position.y + it.y][position.x + it.x] = .BOX;
                }

                mutable_robot = new_robot;
            }
        }
    }

    gps_sum := 0;
    for row, y: mutable_map {
        for tile, x: row {
            if tile == .BOX || tile == .BOX_LEFT {
                gps_sum += x + y * 100;
            }
        }
    }

    return gps_sum;
}

with open('input/15.txt', 'r') as f:
    grid_text, moves = f.read().split("\n\n")

grid = {}
for y, row in enumerate(grid_text.split("\n")):
    for x, cell in enumerate(row):
        location = x + y * 1j
        grid[location] = cell
        if cell == "@":
            roboot_location = location

movemap = {
    ">": 1,
    "<": -1,
    "^": -1j,
    "v": 1j
}

def perform_move(grid, roboot_location, move):
    new_location = roboot_location + movemap[move]
    boxes_to_push = []
    while grid[new_location] == "O":
        new_location = new_location + movemap[move]
        boxes_to_push.append(new_location)

    if grid[new_location] == "#":
        return roboot_location
    if grid[new_location] == ".":
        for box in boxes_to_push:
            grid[box] = "O"
        grid[roboot_location] = "."
        roboot_location = roboot_location + movemap[move]
        grid[roboot_location] = "@"
        return roboot_location

def get_coord(position):
    return int(position.imag) * 100 + int(position.real)

def print_grid(grid, grid_text):
    for y in range(grid_text.count("\n") + 1):
        for x in range(len(grid_text.split("\n")[0])):
            print(grid.get(x + y * 1j, " "), end="")
        print()

for move in "".join(moves.split()):
    roboot_location = perform_move(grid, roboot_location, move)

box_coord_sum = sum([get_coord(position) for position, cell in grid.items() if cell == "O"])
print("Part 1:", box_coord_sum)

widened_grid_text = grid_text.replace("#","##").replace("O", "[]").replace(".", "..").replace("@", "@.")
grid = {}
for y, row in enumerate(widened_grid_text.split("\n")):
    for x, cell in enumerate(row):
        location = x + y * 1j
        grid[location] = cell
        if cell == "@":
            roboot_location = location


def perform_move(grid, roboot_location, move):
    if grid[roboot_location + move] == "#":
        return roboot_location
    
    locations_to_push_from = [roboot_location]
    boxes_to_push = []

    while any([grid[location + move] in "[]" for location in locations_to_push_from]):
        if any([grid[location + move] == "#" for location in locations_to_push_from]):
            return roboot_location
        new_locations_to_push_from = []
        for location in locations_to_push_from:
            boxes = []
            if grid[location + move] == "[":
                boxes = [location + move]
                if move in [1j, -1j]:
                    boxes.append(location + move + 1)
            elif grid[location + move] == "]":
                boxes = [location + move]
                if move in [1j, -1j]:
                    boxes.append(location + move - 1)
            boxes_to_push.extend(boxes)
            new_locations_to_push_from.extend(boxes)
        locations_to_push_from = new_locations_to_push_from

    if any([grid[box+move] == "#" for box in boxes_to_push]):
        return roboot_location

    new_symbols = {}
    for box in boxes_to_push:
        symbol = grid[box]
        new_symbols[box + move] = symbol
    for box in boxes_to_push:
        grid[box] = "."
    for box, symbol in new_symbols.items():
        grid[box] = symbol

    grid[roboot_location] = "."
    roboot_location = roboot_location + move
    grid[roboot_location] = "@"
    return roboot_location

for move in "".join(moves.split()):
    roboot_location = perform_move(grid, roboot_location, movemap[move])

box_coord_sum = sum([get_coord(position) for position, cell in grid.items() if cell == "["])
print("Part 2:", box_coord_sum)

pub fn part2(input: &PuzzleInput) -> String {
    let mut warehouse = Grid2D::new(input.warehouse.width() * 2, input.warehouse.height(), '.');

    // Double the warehouse width
    input.warehouse.iter().for_each(|(coord, c)| {
        let x = coord.x * 2;
        let y = coord.y;

        if *c == 'O' {
            warehouse.set((x, y).into(), '[');
            warehouse.set((x + 1, y).into(), ']');
        } else if *c == '@' {
            warehouse.set((x, y).into(), '@');
            warehouse.set((x + 1, y).into(), '.');
        } else {
            warehouse.set((x, y).into(), *c);
            warehouse.set((x + 1, y).into(), *c);
        }
    });

    let result = run_robot(&PuzzleInput {
        warehouse,
        robot_moves: input.robot_moves.clone(),
    });

    println!("{}", result);

    // GPS coordinates
    let mut sum = 0;
    for (coord, c) in result.iter() {
        if c == &'[' {
            sum += 100 * coord.y + coord.x;
        }
    }

    sum.to_string()
}

pub fn can_push_box(grid: &mut Grid2D<char>, box_pos: Coordinate, dir: Direction) -> bool {
    if grid[box_pos] == '.' {
        return true;
    }

    if grid[box_pos] == '#' {
        return false;
    }

    match dir {
        Direction::Left | Direction::Right => {
            let next = box_pos + dir.into();
            return can_push_box(grid, next, dir);
        }

        Direction::Up | Direction::Down => {
            let next_pos1 = box_pos + dir.into();
            let next_pos2 = box_pos
                + dir.into()
                + if grid[box_pos] == '[' {
                    Direction::Right.into()
                } else {
                    Direction::Left.into()
                };

            return can_push_box(grid, next_pos1, dir) && can_push_box(grid, next_pos2, dir);
        }

        _ => false,
    }
}

pub fn push_box(grid: &mut Grid2D<char>, box_pos: Coordinate, dir: Direction) {
    if grid[box_pos] == '.' {
        return;
    }

    if grid[box_pos] == '#' {
        return;
    }

    if dir == Direction::Left || dir == Direction::Right {
        let mut cur = box_pos;
        let mut next;

        loop {
            next = cur + dir.into();

            if grid[next] == '.' {
                break;
            }

            cur = next;
        }

        while next != box_pos {
            grid.set(next, grid[cur]);
            grid.set(cur, '.');
            next = cur;
            cur = cur.neighbor(dir.opposite());
        }
    } else {
        let second_half_dir = if grid[box_pos] == '[' {
            Direction::Right
        } else {
            Direction::Left
        };

        let opposite_bracket = if grid[box_pos] == '[' { ']' } else { '[' };

        // Recursively push the boxes
        push_box(grid, box_pos + dir.into(), dir);
        push_box(grid, box_pos + dir.into() + second_half_dir.into(), dir);

        // Move the box itself
        grid.set(box_pos + dir.into(), grid[box_pos]);
        grid.set(
            box_pos + dir.into() + second_half_dir.into(),
            opposite_bracket,
        );
        grid.set(box_pos, '.');
        grid.set(box_pos + second_half_dir.into(), '.');
    }
}

pub fn run_robot(input: &PuzzleInput) -> Grid2D<char> {
    let mut grid = input.warehouse.clone();
    let mut robot_pos = grid.iter().find(|(_, c)| **c == '@').unwrap().0;

    for robot_move in input.robot_moves.iter() {
        let dir: Direction = (*robot_move).try_into().unwrap();

        if can_push_box(&mut grid, robot_pos + dir.into(), dir) {
            push_box(&mut grid, robot_pos + dir.into(), dir);
            grid.set(robot_pos, '.');
            grid.set(robot_pos + dir.into(), '@');
            robot_pos = robot_pos + dir.into();
        }
    }

    grid
}