import isaaclab.sim as sim_utils
from isaaclab.actuators import ImplicitActuatorCfg
from isaaclab.assets.articulation import ArticulationCfg

import numpy as np


####### ROBOT DEFINITIONS #######
LEFT_PANDA_DEFAULT_JOINT_POS = {
    "panda_joint1": -0.6109,
    "panda_joint2": -0.6109,
    "panda_joint3": 0.0,
    "panda_joint4": -2.3562,
    "panda_joint5": 0.0,
    "panda_joint6": 1.8326,
    "panda_joint7": 0.7854,

    "panda_finger_joint1": -0.04,
    "panda_finger_joint2": 0.04,

}

RIGHT_PANDA_DEFAULT_JOINT_POS = {
    "panda_joint1": 0.4363,
    "panda_joint2": -0.6109,
    "panda_joint3": 0.0,
    "panda_joint4": -2.3562,
    "panda_joint5": 0.0,
    "panda_joint6": 1.8326,
    "panda_joint7": -1.1345,
    "panda_finger_joint1": -0.04,
    "panda_finger_joint2": 0.04,
}


PANDA_ARTICULATION = ArticulationCfg(
    spawn=sim_utils.UsdFileCfg(
        # usd_path=f"./envs/2_riverway_shelf/models/assembled_robot/full_robot.usd",
        usd_path=f"./envs/lbm_usd_library/stations/cabot_simulation/i justwantrobot.usda",
        rigid_props=sim_utils.RigidBodyPropertiesCfg(
            # disable_gravity=True,
            disable_gravity=False,
            max_depenetration_velocity=5.0,
        ),
        articulation_props=sim_utils.ArticulationRootPropertiesCfg(
            enabled_self_collisions=False, solver_position_iteration_count=36, solver_velocity_iteration_count=0
        ),
    ),
    init_state=ArticulationCfg.InitialStateCfg(
        pos=(0, 0, 0),  # TODO:  this needs to be updated
        rot=(1, 0, 0, 0),
        joint_pos=LEFT_PANDA_DEFAULT_JOINT_POS
    ),
    soft_joint_pos_limit_factor=1,
)



####### EFFORT-ONLY ACTUATOR (for gravity-compensated PD control) #######
# Effort limits from Drake's Franka Panda URDF.
# Joints 1-4: 87 N·m, Joints 5-7: 12 N·m.
PANDA_EFFORT_LIMITS = {
    "panda_joint1": 87.0,
    "panda_joint2": 87.0,
    "panda_joint3": 87.0,
    "panda_joint4": 87.0,
    "panda_joint5": 12.0,
    "panda_joint6": 12.0,
    "panda_joint7": 12.0,
}

# Kp (torque-space). Derived from Drake's IDC acceleration-space gains
# [2000, 1500, 1500, 1500, 1500, 500, 500] scaled by typical M(q) diagonal.
# Stable at dt=10ms (100Hz physics). These need empirical tuning.
PANDA_KP = {
    "panda_joint1": 2000.0,
    "panda_joint2": 1500.0,
    "panda_joint3": 1500.0,
    "panda_joint4": 1500.0,
    "panda_joint5": 1500.0,
    "panda_joint6": 500.0,
    "panda_joint7": 500.0,
}

# Kd (torque-space). Same derivation as Kp.
PANDA_KD = {
    "panda_joint1": 89.4,
    "panda_joint2": 77.5,
    "panda_joint3": 77.0,
    "panda_joint4": 77.5,
    "panda_joint5": 77.5,
    "panda_joint6": 44.7,
    "panda_joint7": 44.7,
}

EXPLICIT_PANDA = PANDA_ARTICULATION.copy()
EXPLICIT_PANDA.actuators = {
    "arm": ImplicitActuatorCfg(
        joint_names_expr=["panda_joint.*"],
        stiffness=0.0,
        damping=0.0,
        velocity_limit=2.175,
        effort_limit_sim=PANDA_EFFORT_LIMITS,
    ),
}


