Exercism exercise (reminding myself of some python idioms)

two_buckets.py

from copy import copy, deepcopy
from dataclasses import dataclass
from functools import partial

#### problem:

# https://exercism.org/tracks/python/exercises/two-bucket

#### basic algorithm:

# brute force ftw! We could be cleverer but in this case that
# doesn't suit my learning goals.

# - breadth-first search through action tree until we hit success
#   - breadth-first means we have to keep a queue of states + steps to be taken
# - nodes that can be abandoned:
#   - ones that put us back into a state we've been in before
#   - ones that put us into the illegal state where starting bucket is empty and
#     other bucket is full
#   - no point emptying an empty bucket
#   - or filling a full bucket

#### questions:

# - is it easier to just always make the starting bucket bucket 0 internally
#   & then convert back at the end? [decision: yes]

#### assumptions:

# - The instructions don't say it, but I'm going to assume that the desired
#   amount has to be able to fit in one bucket, eg I can't have one bucket
#   that holds 7 and another that holds 4, fill them both, and say I've
#   achieved 11.

#### notes:

# we need to pass around enough varied state (unless we want to make it
# global) that we might as well make a dataclass for it.

# we'll treat most of that state as immutable and copy or replace it,
# except:
# 1. we want to keep a global mutable queue (but will put it in the state
#    for convenience)
# 2. we want to keep a global set of already-visited states so we don't
#    duplicate effort

# note that we don't bother with a class for the buckets, we'll always
# represent the buckets with two tuples of ints: one for their capacity
# and one for their current contents, both in liters.

# At some point if we've gone deep enough in the tree we should declare
# failure so we don't run forever.
MAX_DEPTH = 30

@dataclass
class S: # State
    queue: list # of thunks as lambdas - mutable singleton!
    visited_states: set[tuple[int]] # mutable; bucket states we've already seen
    subtree_states: list[str] # path to the current state
    bucket_capacity: tuple[int]
    bucket_state: tuple[int]
    starting_bucket: int
    depth_in_tree: int
    valid: bool
    success: bool

num_to_name = {0: "one", 1: "two"}
name_to_num = {"one": 0, "two": 1}
other_bucket = {0: 1, 1: 0}

## actions:

def pour_a_into_b(state: S, bucket_a: int):
    capacity_b = state.bucket_capacity[other_bucket[bucket_a]]
    old_a = state.bucket_state[bucket_a]
    old_b = state.bucket_state[other_bucket[bucket_a]]
    new_b = old_b + old_a
    new_a = 0
    # Did we overfill b? If so, put some back.
    if new_b > capacity_b:
        new_a = new_b - capacity_b
        new_b = capacity_b
    if bucket_a == 0:
        state.bucket_state = (new_a, new_b)
    else:
        state.bucket_state = (new_b, new_a)
    state.subtree_states = state.subtree_states + [f"Poured {bucket_a} into {other_bucket[bucket_a]} resulting in {state.bucket_state}."]
    return state

def empty(state: S, bucket: int):
    # no point emptying an empty bucket
    if state.bucket_state[bucket] == 0:
        state.valid = False
        return state
    # feels like there should be a better way to express this next bit without an if/else (here & in fill)
    # but I didn't see one in a few minutes of thinking.
    if bucket == 0:
        state.bucket_state = (0, state.bucket_state[1])
    else:
        state.bucket_state = (state.bucket_state[0], 0)
    state.subtree_states = state.subtree_states + [f"Emptied {bucket}, resulting in {state.bucket_state}."]
    return state

def fill(state: S, bucket: int):
    # no point filling a full bucket
    if state.bucket_state[bucket] == state.bucket_capacity[bucket]:
        state.valid = False
        return state
    if bucket == 0:
        state.bucket_state = (state.bucket_capacity[0], state.bucket_state[1])
    else:
        state.bucket_state = (state.bucket_state[0], state.bucket_capacity[1])
    state.subtree_states = state.subtree_states + [f"Filled {bucket}, resulting in {state.bucket_state}."]
    return state

def action_list(state: S):
    """Returns a list of thunks, one for each (action_type, bucket) combo"""
    action_types = [
        pour_a_into_b,
        empty,
        fill
    ]
    actions = []
    for action_type in action_types:
        for bucket in [0, 1]:
            # using partial rather than lambda to avoid very-late-binding
            actions.append(partial(action_type, copy(state), bucket))
    return actions

