"""
RM75-B comfortable workspace evaluator.

You provide:
    - URDF file path '/home/zl/Downloads/urdf_rm75/RM75-B.urdf'
    - your own IK solver inside solve_ik_user()

This script computes:
    - IK success rate
    - joint-limit comfort
    - manipulability
    - singularity / condition number score
    - final comfort score

Recommended install:
    pip install numpy scipy urdfpy pandas matplotlib tqdm

Optional:
    pip install plotly
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from tqdm import tqdm
from scipy.spatial.transform import Rotation as R
from urdfpy import URDF

import sys
from pathlib import Path

# 1. Get the absolute path of the directory containing this current script
current_dir = Path(__file__).resolve().parent

# 2. Get the parent (upper) directory
parent_dir = current_dir.parent

# 3. Add the parent directory to the system path
sys.path.insert(0, str(parent_dir))

from rm75_kine_qp import KinematicsSolver as kine_qp
from rm75_kine_rm import rm75_kine_api as kine_rm
from rm75_mjc import MuJoCoPositionController
from Robotic_Arm.rm_robot_interface import *

import time
from math import radians, degrees, pi, cos, sin

# Cartesian workspace grid, in meters.
# Adjust according to your robot placement and task.
X_RANGE = (-0.6, 0.6)
Y_RANGE = (-0.6, 0.6)
Z_RANGE = (0.0, 0.8)

GRID_RESOLUTION = 0.05  # 5 cm. Use 0.02 for finer but slower.

num_orientations =  120 

# Comfort thresholds
MIN_JOINT_MARGIN = 0.05       # 15% away from joint limits
MAX_CONDITION_NUMBER = 150.0
MIN_MANIPULABILITY_RATIO = 0.10

# Scoring weights
WEIGHT_IK_SUCCESS = 0.70
WEIGHT_JOINT_LIMIT = 0.10
WEIGHT_MANIPULABILITY = 0.1
WEIGHT_SINGULARITY = 0.1




#   pose expression of tool-tip in end-effector, x y z quatx quaty quatz quatw
#   load: kg, mass_center_x in ee frame: m, y, z, then last threes are for filling
tools_in_ee = {
    'scissor': np.array([[0.0, 0.0, 0.19, 0.0, 0.0, 0.0, 1.0],[0.66, 0.0, 0.0, 0.06, 0.0, 0.0, 0.0]],dtype=np.float64),
    'omnipic': np.array([[0.0, 0.0, 0.16, 0.0, 0.0, 0.0, 1.0],[0.43, 0.0, 0.0, 0.06, 0.0, 0.0, 0.0]],dtype=np.float64),
    'minisci': np.array([[0.0, 0.0, 0.19, 0.0, 0.0, 0.0, 1.0],[0.46, 0.0, 0.0, 0.06, 0.0, 0.0, 0.0]],dtype=np.float64),
    'v_minis': np.array([[0.0, 0.1, 0.1, -np.sqrt(2) * 0.5, 0.0, 0.0, np.sqrt(2) * 0.5],[0.46, 0.0, 0.0, 0.06, 0.0, 0.0, 0.0]],dtype=np.float64),
    'no_tool': np.array([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]],dtype=np.float64),
}

# joint limit
# ub = np.array([150.0, 110.0, 170.0, 130, 175.0, 125.0, 179.0]) / 180 * pi
# lb = np.array([-150.0, -30.0, -170.0, -130, -175.0, -125.0, -179.0]) / 180 * pi
#
#
ub = np.array([179.0, 129.0, 179.0, 134, 179.0, 127.0, 359.0])/180*pi
lb = -ub

tool_name = "minisci"

URDF_PATH = str(parent_dir) + '/urdf_rm75/RM75-B.urdf'
MESH_DIR = str(Path(URDF_PATH).parent)

#  -----------      rm75 qp based kine          ------------
robot_kine_qp = kine_qp(urdf_path=URDF_PATH, mesh_dir=MESH_DIR)
robot_kine_qp.add_tool_frames(tools_in_ee)
robot_kine_qp.cfg_j_limit(min_j=lb, max_j=ub, rad_flag=True)

#   ----------      rm75 official algorithm         -----------
robot_kine_rm = kine_rm()
robot_kine_rm.add_tool_frames(tools_in_ee)
robot_kine_rm.cfg_j_limit(min_j=lb, max_j=ub, rad_flag=True)



# ============================================================
# 1. USER SETTINGS
# ============================================================



BASE_LINK = "base_link"
TCP_LINK = "link_7"

JOINT_NAMES = [
    "joint_1",
    "joint_2",
    "joint_3",
    "joint_4",
    "joint_5",
    "joint_6",
    "joint_7",
]



# Numerical Jacobian settings
JACOBIAN_EPS = 1e-5



# ============================================================
# 2. TASK ORIENTATION SAMPLING
# ============================================================

def make_task_orientations(num_orientations=num_orientations, seed=1):
    """
    Random orientation sampling using RM's Euler convention:

        R = Rz @ Ry @ Rx

    Note:
        This samples Euler angles randomly.
        It is useful, but not perfectly uniform over SO(3).
    """

    rng = np.random.default_rng(seed)

    orientations = []

    for _ in range(num_orientations):
        rx = rng.uniform(-np.pi, np.pi)
        ry = rng.uniform(-np.pi / 2.0, np.pi / 2.0)
        rz = rng.uniform(-np.pi, np.pi)

        orientations.append([rx, ry, rz])

    return orientations

#
# def euler_angles_to_rotation_matrix(rx, ry, rz):
#     """
#     Official RM convention:
#         R = Rz @ Ry @ Rx
#     This matches scipy:
#         Rotation.from_euler("xyz", [rx, ry, rz]).as_matrix()
#     """
#     Rx = np.array([
#         [1, 0, 0],
#         [0, np.cos(rx), -np.sin(rx)],
#         [0, np.sin(rx),  np.cos(rx)]
#     ])
#
#     Ry = np.array([
#         [ np.cos(ry), 0, np.sin(ry)],
#         [0,           1, 0],
#         [-np.sin(ry), 0, np.cos(ry)]
#     ])
#
#     Rz = np.array([
#         [np.cos(rz), -np.sin(rz), 0],
#         [np.sin(rz),  np.cos(rz), 0],
#         [0,           0,          1]
#     ])
#
#     return Rz @ Ry @ Rx


