import numpy as np


def parse_kinematic(domain_name, task_name, env, obs):
    task_key = "-".join([t for t in task_name.split("-") if t not in ["variant", "v2"]])
    if domain_name == "metaworld_complex":
        task_key = "complex"

    kinematic = {
        "door-open": parse_door,
        "drawer-open": parse_drawer,
        "drawer-close": parse_drawer,
        "button-press": parse_button,
        "faucet-open": parse_faucet,
        "window-open": parse_window,
        "window-close": parse_window,
        "push": parse_push,
        "pick-place": parse_push,
        "peg-insert-side": parse_peg,
        "complex": parse_complex,
    }[task_key](env, obs)

    return kinematic


def parse_door(env, obs):
    unwrapped_env = env.env.env

    hand_pos = obs[:3]
    joint_pos = unwrapped_env.sim.data.get_body_xpos("door_link")
    handle_pos = unwrapped_env.sim.data.get_body_xpos("handle") + np.array(
        [0.395, -0.12, 0]
    )

    kinematic = f"""<hand pos="{hand_pos[0]:.3f} {hand_pos[1]:.3f} {hand_pos[2]:.3f}"/>
    <joint pos="{joint_pos[0]:.3f} {joint_pos[1]:.3f} {joint_pos[2]:.3f}" type="hinge" axis="0 0 1" range="-2 0"/>
    <handle pos="{handle_pos[0]:.3f} {handle_pos[1]:.3f} {handle_pos[2]:.3f}"/>"""
    return kinematic


def parse_drawer(env, obs):
    unwrapped_env = env.env.env

    hand_pos = obs[:3]
    joint_pos = unwrapped_env.sim.data.get_body_xpos("drawer_link")
    handle_pos = unwrapped_env.sim.data.get_body_xpos("handle") + np.array(
        [0, -0.14, 0]
    )

    kinematic = f"""<hand pos="{hand_pos[0]:.3f} {hand_pos[1]:.3f} {hand_pos[2]:.3f}"/>
    <joint pos="{joint_pos[0]:.3f} {joint_pos[1]:.3f} {joint_pos[2]:.3f}" type="slide" axis="0 1 0" range="-0.16 0"/>
    <handle pos="{handle_pos[0]:.3f} {handle_pos[1]:.3f} {handle_pos[2]:.3f}"/>"""
    return kinematic


def parse_button(env, obs):
    unwrapped_env = env.env.env

    hand_pos = obs[:3]
    joint_pos = unwrapped_env.sim.data.get_body_xpos("button")
    button_pos = unwrapped_env.sim.data.get_site_xpos("buttonStart")

    kinematic = f"""<hand pos="{hand_pos[0]:.3f} {hand_pos[1]:.3f} {hand_pos[2]:.3f}"/>
    <joint pos="{joint_pos[0]:.3f} {joint_pos[1]:.3f} {joint_pos[2]:.3f}" type="slide" axis="0 -1 0" stiffness="0.5" range="-0.06 0" damping="1"/>
    <button pos="{button_pos[0]:.3f} {button_pos[1]:.3f} {button_pos[2]:.3f}"/>"""
    return kinematic


def parse_faucet(env, obs):
    unwrapped_env = env.env.env

    hand_pos = obs[:3]
    joint_pos = unwrapped_env.sim.data.get_body_xpos("faucet_link2")
    faucet_pos = unwrapped_env.sim.data.get_body_xpos("faucet_link2") + np.array(
        [0, -0.15, 0]
    )

    kinematic = f"""<hand pos="{hand_pos[0]:.3f} {hand_pos[1]:.3f} {hand_pos[2]:.3f}"/>
    <joint pos="{joint_pos[0]:.3f} {joint_pos[1]:.3f} {joint_pos[2]:.3f}" type="hinge" range="-1.57 1.57" axis="0 0 1"/>
    <faucet pos="{faucet_pos[0]:.3f} {faucet_pos[1]:.3f} {faucet_pos[2]:.3f}"/>"""
    return kinematic


def parse_window(env, obs):
    unwrapped_env = env.env.env

    hand_pos = obs[:3]
    joint_pos = unwrapped_env.sim.data.get_body_xpos("windowb_a")
    if "close" in env.env_name:
        handle_pos = unwrapped_env.sim.data.get_site_xpos(
            "handleCloseStart"
        ) + np.array([0.2, 0.0, 0.04])
    else:
        handle_pos = unwrapped_env.sim.data.get_site_xpos("handleOpenStart") + np.array(
            [0, 0.0, 0.04]
        )

    kinematic = f"""<hand pos="{hand_pos[0]:.3f} {hand_pos[1]:.3f} {hand_pos[2]:.3f}"/>
    <joint pos="{joint_pos[0]:.3f} {joint_pos[1]:.3f} {joint_pos[2]:.3f}" type="slide" range="0 .2" axis="1 0 0"/>
    <handle pos="{handle_pos[0]:.3f} {handle_pos[1]:.3f} {handle_pos[2]:.3f}"/>"""
    return kinematic