def continuable_state(state: S):
    return ((state.queue)                      and
            (state.depth_in_tree <= MAX_DEPTH) and
            (state.success == False))

def create_initial_state(bucket_one_l, bucket_two_l, goal_l, start_bucket_name):
    """
    Note: we reorder for convenience so the start bucket is always the first one
    """
    # Some initial error checking:
    try:
        start_bucket = name_to_num[start_bucket_name]
    except:
        raise ValueError(f"Error: '{start_bucket_name}' is not a valid bucket name.")
    if (goal_l > bucket_one_l) and (goal_l > bucket_two_l):
        raise ValueError(f"Neither of these buckets (with capacities {bucket_one_l} and {bucket_two_l}) can hold {goal_l} liters!")

    queue = []
    visited_states = {(0,0)} # set of tuples
    # Internally, the first bucket is always the starting bucket
    if start_bucket == 0:
        capacities = (bucket_one_l, bucket_two_l)
    else:
        capacities = (bucket_two_l, bucket_one_l)
    bucket_state = (0, 0)
    initial_state = S(queue=queue,
                      visited_states=visited_states,
                      subtree_states=[],
                      bucket_capacity=capacities,
                      bucket_state=bucket_state,
                      starting_bucket=start_bucket,
                      depth_in_tree=0,
                      valid=True,
                      success=False,)
    return initial_state

def measure(bucket_one_l, bucket_two_l, goal_l, start_bucket_name):
    initial_state = create_initial_state(bucket_one_l,
                                         bucket_two_l,
                                         goal_l,
                                         start_bucket_name)
    # first action is fixed by the description. Recall that the start bucket,
    # whichever it is, has been placed at position 0 for convenience.
    first_action = lambda: fill(initial_state, 0)
    initial_state.queue.append(first_action)
    state = initial_state

    while continuable_state(state):
        thunk = state.queue.pop(0)

        state = thunk()
        state.depth_in_tree += 1

        #     Nodes that can be abandoned:
        if (
            # - ones that put us back into a state we've been in before (not 'invalid' exactly, just pointless)
            (state.bucket_state in state.visited_states) or
            # - ones that put us into the illegal state where starting bucket is empty and
            #   other bucket is full
            (state.bucket_state[0] == 0 and state.bucket_state[1] == state.bucket_capacity[1]) or
            # - or either bucket is negative (should never happen!)
            (state.bucket_state[0] < 0 or state.bucket_state[1] < 0 ) or
            # - or either bucket is over capacity (should never happen!)
            (state.bucket_state[0] > state.bucket_capacity[0] or state.bucket_state[1] > state.bucket_capacity[1])
        ):
            state.valid = False

        if state.valid:
            state.visited_states.add(state.bucket_state)
            if state.bucket_state[0] == goal_l or state.bucket_state[1] == goal_l:
                state.success = True
            new_actions = action_list(state)
            state.queue.extend(new_actions)

    # Yay!
    if state.success == True:
        # translate back from 0 as starting bucket to whichever one was given
        winning_bucket_internal = 0 if state.bucket_state[0] == goal_l else 1
        other_bucket_l = state.bucket_state[other_bucket[winning_bucket_internal]]
        winning_bucket = (winning_bucket_internal
                          if state.starting_bucket == 0
                          else other_bucket[winning_bucket_internal])
        temp_state = deepcopy(state) # just for printing
        temp_state.queue = f"Queue had {len(temp_state.queue)} items."
        temp_state.visited_states = []
        temp_state.subtree_states = []
        print()
        print(f"Success!")
        print(f"winning_bucket: {num_to_name[winning_bucket]}")
        print(f"other_bucket_l: {other_bucket_l}")
        print(f"winning state: {temp_state}")
        print()
        print(f"Path to winning (capacities {state.bucket_capacity}, goal {goal_l}):")
        print(f"--------------------------------------------")
        print("\n".join(state.subtree_states))
        return state.depth_in_tree, num_to_name[winning_bucket], other_bucket_l,

    # we've fallen out of the while loop because the queue is empty or we've
    # gone unacceptably many layers deep
    if not state.queue:
        raise ValueError("O NOES we've run out of options, this appears to be impossible.")
    if state.depth_in_tree > MAX_DEPTH:
        raise ValueError(f"Giving up because we're {MAX_DEPTH} levels deep.")

    raise ValueError(f"Uhhhhhh something is horribly wrong how did we get here. State: {state}")