# ============================================================
# 3. YOUR IK FUNCTION GOES HERE
# ============================================================

def solve_ik_user(target_position, target_rotation):
    """
    Replace this function with your own IK solver.

    Parameters
    ----------
    target_position : np.ndarray, shape (3,)
        Desired TCP position in base_link frame.

    target_rotation : np.ndarray, shape (3, 3)
        Desired TCP rotation matrix in base_link frame.

    Returns
    -------
    None
        If IK fails.

    or

    np.ndarray, shape (7,)
        One IK solution.

    or

    list[np.ndarray]
        Multiple IK solutions.

    Important:
        Joint order must be:
        [joint_1, joint_2, joint_3, joint_4, joint_5, joint_6, joint_7]
    """

    # ========================================================
    # INSERT YOUR IK CODE HERE
    # ========================================================
    initial_guess = [0.1] * 7

    ret_qp, q = robot_kine_qp.inverse_kinematics(target_position=target_position, target_rpy=target_rotation, initial_guess=initial_guess, tool=tool_name, max_iter=250)
    # print(f'---- with qp ik, ret_qp: {ret_qp}, q = {q}')
    if ret_qp == 0:
        return q


    ret_rm, q = robot_kine_rm.inverse_kinematics(target_position=target_position, target_rpy=target_rotation, initial_guess=initial_guess, tool=tool_name)
    # print(f'==== with rm ik, ret_rm: {ret_rm}, q = {q}')
    if ret_rm == 0:
        pose_rm = robot_kine_rm.forward_kinematics(joint_angles=q, tool=tool_name)
        # print(f'target position = {target_position}\ntarget_rpy = {target_rotation} \npose_rm = {pose_rm}')
        return q


    return None


# ============================================================
# 4. URDF / FK UTILITIES
# ============================================================

def load_robot_and_limits(urdf_path):
    robot = URDF.load(urdf_path)

    joints = []
    lower = []
    upper = []

    joint_map = {j.name: j for j in robot.joints}

    for name in JOINT_NAMES:
        joint = joint_map[name]
        joints.append(joint)

        if joint.limit is None:
            raise ValueError(f"Joint {name} has no limit in URDF.")

        lower.append(joint.limit.lower)
        upper.append(joint.limit.upper)

    lower = np.asarray(lower, dtype=float)
    upper = np.asarray(upper, dtype=float)

    return robot, lower, upper


