from pyhop import hop

# Define state  ## --> problem file 
state1 = hop.State("state1")

state1.robot_at = "table1"
state1.gripper_free = True 
state1.holding = None  ## what the robot is holding

state1.on_table1 = {
    "knife"   : 10,
    "fork"    : 10,
    "napkin"  : 10,
    "plate"   : 10,
    "cup"     : 10
}
state1.on_table2 = {
    "knife"   : 0,
    "fork"    : 0,
    "napkin"  : 0,
    "plate"   : 0,
    "cup"     : 0
}

## Primitive operations

def goto(state, table):
    state.robot_at = table
    return state


def pick(state, item, table):
    """Robot picks up an object from any table"""
    if state.gripper_free != True:
        return False
    # Use the table name to access the corresponding dictionary
    table_dict = getattr(state, f'on_{table}')  # on_table1 or on_table2
    if state.robot_at == table and table_dict[item] > 0:
        table_dict[item] -= 1
        state.gripper_free = False
        state.holding = item
        return state
    return False


def drop(state, item, table):
    """Robot drops an object on any table"""
    if state.gripper_free == True:
        return False
    # Use the table name to access the corresponding dictionary
    table_dict = getattr(state, f'on_{table}')  # on_table1 or on_table2
    if state.robot_at == table and state.holding == item:
        table_dict[item] += 1
        state.gripper_free = True
        state.holding = None
        return state
    return False


hop.declare_operators(goto, pick, drop)

## Methods

def bring_item(state, item, table_src, table_dest):
    return [('goto', table_src),
            ('pick', item, table_src),
            ('goto', table_dest),
            ('drop', item, table_dest)]


def prepare_place(state, table_src, table_dest):
    return [('bring_item', 'knife', table_src, table_dest),
            ('bring_item', 'fork', table_src, table_dest),
            ('bring_item', 'napkin', table_src, table_dest),
            ('bring_item', 'plate', table_src, table_dest),
            ('bring_item', 'cup', table_src, table_dest)]


## Clear the table
def clear_place(state, table_to_clean, table_dest):
    lst = []
    table_dict = getattr(state, f'on_{table_to_clean}')
    for item in ['knife', 'fork', 'napkin', 'plate', 'cup']:
        if table_dict[item] > 0:
            lst.append(('bring_item', item, table_to_clean, table_dest))
    return lst

hop.declare_methods('bring_item', bring_item)
hop.declare_methods('prepare_place', prepare_place)
hop.declare_methods('clear_place', clear_place)

# Plan a sequence of actions to reach a goal
hop.plan(state1, [('prepare_place', 'table1', 'table2')], 
         hop.get_operators(),
         hop.get_methods(),
         verbose=1)

hop.plan(state1, [('clear_place', 'table2', 'table1')],
         hop.get_operators(),
         hop.get_methods(),
         verbose=1)


#####
# Execute plan
def execute_plan(state, plan):
    for action in plan:
        op_name, *args = action
        operator = hop.get_operators().get(op_name)
        if operator is None:
            raise ValueError(f"Operator {op_name} not found")
        new_state = operator(state, *args)
        if new_state is False:
            raise RuntimeError(f"Action {action} failed")
    return state


# PLANNING

print("Prepare table2")

plan1 = hop.plan(state1, [('prepare_place', 'table1', 'table2')], 
                 hop.get_operators(),
                 hop.get_methods(),
                 verbose=1)

print(f"Plan found: {len(plan1)} actions")
print(f"Table1 after planning: {state1.on_table1}")
print(f"Table2 after planning: {state1.on_table2}\n")
execute_plan(state1, plan1)
print(f"Table1 after execution: {state1.on_table1}")
print(f"Table2 after execution: {state1.on_table2}\n")


print("Clear table2")

plan2 = hop.plan(state1, [('clear_place', 'table2', 'table1')],
                 hop.get_operators(),
                 hop.get_methods(),
                 verbose=1)

print(f"Plan found: {len(plan2)} actions")
print(f"Final Table1: {state1.on_table1}")
print(f"Final Table2: {state1.on_table2}")

execute_plan(state1, plan2)

print(f"Final Table1 after execution: {state1.on_table1}")
print(f"Final Table2 after execution: {state1.on_table2}")