from pathlib import Path


def parse_data(file_name):
    file_path = Path(__file__).resolve().parent / f"{file_name}.txt"
    wall, box, instructions = set(), set(), []
    robot = 0

    with file_path.open() as f:
        map_data, instruction_data = f.read().split("\n\n")

        for row_idx, row in enumerate(map_data.split()):
            for col_idx, item in enumerate(row.strip()):
                position = col_idx + row_idx * 1j
                if item == "#":
                    wall.add(position)
                elif item == "O":
                    box.add(position)
                elif item == "@":
                    robot = position

        instructions = [
            {"<": -1, "^": -1j, ">": 1, "v": 1j}[d]
            for d in instruction_data.strip()
            if d in "<^>v"
        ]

    return wall, box, instructions, robot


def simulate(wall, box, instructions, robot):
    box_positions = box.copy()
    robot_position = robot

    for move in instructions:
        stack = [robot_position]
        current_boxes, new_box_positions = set(), set()

        while stack:
            current_position = stack.pop()
            next_position = current_position + move

            if next_position in wall:
                break

            if next_position in box_positions:
                stack.append(next_position)
                current_boxes.add(next_position)
                new_box_positions.add(next_position + move)
        else:
            robot_position += move
            box_positions = (box_positions - current_boxes) | new_box_positions

    return robot_position, box_positions


def increase_wide(walls, boxes, robot):
    robot += robot.real
    new_wall = {w.real + w + offset for w in walls for offset in (0, 1)}
    box_left = {b.real + b for b in boxes}
    box_right = {b.real + b + 1 for b in boxes}
    return new_wall, box_left, box_right, robot


def simulate_new(wall, box_left, box_right, instructions, robot):
    box_positions_left = box_left.copy()
    box_positions_right = box_right.copy()
    robot_position = robot

    for move in instructions:
        stack = [robot_position]
        current_boxes_left, new_box_positions_left = set(), set()
        current_boxes_right, new_box_positions_right = set(), set()

        while stack:
            current_position = stack.pop()
            next_position = current_position + move
            if next_position in wall:
                break

            if (
                next_position in box_positions_left
                and next_position + 1 in box_positions_right
            ):
                if move != 1:
                    stack.append(next_position)
                stack.append(next_position + 1)
                current_boxes_left.add(next_position)
                current_boxes_right.add(next_position + 1)
                new_box_positions_left.add(next_position + move)
                new_box_positions_right.add(next_position + 1 + move)

            if (
                next_position in box_positions_right
                and next_position - 1 in box_positions_left
            ):
                if move != -1:
                    stack.append(next_position)
                stack.append(next_position - 1)
                current_boxes_left.add(next_position - 1)
                current_boxes_right.add(next_position)
                new_box_positions_left.add(next_position - 1 + move)
                new_box_positions_right.add(next_position + move)
        else:
            robot_position += move
            box_positions_left = (
                box_positions_left - current_boxes_left
            ) | new_box_positions_left
            box_positions_right = (
                box_positions_right - current_boxes_right
            ) | new_box_positions_right

    return robot_position, box_positions_left, box_positions_right


def main():
    file_name = "data"
    wall, box, instructions, robot = parse_data(file_name)
    final_robot, final_boxes = simulate(wall, box, instructions, robot)
    print(sum((int(pos.real) + 100 * int(pos.imag)) for pos in final_boxes))

    new_wall, box_left, box_right, new_robot = increase_wide(wall, box, robot)
    final_robot_2, box_left, box_right = simulate_new(
        new_wall, box_left, box_right, instructions, new_robot
    )
    print(sum((int(pos.real) + 100 * int(pos.imag)) for pos in box_left))


if __name__ == "__main__":
    main()