# def q_to_cfg(q):
#     """
#     Convert joint vector to urdfpy FK config dictionary.
#     """
#     return {name: float(q[i]) for i, name in enumerate(JOINT_NAMES)}


# def fk_transform(robot, q):
#     """
#     Forward kinematics from base_link to TCP_LINK.
#
#     Returns
#     -------
#     T : np.ndarray, shape (4, 4)
#     """
#     cfg = q_to_cfg(q)
#     fk = robot.link_fk(cfg=cfg)
#     tcp_link = robot.link_map[TCP_LINK]
#     return fk[tcp_link]
#
#
# def fk_position(robot, q):
#     T = fk_transform(robot, q)
#     return T[:3, 3]


# ============================================================
# 5. COMFORT METRICS
# ============================================================

def is_within_joint_limits(q, lower, upper, tol=1e-8):
    q = np.asarray(q)
    return np.all(q >= lower - tol) and np.all(q <= upper + tol)


def joint_limit_score(q, lower, upper):
    """
    Score in [0, 1].
    1 means every joint is at center of its range.
    0 means at least one joint is at its limit.
    """

    q = np.asarray(q)
    mid = 0.5 * (lower + upper)
    half_range = 0.5 * (upper - lower)

    per_joint_score = 1.0 - np.abs(q - mid) / half_range
    per_joint_score = np.clip(per_joint_score, 0.0, 1.0)

    # Conservative: one bad joint makes the whole pose less comfortable.
    return float(np.min(per_joint_score))


def joint_margin(q, lower, upper):
    """
    Minimum normalized distance to joint limits.

    0.15 means the closest joint is 15% away from its limit.
    """
    q = np.asarray(q)
    margin_lower = (q - lower) / (upper - lower)
    margin_upper = (upper - q) / (upper - lower)
    margin = np.minimum(margin_lower, margin_upper)
    return float(np.min(margin))


def q_to_cfg(q):
    """
    Convert joint vector to urdfpy FK config dictionary.
    """
    return {name: float(q[i]) for i, name in enumerate(JOINT_NAMES)}


def fk_transform(robot, q):
    """
    Forward kinematics from base_link to TCP_LINK.

    Returns
    -------
    T : np.ndarray, shape (4, 4)
    """
    cfg = q_to_cfg(q)
    fk = robot.link_fk(cfg=cfg)
    tcp_link = robot.link_map[TCP_LINK]
    return fk[tcp_link]


def numerical_geometric_jacobian(robot, q, eps=1e-5):
    """
    Numerical 6D geometric-like Jacobian, shape (6, 7).

    Top 3 rows:
        linear velocity approximation

    Bottom 3 rows:
        angular velocity approximation as rotation-vector difference

    This is useful for manipulability and singularity checks.
    """
    q = np.asarray(q, dtype=float)
    n = len(q)

    J = np.zeros((6, n))

    T0 = fk_transform(robot, q)
    p0 = T0[:3, 3]
    R0 = T0[:3, :3]

    for i in range(n):
        q_plus = q.copy()
        q_minus = q.copy()

        q_plus[i] += eps
        q_minus[i] -= eps

        T_plus = fk_transform(robot, q_plus)
        T_minus = fk_transform(robot, q_minus)

        p_plus = T_plus[:3, 3]
        p_minus = T_minus[:3, 3]

        R_plus = T_plus[:3, :3]
        R_minus = T_minus[:3, :3]

        # Linear part
        J[:3, i] = (p_plus - p_minus) / (2.0 * eps)

        # Angular part
        # Relative rotation from minus to plus.
        dR = R_plus @ R_minus.T
        rotvec = R.from_matrix(dR).as_rotvec()
        J[3:, i] = rotvec / (2.0 * eps)

    return J


def manipulability_score_from_jacobian(J):
    """
    Yoshikawa-style manipulability.

    For a 6x7 Jacobian:
        w = sqrt(det(J J.T))

    To improve numerical robustness, compute from singular values.
    """
    singular_values = np.linalg.svd(J, compute_uv=False)

    # Product of singular values.
    # For a 6x7 Jacobian, there are 6 singular values.
    w = float(np.prod(singular_values))

    return w


def condition_number_from_jacobian(J, min_sigma=1e-9):
    singular_values = np.linalg.svd(J, compute_uv=False)
    sigma_max = np.max(singular_values)
    sigma_min = np.min(singular_values)

    if sigma_min < min_sigma:
        return np.inf

    return float(sigma_max / sigma_min)


