import numpy as np
from scipy.spatial.transform import Rotation as R

from env import setup_environment, shutdown_environment
from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions

# bring every predefined skill into the namespace
from skill_code import *   # provides rotate, move, pick, pull, place, …

def run_combined_task():
    print("\n======================  START COMBINED TASK  ======================")

    # ------------------------------------------------------------
    #  Environment initialisation
    # ------------------------------------------------------------
    env, task = setup_environment()
    try:
        descriptions, obs = task.reset()

        # optional video capture
        init_video_writers(obs)
        task.step           = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # ------------------------------------------------------------
        #  World-state information (exact Cartesian positions)
        # ------------------------------------------------------------
        pos_dict          = get_object_positions()

        bottom_side_pos   = pos_dict['bottom_side_pos']
        bottom_anchor_pos = pos_dict['bottom_anchor_pos']
        rubbish_pos       = pos_dict['rubbish']
        bin_pos           = pos_dict['bin']

        pull_distance     = 0.20        # 20 cm full extension

        # ------------------------------------------------------------
        #  STEP-1…5 — the mandatory “Frozen” part (do NOT change)
        # ------------------------------------------------------------
        target_quat = R.from_euler('z', 90, degrees=True).as_quat()

        # [Frozen Code Start]
        obs, reward, done = rotate(env, task, target_quat)
        obs, reward, done = move(env, task, bottom_side_pos)
        obs, reward, done = move(env, task, bottom_anchor_pos)
        obs, reward, done = pick(env, task, bottom_anchor_pos,
                                         approach_distance=0.05,  # close approach
                                         approach_axis='-z')
        obs, reward, done = pull(env, task, pull_distance=pull_distance,
                                         pull_axis='x')
        # [Frozen Code End]

        if done:
            print("[Abort] Task terminated unexpectedly during the mandatory steps.")
            return

        # ------------------------------------------------------------
        #  Helper – release drawer handle (open gripper)
        # ------------------------------------------------------------
        print("\n[Plan] Helper – release drawer handle")
        current_pos = obs.gripper_pose[:3]
        obs, reward, done = place(env, task, current_pos,
                                  approach_distance=0.00,
                                  approach_axis='z')
        if done:
            print("[Abort] Task finished unexpectedly during release.")
            return

        # ------------------------------------------------------------
        #  STEP-6  pick  – lift the rubbish
        # ------------------------------------------------------------
        print("\n[Plan] Step 6 – pick rubbish at:", rubbish_pos)
        obs, reward, done = pick(env, task, rubbish_pos,
                                 approach_distance=0.10,
                                 approach_axis='z')
        if done:
            print("[Abort] Task finished unexpectedly after picking rubbish.")
            return

        # ------------------------------------------------------------
        #  STEP-7  place  – drop rubbish into bin
        # ------------------------------------------------------------
        print("\n[Plan] Step 7 – place rubbish into bin at:", bin_pos)
        obs, reward, done = place(env, task, bin_pos,
                                  approach_distance=0.10,
                                  approach_axis='z')

        # ------------------------------------------------------------
        #  Final status message
        # ------------------------------------------------------------
        if done:
            print("\n[SUCCESS] Task completed!  Reward:", reward)
        else:
            print("\n[INFO] Task not flagged as done – verify success conditions.")

    except Exception as e:
        print("[ERROR] Exception during task execution:", str(e))

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

if __name__ == "__main__":
    run_combined_task()