def parse_push(env, obs):
    unwrapped_env = env.env.env

    hand_pos = obs[:3]
    object_pos = unwrapped_env.sim.data.get_body_xpos("obj")
    goal_pos = obs[-3:]

    kinematic = f"""<hand pos="{hand_pos[0]:.3f} {hand_pos[1]:.3f} {hand_pos[2]:.3f}"/>
    <object pos="{object_pos[0]:.3f} {object_pos[1]:.3f} {object_pos[2]:.3f}" type="cylinder" pos="0 0 0" size="0.02 0.02"/>
    <goal pos="{goal_pos[0]:.3f} {goal_pos[1]:.3f} {goal_pos[2]:.3f}"/>"""
    return kinematic


def parse_peg(env, obs):
    unwrapped_env = env.env.env

    hand_pos = obs[:3]
    peg_pos = unwrapped_env.sim.data.get_body_xpos("peg")
    goal_pos = obs[-3:]

    kinematic = f"""<hand pos="{hand_pos[0]:.3f} {hand_pos[1]:.3f} {hand_pos[2]:.3f}"/>
    <peg pos="{peg_pos[0]:.3f} {peg_pos[1]:.3f} {peg_pos[2]:.3f}" type="box" euler="0 1.57 0" size="0.015 0.015 0.12"/>
    <goal pos="-0.350 {goal_pos[1]:.3f} 0.160"/>"""
    return kinematic


def parse_complex(env, obs):
    unwrapped_env = env.env.env

    hand_pos = obs[:3]

    puck_obs = obs[4:21]
    puck_pos = np.array([*puck_obs[0:2], 0.045]) + np.array([0, -0.01, 0])
    puck_goal_pos = puck_obs[-3:] + np.array([-0.1, 0, 0])

    handle_joint_pos = unwrapped_env.sim.data.get_body_xpos("drawer_link")
    handle_pos = unwrapped_env.sim.data.get_body_xpos("drawer_link") + np.array(
        [0, -0.14, 0]
    )

    button_joint_pos = unwrapped_env.sim.data.get_body_xpos("button")
    button_pos = unwrapped_env.sim.data.get_site_xpos("buttonStart")

    peg_obs = obs[55:72]
    peg_goal_pos = peg_obs[-3:]  #
    peg_pos = peg_obs[:3]

    goal_pos = obs[-3:]

    kinematic = f"""<hand pos="{hand_pos[0]:.3f} {hand_pos[1]:.3f} {hand_pos[2]:.3f}"/>
    <peg pos="{peg_pos[0]:.3f} {peg_pos[1]:.3f} {peg_pos[2]:.3f}" type="box" euler="0 1.57 0" size="0.015 0.015 0.12"/>
    <goal pos="-0.350 {goal_pos[1]:.3f} 0.160"/>"""

    kinematic = f"""<hand pos="{hand_pos[0]:.3f} {hand_pos[1]:.3f} {hand_pos[2]:.3f}"/>
    <puck>
        <puck pos="{puck_pos[0]:.3f} {puck_pos[1]:.3f} {puck_pos[2]:.3f}" type="box" size="0.05 0.05 0.015" mass=".05"/>
        <goal pos="{puck_goal_pos[0]:.3f} {puck_goal_pos[1]:.3f} {puck_goal_pos[2]:.3f}"/>
    </puck>
    <drawer>
        <joint pos="{handle_joint_pos[0]:.3f} {handle_joint_pos[1]:.3f} {handle_joint_pos[2]:.3f}" type="slide" axis="0 1 0" range="-0.16 0"/>
        <handle pos="{handle_pos[0]:.3f} {handle_pos[1]:.3f} {handle_pos[2]:.3f}"/>
    </drawer>
    <button>
        <joint pos="{button_joint_pos[0]:.3f} {button_joint_pos[1]:.3f} {button_joint_pos[2]:.3f}" type="slide" axis="0 -1 0" stiffness="0.5" range="-0.06 0" damping="1"/>
        <button pos="{button_pos[0]:.3f} {button_pos[1]:.3f} {button_pos[2]:.3f}"/>
    </button>
    <peg>
        <peg pos="{peg_pos[0]:.3f} {peg_pos[1]:.3f} {peg_pos[2]:.3f}" type="box" euler="0 1.57 0" size="0.015 0.015 0.12"/>
        <goal pos="{peg_goal_pos[0]:.3f} {peg_goal_pos[1]:.3f} {peg_goal_pos[2] + 0.03:.3f}"/>
    </peg>
    """
    return kinematic