def singularity_score(condition_number):
    """
    Score in [0, 1].
    Higher is better.

    condition_number = 1 is ideal.
    Very large means near singularity.
    """
    if not np.isfinite(condition_number):
        return 0.0

    return float(1.0 / condition_number)


# ============================================================
# 6. IK RESULT HANDLING
# ============================================================

def normalize_ik_solutions(ik_result):
    """
    Your IK returns:
        - None if failed
        - one list/array of 7 joint values if successful
    """
    if ik_result is None:
        return []

    q = np.asarray(ik_result, dtype=float).reshape(-1)

    if q.shape[0] != 7:
        return []

    return [q]


def evaluate_single_solution(robot, q, lower, upper):
    """
    Evaluate one IK solution.

    Returns a dictionary with metrics.
    """

    if q.shape[0] != 7:
        return None

    if not is_within_joint_limits(q, lower, upper):
        return None

    jl_score = joint_limit_score(q, lower, upper)
    jl_margin = joint_margin(q, lower, upper)

    J = numerical_geometric_jacobian(robot, q, eps=JACOBIAN_EPS)

    manip = manipulability_score_from_jacobian(J)
    cond = condition_number_from_jacobian(J)
    sing_score = singularity_score(cond)

    valid_by_thresholds = (
        jl_margin >= MIN_JOINT_MARGIN
        and cond <= MAX_CONDITION_NUMBER
    )

    return {
        "q": q,
        "joint_limit_score": jl_score,
        "joint_margin": jl_margin,
        "manipulability": manip,
        "condition_number": cond,
        "singularity_score": sing_score,
        "valid_by_thresholds": valid_by_thresholds,
    }


# ============================================================
# 7. MAIN WORKSPACE EVALUATION
# ============================================================

def make_grid():
    xs = np.arange(X_RANGE[0], X_RANGE[1] + 1e-9, GRID_RESOLUTION)
    ys = np.arange(Y_RANGE[0], Y_RANGE[1] + 1e-9, GRID_RESOLUTION)
    zs = np.arange(Z_RANGE[0], Z_RANGE[1] + 1e-9, GRID_RESOLUTION)

    points = []

    for x in xs:
        for y in ys:
            for z in zs:
                points.append(np.array([x, y, z], dtype=float))

    return points