IMPLICIT_PANDA = PANDA_ARTICULATION.copy()
IMPLICIT_PANDA.actuators = {
        "panda_shoulder": ImplicitActuatorCfg(
            joint_names_expr=["panda_joint[1-4]"],
            effort_limit_sim=87.0,
            stiffness=400.0,
            damping=80.0,
            # stiffness=80.0,
            # damping=4.0,
        ),
        "panda_forearm": ImplicitActuatorCfg(
            joint_names_expr=["panda_joint[5-7]"],
            effort_limit_sim=12.0,
            stiffness=400.0,
            damping=80.0,
            # stiffness=80.0,
            # damping=4.0,
        ),
        "panda_hand": ImplicitActuatorCfg(
            joint_names_expr=["panda_finger_joint.*"],
            effort_limit_sim=200.0,
            stiffness=2000.0,
            damping=200.0
            # velocity_limit=0.01
        ),
}
####### ROBOT DEFINITIONS #######


####### ACTION DEFINITIONS #######
from isaaclab.utils import configclass
from isaaclab.controllers.differential_ik_cfg import DifferentialIKControllerCfg
from isaaclab.envs.mdp.actions.actions_cfg import DifferentialInverseKinematicsActionCfg
from sim_improvement.environments.mdp import (
    GravityCompensatedJointPositionActionCfg,
    GravityCompensatedRelativeJointPositionActionCfg,
    GravityCompensatedDifferentialIKActionCfg,
    EnvFrameDiffIKActionCfg,
    JointPositionActionCfg,
    AbsBinaryJointPositionActionCfg,
    JointEffortActionCfg,
)
from .actions import DualArmEnvFrameDiffIKActionCfg

LBM_JOINT_EFFORT = JointEffortActionCfg(
    asset_name="robot", 
    joint_names=["panda_joint.*"],
)

LBM_GRAVITY_COMP_DIFF_IK = GravityCompensatedDifferentialIKActionCfg(
    asset_name="robot",
    joint_names=["panda_joint.*"],
    body_name="panda_link8",
    controller=DifferentialIKControllerCfg(
        command_type="pose",
        use_relative_mode=False,
        ik_method="dls",
        ik_params={"lambda_val": 0.1},
    ),
    scale=1.0,
    kp=PANDA_KP,
    kd=PANDA_KD,
    effort_limit=PANDA_EFFORT_LIMITS,
)

ARM_IK_ACTION = EnvFrameDiffIKActionCfg(
    asset_name="robot",
    joint_names=["panda_joint.*"],
    body_name="panda_link8",
    open_loop=False,
    # open_loop=True,
    controller=DifferentialIKControllerCfg(
        command_type="pose", 
        use_relative_mode=False, 
        ik_method="dls",
        ik_params={"lambda_val": 0.1},
    ),
    max_joint_vel=10.0,  # safety net only — catches singularity-induced blowup
)

DUAL_ARM_IK_ACTION = DualArmEnvFrameDiffIKActionCfg(
    left_asset_name="left_panda",
    right_asset_name="right_panda",
    joint_names=["panda_joint.*"],
    body_name="panda_link8",
    open_loop=False,
    # qp_net_path="/home/arhanjain/projects/sim-improvement/src/openpi/checkpoints/qp_net/latest.pt",
    controller=DifferentialIKControllerCfg(
        command_type="pose",
        use_relative_mode=False,
        ik_method="dls",
        # ik_params={"lambda_val": 0.1},
    ),
    max_joint_vel=10.0,  # safety net only — catches singularity-induced blowup
)

GRIPPER_ACTION = JointPositionActionCfg(
    asset_name="robot",
    joint_names=["panda_finger_joint1"],
    scale=-0.5,
    use_default_offset=False,
)
# (0, 0.08) -> (0, -0.04)
BINARY_GRIPPER_ACTION =  AbsBinaryJointPositionActionCfg(
    asset_name="robot",
    joint_names=["panda_finger_joint1"],
    open_command_expr={"panda_finger_joint1": -0.055},
    close_command_expr={"panda_finger_joint1": 0.0},
    threshold=0.04,
    positive_threshold=True,
)

@configclass
class LBMDiffIKAction:
    arm = ARM_IK_ACTION
    # gripper = GRIPPER_ACTION
    gripper = BINARY_GRIPPER_ACTION


####### ACTION DEFINITIONS #######