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

from env import setup_environment, shutdown_environment

# Low-level skills that already exist
from skill_code import move, pick, place, rotate, pull

# (Optional) video utilities
from video import init_video_writers, recording_step, recording_get_observation

# Helper that returns {object_name(str): np.ndarray([x, y, z])}
from object_positions import get_object_positions


# ---------------------------------------------------------------------------
# Small utility helpers
# ---------------------------------------------------------------------------

def resolve_name(name, positions, aliases=None):
    """
    Convert a symbolic name appearing in the high-level plan into an XYZ pose
    contained in the dictionary returned by get_object_positions().

    Parameters
    ----------
    name : str
        Canonical symbolic name such as ‘side-pos-bottom’
    positions : dict
        Mapping produced by get_object_positions()
    aliases : list[str] | None
        Alternative spellings to try, e.g. [‘bottom_side_pos’]

    Returns
    -------
    np.ndarray(3,)
    """
    candidates = [name]
    if aliases:
        candidates.extend(aliases)

    for cand in candidates:
        for variant in (cand,
                        cand.replace('-', '_'),
                        cand.replace('_', '-')):
            if variant in positions:
                return np.asarray(positions[variant], dtype=float)

    raise KeyError(f"[resolve_name] Could not find {candidates} among "
                   f"available keys: {list(positions.keys())}")


def quaternion_from_euler(rx_deg=0.0, ry_deg=0.0, rz_deg=0.0, seq='xyz'):
    """Convert Euler angles (deg) → xyzw quaternion"""
    from scipy.spatial.transform import Rotation as R
    return R.from_euler(seq,
                        np.deg2rad([rx_deg, ry_deg, rz_deg])).as_quat()


# ---------------------------------------------------------------------------
# Main routine that executes the oracle (Specification) plan
# ---------------------------------------------------------------------------

def run_skeleton_task():
    print("\n====================  START TASK  ====================")

    # ----------------------------------------------------------------------
    # 1)  Environment set-up
    # ----------------------------------------------------------------------
    env, task = setup_environment()
    try:
        _, obs = task.reset()

        # Optional video recording
        init_video_writers(obs)
        task.step = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # Cache static object positions
        positions = get_object_positions()

        # ------------------------------------------------------------------
        # 2)  Execute the nine steps described in the given Specification
        # ------------------------------------------------------------------
        done = False
        reward = 0.0

        # --------  STEP-1  -------------------------------------------------
        quat_90_z = quaternion_from_euler(rz_deg=90.0)
        target = resolve_name('side-pos-bottom', positions,
                              aliases=['bottom_side_pos'])

        # [Frozen Code Start]
        obs, reward, done = rotate(env, task, target_quat=quat_90_z)
        obs, reward, done = move(env, task, target_pos=target)
        obs, reward, done = move(env, task, target_pos=target)
        # [Frozen Code End]

        # --------  STEP-3 (corrected)  ------------------------------------
        # The second frozen move repeated the same pose; now move to the
        # correct anchor-pos-bottom before grasping the drawer handle.
        if not done:
            target_anchor = resolve_name('anchor-pos-bottom', positions,
                                         aliases=['bottom_anchor_pos'])
            print("\n[PLAN] Extra move → anchor-pos-bottom  ", target_anchor)
            obs, reward, done = move(env, task, target_pos=target_anchor)

        # --------  STEP-4  pick-drawer (grasp handle) ---------------------
        if not done:
            print("\n[PLAN] Step-4  grasp drawer handle at anchor-pos-bottom")
            obs, reward, done = pick(env, task,
                                     target_pos=target_anchor,
                                     approach_distance=0.10,
                                     approach_axis='z')

        # --------  STEP-5  pull drawer fully open -------------------------
        if not done:
            print("\n[PLAN] Step-5  pull drawer +0.25 m along +X")
            obs, reward, done = pull(env, task,
                                     pull_distance=0.25,
                                     pull_axis='x')

        # --------  STEP-6  pick tomato1 -----------------------------------
        if not done:
            tomato1_pos = resolve_name('tomato1', positions,
                                       aliases=['item1'])
            print("\n[PLAN] Step-6  pick tomato1 at", tomato1_pos)
            obs, reward, done = pick(env, task,
                                     target_pos=tomato1_pos,
                                     approach_distance=0.12,
                                     approach_axis='z')

        # --------  STEP-7  place tomato1 on plate -------------------------
        if not done:
            plate_pos = resolve_name('plate', positions)
            print("\n[PLAN] Step-7  place tomato1 on plate at", plate_pos)
            obs, reward, done = place(env, task,
                                      target_pos=plate_pos,
                                      approach_distance=0.12,
                                      approach_axis='z')

        # --------  STEP-8  pick tomato2 -----------------------------------
        if not done:
            tomato2_pos = resolve_name('tomato2', positions,
                                       aliases=['item2'])
            print("\n[PLAN] Step-8  pick tomato2 at", tomato2_pos)
            obs, reward, done = pick(env, task,
                                     target_pos=tomato2_pos,
                                     approach_distance=0.12,
                                     approach_axis='z')

        # --------  STEP-9  place tomato2 on plate -------------------------
        if not done:
            print("\n[PLAN] Step-9  place tomato2 on plate at", plate_pos)
            obs, reward, done = place(env, task,
                                      target_pos=plate_pos,
                                      approach_distance=0.12,
                                      approach_axis='z')

        # ------------------------------------------------------------------
        # 3)  Final status
        # ------------------------------------------------------------------
        if done:
            print("\n[TASK] Environment signalled done=True. Final reward:", reward)
        else:
            print("\n[TASK] Scripted plan finished. done flag:", done,
                  " reward:", reward)

    finally:
        shutdown_environment(env)
        print("=====================  END TASK  =====================")


# ---------------------------------------------------------------------------
# Entry-point
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    run_skeleton_task()