def evaluate_workspace():
    robot, lower, upper = load_robot_and_limits(URDF_PATH)
    orientations = make_task_orientations()
    grid_points = make_grid()

    rows = []

    # First pass stores raw manipulability.
    # Later we normalize manipulability by max observed value.
    all_valid_solution_metrics = []

    print(f"Loaded robot from: {URDF_PATH}")
    print(f"Grid points: {len(grid_points)}")
    print(f"Orientations per point: {len(orientations)}")
    print("Evaluating IK reachability and raw metrics...")

    for point in tqdm(grid_points):
        point_solution_metrics = []

        attempted = 0
        ik_success_count = 0


        for rpy in orientations:
            attempted += 1


            ik_result = solve_ik_user(point, rpy)

            # print(f'\n point is {point}, rpy is {rpy}, and ik result q: {ik_result}')
            candidate_solutions = normalize_ik_solutions(ik_result)

            if len(candidate_solutions) == 0:
                continue

            evaluated_solutions = []

            for q in candidate_solutions:
                # pose = robot_kine_qp.forward_kinematics(joint_angles=q, tool=tool_name)
                # print(f'the fk of q is {pose}\n')
                metrics = evaluate_single_solution(robot, q, lower, upper)
                # print(f'matrics: {metrics}, q = {q}, lower = {lower}, upper = {upper}')
                if metrics is not None:
                    evaluated_solutions.append(metrics)

            if len(evaluated_solutions) == 0:
                continue

            ik_success_count += 1

            # Use the best solution for this pose.
            # At this stage, manipulability is not normalized,
            # so use joint score + singularity score as temporary ranking.
            best = max(
                evaluated_solutions,
                key=lambda m: 0.6 * m["joint_limit_score"] + 0.4 * m["singularity_score"]
            )

            point_solution_metrics.append(best)
            all_valid_solution_metrics.append(best)
        print(f'this position+all orientations, the point_solution_metrics = {point_solution_metrics}')
        ik_success_rate = ik_success_count / attempted if attempted > 0 else 0.0

        if len(point_solution_metrics) == 0:
            rows.append({
                "x": point[0],
                "y": point[1],
                "z": point[2],
                "ik_success_rate": 0.0,
                "joint_limit_score": 0.0,
                "joint_margin": 0.0,
                "manipulability": 0.0,
                "manipulability_score": 0.0,
                "condition_number": np.inf,
                "singularity_score": 0.0,
                "comfort_score": 0.0,
                "comfortable": False,
                "reachable": False,
            })
        else:
            # Average over task orientations.
            rows.append({
                "x": point[0],
                "y": point[1],
                "z": point[2],
                "ik_success_rate": ik_success_rate,
                "joint_limit_score": np.mean([m["joint_limit_score"] for m in point_solution_metrics]),
                "joint_margin": np.mean([m["joint_margin"] for m in point_solution_metrics]),
                "manipulability": np.mean([m["manipulability"] for m in point_solution_metrics]),
                "manipulability_score": 0.0,  # filled later
                "condition_number": np.mean([m["condition_number"] for m in point_solution_metrics]),
                "singularity_score": np.mean([m["singularity_score"] for m in point_solution_metrics]),
                "comfort_score": 0.0,  # filled later
                "comfortable": False,
                "reachable": True,
            })

    df = pd.DataFrame(rows)

    # Normalize manipulability by maximum observed value.
    max_manip = df["manipulability"].replace([np.inf, -np.inf], np.nan).max()

    if max_manip is None or not np.isfinite(max_manip) or max_manip <= 0:
        max_manip = 1.0

    df["manipulability_score"] = df["manipulability"] / max_manip
    df["manipulability_score"] = df["manipulability_score"].clip(0.0, 1.0)

    # Final comfort score.
    df["comfort_score"] = (
        WEIGHT_IK_SUCCESS * df["ik_success_rate"]
        + WEIGHT_JOINT_LIMIT * df["joint_limit_score"]
        + WEIGHT_MANIPULABILITY * df["manipulability_score"]
        + WEIGHT_SINGULARITY * df["singularity_score"]
    )

    # Comfortable binary classification.
    df["comfortable"] = (
        (df["reachable"] == True)
        & (df["ik_success_rate"] >= 0.80)
        & (df["joint_margin"] >= MIN_JOINT_MARGIN)
        & (df["condition_number"] <= MAX_CONDITION_NUMBER)
        & (df["manipulability_score"] >= MIN_MANIPULABILITY_RATIO)
    )

    return df


# ============================================================
# 8. PLOTTING
# ============================================================

def plot_workspace(df):
    """
    3D scatter plot:
        gray/low = low comfort
        brighter = higher comfort
    """

    reachable = df[df["reachable"] == True]

    if len(reachable) == 0:
        print("No reachable points found. Check your IK function.")
        return

    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")

    sc = ax.scatter(
        reachable["x"],
        reachable["y"],
        reachable["z"],
        c=reachable["comfort_score"],
        s=12,
        alpha=0.8,
    )

    ax.set_title("RM75-B Comfortable Workspace")
    ax.set_xlabel("X [m]")
    ax.set_ylabel("Y [m]")
    ax.set_zlabel("Z [m]")

    fig.colorbar(sc, ax=ax, label="Comfort score")
    plt.show()


def plot_comfortable_only(df):
    comfortable = df[df["comfortable"] == True]

    if len(comfortable) == 0:
        print("No comfortable points found under current thresholds.")
        return

    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")

    ax.scatter(
        comfortable["x"],
        comfortable["y"],
        comfortable["z"],
        c=comfortable["comfort_score"],
        s=16,
        alpha=0.9,
    )

    ax.set_title("RM75-B Comfortable Region Only")
    ax.set_xlabel("X [m]")
    ax.set_ylabel("Y [m]")
    ax.set_zlabel("Z [m]")

    plt.show()


# ============================================================
# 9. ENTRY POINT
# ============================================================

if __name__ == "__main__":
    df = evaluate_workspace()

    output_csv = "rm75b_comfort_workspace.csv"
    df.to_csv(output_csv, index=False)

    print(f"\nSaved result to: {output_csv}")

    print("\nSummary:")
    print(f"Total grid points:       {len(df)}")
    print(f"Reachable points:        {df['reachable'].sum()}")
    print(f"Comfortable points:      {df['comfortable'].sum()}")

    if df["reachable"].sum() > 0:
        print(f"Max comfort score:       {df['comfort_score'].max():.3f}")
        print(f"Mean comfort score:      {df[df['reachable']]['comfort_score'].mean():.3f}")

    plot_workspace(df)
    plot_comfortable_only(df)