# run_skeleton_task.py (Completed)

import traceback
import numpy as np
from pyrep.objects.shape import Shape
from pyrep.objects.proximity_sensor import ProximitySensor

from env import setup_environment, shutdown_environment

# We are explicitly importing every skill that can be used.
# DO NOT define new skills – the judge provides the real implementations.
from skill_code import pick, place, move, rotate, pull

from video import init_video_writers, recording_step, recording_get_observation

from object_positions import get_object_positions


# -----------------------------------------------------------
# Helper Utilities
# -----------------------------------------------------------

def safe_skill_call(skill_fn, *pos_args, **kw_args):
    """
    Because we do not know the exact signature of the provided skills,
    try several common call-styles before giving up.  The order is:

        1) skill_fn(*pos_args, **kw_args)
        2) skill_fn(*pos_args)
        3) skill_fn(**kw_args)

    Any TypeError coming from argument-mismatch will be ignored and the
    next variant tested.  All other exceptions are re-raised so that the
    user is aware of real execution problems (e.g. IK failure).
    """
    variants = (
        lambda: skill_fn(*pos_args, **kw_args),
        lambda: skill_fn(*pos_args),
        lambda: skill_fn(**kw_args),
    )
    last_err = None
    for v in variants:
        try:
            return v()
        except TypeError as e:
            # Argument mismatch – try the next variant
            last_err = e
            continue
    # If we reach here the call could not be matched
    if last_err:
        raise last_err


def try_rotate_gripper(env, task, gripper_name):
    """
    Exploration phase:  we do not initially know if predicate (rotated ?g ?a)
    holds for ANY angle.  We therefore try to rotate the gripper between two
    canonical angles.  If the rotate() skill executes without throwing a
    planning-specific exception we consider the predicate discovered.
    """
    print("[Exploration]   Attempting to discover predicate (rotated …)")
    discovered = False

    # Two canonical angles that almost always exist in drawer tasks.
    ANGLES = ["zero_deg", "ninety_deg"]

    for from_a in ANGLES:
        for to_a in ANGLES:
            if from_a == to_a:
                continue
            try:
                safe_skill_call(
                    rotate,
                    env,
                    task,
                    gripper_name=gripper_name,
                    from_angle=from_a,
                    to_angle=to_a
                )
                print(f"  ✔  Rotate succeeded: {gripper_name} {from_a} → {to_a}")
                discovered = True
                return discovered
            except Exception as e:
                # Any exception means either wrong arguments or the rotation is
                # not possible – we continue trying.
                continue

    if not discovered:
        print("  ✘  Could not execute rotate – predicate may still be unknown.")
    return discovered


def open_drawer_sequence(env, task, positions):
    """
    A generic plan that will – if the simulated workspace follows the usual
    naming conventions – rotate the gripper, move to the drawer handle, grab
    it and pull the drawer open.
    """
    # ---------------------------------------------
    # Heuristic lookup for important object names
    # ---------------------------------------------
    drawer_name  = next((n for n in positions if "drawer"  in n.lower()), None)
    gripper_name = next((n for n in positions if "gripper" in n.lower()), None)
    side_pos     = next((n for n in positions if "side"    in n.lower()), None)
    anchor_pos   = next((n for n in positions if "anchor"  in n.lower()), None)

    # If any of the required objects/positions are missing we abort the plan.
    required = [drawer_name, gripper_name, side_pos, anchor_pos]
    if any(x is None for x in required):
        print("[Plan]   Environment does not expose the necessary named objects – skipping drawer plan.")
        return

    print("[Plan]   Executing drawer-opening plan…")

    # 1. Ensure the gripper is at zero_deg and then rotate to ninety_deg
    try_rotate_gripper(env, task, gripper_name)

    # 2. Move gripper to the side of the drawer
    try:
        print("  – Moving gripper to side position")
        safe_skill_call(
            move,
            env,
            task,
            gripper_name=gripper_name,
            target_pos=positions[side_pos]
        )
    except Exception as e:
        print("    ! move-to-side failed; aborting drawer plan.")
        traceback.print_exc()
        return

    # 3. Move from side position to anchor position
    try:
        print("  – Moving gripper to anchor (handle) position")
        safe_skill_call(
            move,
            env,
            task,
            gripper_name=gripper_name,
            target_pos=positions[anchor_pos]
        )
    except Exception as e:
        print("    ! move-to-anchor failed; aborting drawer plan.")
        traceback.print_exc()
        return

    # 4. Grab the drawer handle (pick-drawer reduced to standard pick interface)
    try:
        print("  – Gripping the drawer handle")
        safe_skill_call(
            pick,
            env,
            task,
            target_obj=drawer_name,
            target_pos=positions[anchor_pos]
        )
    except Exception as e:
        print("    ! pick-drawer failed; aborting drawer plan.")
        traceback.print_exc()
        return

    # 5. Pull the drawer
    try:
        print("  – Pulling the drawer open")
        safe_skill_call(
            pull,
            env,
            task,
            target_obj=drawer_name
        )
    except Exception as e:
        print("    ! pull failed; the drawer may already be open or locked.")
        traceback.print_exc()
        return

    print("[Plan]   Drawer opening procedure finished.")


# -----------------------------------------------------------
# Main Entry Point
# -----------------------------------------------------------

def run_skeleton_task():
    print("===== Starting Skeleton Task =====")

    env, task = setup_environment()
    try:
        # Reset task to initial state and initialise video recording
        descriptions, obs = task.reset()
        init_video_writers(obs)

        # Wrap step() and get_observation() so any skill that uses them will
        # automatically record video frames.
        task.step           = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # -------------------------------------------------
        # Retrieve object positions for later reasoning
        # -------------------------------------------------
        print("[Init]    Retrieving object positions from helper module")
        try:
            positions = get_object_positions()
            print(f"[Init]    Found {len(positions)} named poses")
        except Exception as e:
            print("[Init]    Could not retrieve positions – continuing with empty dict")
            positions = {}

        # -------------------------------------------------
        # EXPLORATION:  find missing predicate 'rotated'
        # -------------------------------------------------
        candidate_gripper = next((n for n in positions if "gripper" in n.lower()), None)
        if candidate_gripper:
            predicate_found = try_rotate_gripper(env, task, candidate_gripper)
            if predicate_found:
                print("[Exploration]   Predicate (rotated …) confirmed.")
            else:
                print("[Exploration]   Predicate (rotated …) still uncertain – subsequent skills may fail.")
        else:
            print("[Exploration]   No gripper object found – skipping rotation exploration.")

        # -------------------------------------------------
        # MAIN PLAN
        # -------------------------------------------------
        open_drawer_sequence(env, task, positions)

        # Additional plans (pick-and-place, disposal, etc.) could follow here
        # if the judge task required it.  For now we concentrate on the drawer
        # because the feedback indicated a missing 'rotated' predicate that is
        # crucial for that scenario.

    finally:
        shutdown_environment(env)

    print("===== End of Skeleton Task =====")


if __name__ == "__main__":
    run_skeleton_task()