# Copyright [2021-2025] Thanh Nguyen
# Copyright [2022-2023] [CNRS, Toward SAS]

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

# http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import pinocchio as pin
import yaml
from yaml.loader import SafeLoader
from os.path import abspath
import matplotlib.pyplot as plt
import logging
from scipy.optimize import least_squares
from abc import ABC

import pandas as pd
from ..tools.regressor import eliminate_non_dynaffect
from ..tools.qrdecomposition import (
    get_baseParams,
    get_baseIndex,
    build_baseRegressor,
)

TOL_QR = 1e-8
FULL_PARAMTPL = ["d_px", "d_py", "d_pz", "d_phix", "d_phiy", "d_phiz"]
JOINT_OFFSETTPL = [
    "offsetPX",
    "offsetPY",
    "offsetPZ",
    "offsetRX",
    "offsetRY",
    "offsetRZ",
]
ELAS_TPL = [
    "k_PX",
    "k_PY",
    "k_PZ",
    "k_RX",
    "k_RY",
    "k_RZ",
]  # ONLY REVOLUTE JOINT FOR NOW
EE_TPL = ["pEEx", "pEEy", "pEEz", "phiEEx", "phiEEy", "phiEEz"]
BASE_TPL = [
    "base_px",
    "base_py",
    "base_pz",
    "base_phix",
    "base_phiy",
    "base_phiz",
]


# INITIALIZATION TOOLS
def get_param_from_yaml(robot, calib_data) -> dict:
    """Parse calibration parameters from YAML configuration file.

    Processes robot and calibration data to build a parameter dictionary containing
    all necessary settings for robot calibration. Handles configuration of:
    - Frame identifiers and relationships
    - Marker/measurement settings
    - Joint indices and configurations
    - Non-geometric parameters
    - Eye-hand calibration setup

    Args:
        robot (pin.RobotWrapper): Robot instance containing model and data
        calib_data (dict): Calibration parameters parsed from YAML file containing:
            - markers: List of marker configurations
            - calib_level: Calibration model type
            - base_frame: Starting frame name
            - tool_frame: End frame name
            - free_flyer: Whether base is floating
            - non_geom: Whether to include non-geometric params

    Returns:
        dict: Parameter dictionary containing:
            - robot_name: Name of robot model
            - NbMarkers: Number of markers
            - measurability: Measurement DOFs per marker
            - start_frame, end_frame: Frame names
            - base_to_ref_frame: Optional camera frame
            - IDX_TOOL: Tool frame index
            - actJoint_idx: Active joint indices
            - param_name: List of parameter names
            - Additional settings from YAML

    Side Effects:
        Prints warning messages if optional frames undefined
        Prints final parameter dictionary

    Example:
        >>> calib_data = yaml.safe_load(config_file)
        >>> params = get_param_from_yaml(robot, calib_data)
        >>> print(params['NbMarkers'])
        2
    """
    # NOTE: since joint 0 is universe and it is trivial,
    # indices of joints are different from indices of joint configuration,
    # different from indices of joint velocities
    calib_config = dict()
    robot_name = robot.model.name
    frames = [f.name for f in robot.model.frames]
    calib_config["robot_name"] = robot_name

    # End-effector sensing measurability:
    NbMarkers = len(calib_data["markers"])
    measurability = calib_data["markers"][0]["measure"]
    calib_idx = measurability.count(True)
    calib_config["NbMarkers"] = NbMarkers
    calib_config["measurability"] = measurability
    calib_config["calibration_index"] = calib_idx

    # Calibration model
    calib_config["calib_model"] = calib_data["calib_level"]

    # Get start and end frames
    start_frame = calib_data["base_frame"]
    end_frame = calib_data["tool_frame"]

    # Validate frames exist
    err_msg = "{}_frame {} does not exist"
    if start_frame not in frames:
        raise AssertionError(err_msg.format("Start", start_frame))
    if end_frame not in frames:
        raise AssertionError(err_msg.format("End", end_frame))

    calib_config["start_frame"] = start_frame
    calib_config["end_frame"] = end_frame

    # Handle eye-hand calibration frames
    try:
        base_to_ref_frame = calib_data["base_to_ref_frame"]
        ref_frame = calib_data["ref_frame"]
    except KeyError:
        base_to_ref_frame = None
        ref_frame = None
        print("base_to_ref_frame and ref_frame are not defined.")

    # Validate base-to-ref frame if provided
    if base_to_ref_frame is not None:
        if base_to_ref_frame not in frames:
            err_msg = "base_to_ref_frame {} does not exist"
            raise AssertionError(err_msg.format(base_to_ref_frame))

    # Validate ref frame if provided
    if ref_frame is not None:
        if ref_frame not in frames:
            err_msg = "ref_frame {} does not exist"
            raise AssertionError(err_msg.format(ref_frame))

    calib_config["base_to_ref_frame"] = base_to_ref_frame
    calib_config["ref_frame"] = ref_frame

    # Get initial poses
    try:
        base_pose = calib_data["base_pose"]
        tip_pose = calib_data["tip_pose"]
    except KeyError:
        base_pose = None
        tip_pose = None
        print("base_pose and tip_pose are not defined.")

    calib_config["base_pose"] = base_pose
    calib_config["tip_pose"] = tip_pose

    # q0: default zero configuration
    calib_config["q0"] = robot.q0
    calib_config["NbSample"] = calib_data["nb_sample"]

    # IDX_TOOL: frame ID of the tool
    IDX_TOOL = robot.model.getFrameId(end_frame)
    calib_config["IDX_TOOL"] = IDX_TOOL

    # tool_joint: ID of the joint right before the tool's frame (parent)
    tool_joint = robot.model.frames[IDX_TOOL].parent
    calib_config["tool_joint"] = tool_joint

    # indices of active joints: from base to tool_joint
    actJoint_idx = get_sup_joints(robot.model, start_frame, end_frame)
    calib_config["actJoint_idx"] = actJoint_idx

    # indices of joint configuration corresponding to active joints
    config_idx = [robot.model.joints[i].idx_q for i in actJoint_idx]
    calib_config["config_idx"] = config_idx

    # number of active joints
    NbJoint = len(actJoint_idx)
    calib_config["NbJoint"] = NbJoint

    # initialize a list of calibrating parameters name
    param_name = []
    if calib_data["non_geom"]:
        # list of elastic gain parameter names
        elastic_gain = []
        joint_axes = ["PX", "PY", "PZ", "RX", "RY", "RZ"]
        for j_id, joint_name in enumerate(robot.model.names.tolist()):
            if joint_name == "universe":
                axis_motion = "null"
            else:
                # for ii, ax in enumerate(AXIS_MOTION[j_id]):
                #     if ax == 1:
                #         axis_motion = axis[ii]
                shortname = robot.model.joints[
                    j_id
                ].shortname()  # ONLY TAKE PRISMATIC AND REVOLUTE JOINT
                for ja in joint_axes:
                    if ja in shortname:
                        axis_motion = ja
                    elif "RevoluteUnaligned" in shortname:
                        axis_motion = "RZ"  # hard coded fix for canopies

            elastic_gain.append("k_" + axis_motion + "_" + joint_name)
        for i in actJoint_idx:
            param_name.append(elastic_gain[i])
    calib_config["param_name"] = param_name

    calib_config.update(
        {
            "free_flyer": calib_data["free_flyer"],
            "non_geom": calib_data["non_geom"],
            "eps": 1e-3,
            "PLOT": 0,
        }
    )
    try:
        calib_config.update(
            {
                "coeff_regularize": calib_data["coeff_regularize"],
                "data_file": calib_data["data_file"],
                "sample_configs_file": calib_data["sample_configs_file"],
                "outlier_eps": calib_data["outlier_eps"],
            }
        )
    except KeyError:
        calib_config.update(
            {
                "coeff_regularize": None,
                "data_file": None,
                "sample_configs_file": None,
                "outlier_eps": None,
            }
        )
    return calib_config


def unified_to_legacy_config(robot, unified_calib_config) -> dict:
    """Convert unified configuration format to legacy calib_config format.
    
    Maps the new unified configuration structure to the exact format expected
    by get_param_from_yaml. This ensures backward compatibility while using
    the new unified parser.
    
    Args:
        robot (pin.RobotWrapper): Robot instance containing model and data
        unified_calib_config (dict): Configuration from create_task_config
        
    Returns:
        dict: Legacy format calibration configuration matching
              get_param_from_yaml output
        
    Raises:
        KeyError: If required fields are missing from unified config
        
    Example:
        >>> unified_config = create_task_config(robot, parsed_config,
        ...                                    "calibration")
        >>> legacy_config = unified_to_legacy_config(robot, unified_config)
    """
    # Extract basic robot information
    calib_config = dict()
    robot_name = robot.model.name
    frames = [f.name for f in robot.model.frames]
    calib_config["robot_name"] = robot_name
    
    # Extract markers and measurements information
    markers = unified_calib_config.get("measurements", {}).get("markers", [{}])
    if not markers:
        raise KeyError("No markers defined in unified configuration")
    
    first_marker = markers[0]
    NbMarkers = len(markers)
    measurability = first_marker.get(
        "measurable_dof", [True, True, True, False, False, False]
    )
    calib_idx = measurability.count(True)
    
    calib_config["NbMarkers"] = NbMarkers
    calib_config["measurability"] = measurability
    calib_config["calibration_index"] = calib_idx
    
    # Extract calibration model
    parameters = unified_calib_config.get("parameters", {})
    calib_config["calib_model"] = parameters.get(
        "calibration_level", "full_params"
    )
    
    # Extract kinematic frames
    kinematics = unified_calib_config.get("kinematics", {})
    start_frame = kinematics.get("base_frame", "universe")
    end_frame = kinematics.get("tool_frame")
    
    if not end_frame:
        raise KeyError("tool_frame not specified in unified configuration")
    
    # Validate frames exist
    err_msg = "{}_frame {} does not exist"
    if start_frame not in frames:
        raise AssertionError(err_msg.format("Start", start_frame))
    if end_frame not in frames:
        raise AssertionError(err_msg.format("End", end_frame))
    
    calib_config["start_frame"] = start_frame
    calib_config["end_frame"] = end_frame
    
    # Handle eye-hand calibration frames (optional)
    calib_config["base_to_ref_frame"] = unified_calib_config.get(
        "base_to_ref_frame"
    )
    calib_config["ref_frame"] = unified_calib_config.get("ref_frame")
    if not calib_config["base_to_ref_frame"] and not calib_config["ref_frame"]:
        print("base_to_ref_frame and ref_frame are not defined.")
    
    # Extract poses
    poses = unified_calib_config.get("measurements", {}).get("poses", {})
    calib_config["base_pose"] = poses.get("base_pose")
    calib_config["tip_pose"] = poses.get("tool_pose")
    if not calib_config["base_pose"] and not calib_config["tip_pose"]:
        print("base_pose and tip_pose are not defined.")
    
    # Robot configuration
    calib_config["q0"] = robot.q0
    
    # Extract data information
    data_info = unified_calib_config.get("data", {})
    calib_config["NbSample"] = data_info.get("number_of_samples", 500)
    
    # Tool frame information
    IDX_TOOL = robot.model.getFrameId(end_frame)
    calib_config["IDX_TOOL"] = IDX_TOOL
    tool_joint = robot.model.frames[IDX_TOOL].parent
    calib_config["tool_joint"] = tool_joint
    
    # Active joints - extract from joints.active_joints if available
    joints_info = unified_calib_config.get("joints", {})
    active_joint_names = joints_info.get("active_joints", [])
    
    if active_joint_names:
        # Map joint names to indices
        actJoint_idx = []
        for joint_name in active_joint_names:
            if joint_name in robot.model.names:
                joint_id = robot.model.getJointId(joint_name)
                actJoint_idx.append(joint_id)
        calib_config["actJoint_idx"] = actJoint_idx
    else:
        # Fall back to computing from frames if no active joints specified
        actJoint_idx = get_sup_joints(robot.model, start_frame, end_frame)
        calib_config["actJoint_idx"] = actJoint_idx
    
    # Configuration indices
    config_idx = [robot.model.joints[i].idx_q for i in
                  calib_config["actJoint_idx"]]
    calib_config["config_idx"] = config_idx
    
    # Number of active joints
    NbJoint = len(calib_config["actJoint_idx"])
    calib_config["NbJoint"] = NbJoint
    
    # Parameter names (initialize empty, will be filled by create_param_list)
    param_name = []
    
    # Handle non-geometric parameters
    non_geom = parameters.get("include_non_geometric", False)
    if non_geom:
        # Create elastic gain parameter names
        elastic_gain = []
        joint_axes = ["PX", "PY", "PZ", "RX", "RY", "RZ"]
        for j_id, joint_name in enumerate(robot.model.names.tolist()):
            if joint_name == "universe":
                axis_motion = "null"
            else:
                shortname = robot.model.joints[j_id].shortname()
                for ja in joint_axes:
                    if ja in shortname:
                        axis_motion = ja
                    elif "RevoluteUnaligned" in shortname:
                        axis_motion = "RZ"  # hard coded fix for canopies
            elastic_gain.append("k_" + axis_motion + "_" + joint_name)
        for i in calib_config["actJoint_idx"]:
            param_name.append(elastic_gain[i])
    
    calib_config["param_name"] = param_name
    
    # Basic calibration settings
    calib_config.update({
        "free_flyer": parameters.get("free_flyer", False),
        "non_geom": non_geom,
        "eps": 1e-3,
        "PLOT": 0,
    })
    
    # Optional parameters with defaults
    calib_config["coeff_regularize"] = parameters.get(
        "regularization_coefficient", 0.01
    )
    calib_config["data_file"] = data_info.get("source_file")
    calib_config["sample_configs_file"] = data_info.get(
        "sample_configurations_file"
    )
    calib_config["outlier_eps"] = parameters.get("outlier_threshold", 0.05)
    
    return calib_config


def get_joint_offset(model, joint_names):
    """Get dictionary of joint offset parameters.

    Maps joint names to their offset parameters, handling special cases for
    different joint types and multiple DOF joints.

    Args:
        model: Pinocchio robot model
        joint_names: List of joint names from model.names

    Returns:
        dict: Mapping of joint offset parameter names to initial zero values.
            Keys have format: "{offset_type}_{joint_name}"

    Example:
        >>> offsets = get_joint_offset(robot.model, robot.model.names[1:])
        >>> print(offsets["offsetRZ_joint1"])
        0.0
    """
    joint_off = []
    joint_names = list(model.names[1:])
    joints = list(model.joints[1:])
    assert len(joint_names) == len(
        joints
    ), "Number of jointnames does not match number of joints! Please check\
        imported model."
    for id, joint in enumerate(joints):
        name = joint_names[id]
        shortname = joint.shortname()
        if model.name == "canopies":
            if "RevoluteUnaligned" in shortname:
                shortname = shortname.replace("RevoluteUnaligned", "RZ")
        for i in range(joint.nv):
            if i > 0:
                offset_param = (
                    shortname.replace("JointModel", "offset")
                    + "{}".format(i + 1)
                    + "_"
                    + name
                )
            else:
                offset_param = (
                    shortname.replace("JointModel", "offset") + "_" + name
                )
            joint_off.append(offset_param)

    phi_jo = [0] * len(joint_off)  # default zero values
    joint_off = dict(zip(joint_off, phi_jo))
    return joint_off


def get_fullparam_offset(joint_names):
    """Get dictionary of geometric parameter variations.

    Creates mapping of geometric offset parameters for each joint's
    position and orientation.

    Args:
        joint_names: List of joint names from robot model

    Returns:
        dict: Mapping of geometric parameter names to initial zero values.
            Keys have format: "d_{param}_{joint_name}" where param is:
            - px, py, pz: Position offsets
            - phix, phiy, phiz: Orientation offsets

    Example:
        >>> geo_params = get_fullparam_offset(robot.model.names[1:])
        >>> print(geo_params["d_px_joint1"])
        0.0
    """
    geo_params = []

    for i in range(len(joint_names)):
        for j in FULL_PARAMTPL:
            # geo_params.append(j + ("_%d" % i))
            geo_params.append(j + "_" + joint_names[i])

    phi_gp = [0] * len(geo_params)  # default zero values
    geo_params = dict(zip(geo_params, phi_gp))
    return geo_params


def add_base_name(calib_config):
    """Add base frame parameters to parameter list.

    Updates calib_config["param_name"] with base frame parameters depending on
    calibration model type.

    Args:
        calib_config: Parameter dictionary containing:
            - calib_model: "full_params" or "joint_offset"
            - param_name: List of parameter names to update

    Side Effects:
        Modifies calib_config["param_name"] in place by:
        - For full_params: Replaces first 6 entries with base parameters
        - For joint_offset: Prepends base parameters to list
    """
    if calib_config["calib_model"] == "full_params":
        calib_config["param_name"][0:6] = BASE_TPL
    elif calib_config["calib_model"] == "joint_offset":
        calib_config["param_name"] = BASE_TPL + calib_config["param_name"]


def add_pee_name(calib_config):
    """Add end-effector marker parameters to parameter list.

    Adds parameters for each active measurement DOF of each marker.

    Args:
        calib_config: Parameter dictionary containing:
            - NbMarkers: Number of markers
            - measurability: List of booleans for active DOFs
            - param_name: List of parameter names to update

    Side Effects:
        Modifies calib_config["param_name"] in place by appending marker parameters
        in format: "{param_type}_{marker_num}"
    """
    PEE_names = []
    for i in range(calib_config["NbMarkers"]):
        for j, state in enumerate(calib_config["measurability"]):
            if state:
                PEE_names.extend(["{}_{}".format(EE_TPL[j], i + 1)])
    calib_config["param_name"] = calib_config["param_name"] + PEE_names


def add_eemarker_frame(frame_name, p, rpy, model, data):
    """Add a new frame attached to the end-effector.

    Creates and adds a fixed frame to the robot model at the end-effector location,
    typically used for marker or tool frames.

    Args:
        frame_name (str): Name for the new frame
        p (ndarray): 3D position offset from parent frame
        rpy (ndarray): Roll-pitch-yaw angles for frame orientation
        model (pin.Model): Robot model to add frame to
        data (pin.Data): Robot data structure

    Returns:
        int: ID of newly created frame

    Note:
        Currently hardcoded to attach to "arm_7_joint". This should be made
        configurable in future versions.
    """
    p = np.array([0.1, 0.1, 0.1])
    R = pin.rpy.rpyToMatrix(rpy)
    frame_placement = pin.SE3(R, p)

    parent_jointId = model.getJointId("arm_7_joint")
    prev_frameId = model.getFrameId("arm_7_joint")
    ee_frame_id = model.addFrame(
        pin.Frame(
            frame_name,
            parent_jointId,
            prev_frameId,
            frame_placement,
            pin.FrameType(0),
            pin.Inertia.Zero(),
        ),
        False,
    )
    return ee_frame_id


# DATA PROCESSING TOOLS
# data collected and saved in various formats and structure, it has
# always been a pain in the ass to handles. Need to think a way to simplify
# and optimize this procedure.
def read_config_data(model, path_to_file):
    """Read joint configurations from CSV file.

    Args:
        model (pin.Model): Robot model containing joint information
        path_to_file (str): Path to CSV file containing joint configurations

    Returns:
        ndarray: Matrix of shape (n_samples, n_joints-1) containing joint positions
    """
    df = pd.read_csv(path_to_file)
    q = np.zeros([len(df), model.njoints - 1])
    for i in range(len(df)):
        for j, name in enumerate(model.names[1:].tolist()):
            jointidx = get_idxq_from_jname(model, name)
            q[i, jointidx] = df[name][i]
    return q


def load_data(path_to_file, model, calib_config, del_list=[]):
    """Load joint configuration and marker data from CSV file.

    Reads marker positions/orientations and joint configurations from a CSV file.
    Handles data validation, bad sample removal, and conversion to numpy arrays.

    Args:
        path_to_file (str): Path to CSV file containing recorded data
        model (pin.Model): Robot model containing joint information
        calib_config (dict): Parameter dictionary containing:
            - NbMarkers: Number of markers to load
            - measurability: List indicating which DOFs are measured
            - actJoint_idx: List of active joint indices
            - config_idx: Configuration vector indices
            - q0: Default configuration vector
        del_list (list, optional): Indices of bad samples to remove. Defaults to [].

    Returns:
        tuple:
            - PEEm_exp (ndarray): Flattened marker measurements of shape (n_active_dofs,)
            - q_exp (ndarray): Joint configurations of shape (n_samples, n_joints)

    Note:
        CSV file must contain columns:
        - For each marker i: [xi, yi, zi, phixi, phiyi, phizi]
        - Joint names matching model.names for active joints

    Raises:
        KeyError: If required columns are missing from CSV

    Side Effects:
        - Prints joint headers
        - Updates calib_config["NbSample"] with number of valid samples
    """
    # read_csv
    df = pd.read_csv(path_to_file)

    # create headers for marker position
    PEE_headers = []
    pee_tpl = ["x", "y", "z", "phix", "phiy", "phiz"]
    for i in range(calib_config["NbMarkers"]):
        for j, state in enumerate(calib_config["measurability"]):
            if state:
                PEE_headers.extend(["{}{}".format(pee_tpl[j], i + 1)])

    # create headers for joint configurations
    joint_headers = [model.names[i] for i in calib_config["actJoint_idx"]]
    
    # check if all created headers present in csv file
    csv_headers = list(df.columns)
    for header in PEE_headers + joint_headers:
        if header not in csv_headers:
            print("%s does not exist in the file." % header)
            break

    # Extract marker position/location
    pose_ee = df[PEE_headers].to_numpy()

    # Extract joint configurations
    q_act = df[joint_headers].to_numpy()

    # remove bad data
    if del_list:
        pose_ee = np.delete(pose_ee, del_list, axis=0)
        q_act = np.delete(q_act, del_list, axis=0)

    # update number of data points
    calib_config["NbSample"] = q_act.shape[0]

    PEEm_exp = pose_ee.T
    PEEm_exp = PEEm_exp.flatten("C")

    q_exp = np.empty((calib_config["NbSample"], calib_config["q0"].shape[0]))
    for i in range(calib_config["NbSample"]):
        config = calib_config["q0"]
        config[calib_config["config_idx"]] = q_act[i, :]
        q_exp[i, :] = config

    return PEEm_exp, q_exp


# COMMON TOOLS


def get_idxq_from_jname(model, joint_name):
    """Get index of joint in configuration vector.

    Args:
        model (pin.Model): Robot model
        joint_name (str): Name of joint to find index for

    Returns:
        int: Index of joint in configuration vector q

    Raises:
        AssertionError: If joint name does not exist in model
    """
    assert joint_name in model.names, "Given joint name does not exist."
    jointId = model.getJointId(joint_name)
    joint_idx = model.joints[jointId].idx_q
    return joint_idx


def cartesian_to_SE3(X):
    """Convert cartesian coordinates to SE3 transformation.

    Args:
        X (ndarray): (6,) array with [x,y,z,rx,ry,rz] coordinates

    Returns:
        pin.SE3: SE3 transformation with:
            - translation from X[0:3]
            - rotation matrix from RPY angles X[3:6]
    """
    X = np.array(X)
    X = X.flatten("C")
    translation = X[0:3]
    rot_matrix = pin.rpy.rpyToMatrix(X[3:6])
    placement = pin.SE3(rot_matrix, translation)
    return placement


def xyzquat_to_SE3(xyzquat):
    """Convert XYZ position and quaternion orientation to SE3 transformation.

    Takes a 7D vector containing XYZ position and WXYZ quaternion and creates
    an SE3 transformation matrix.

    Args:
        xyzquat (ndarray): (7,) array containing:
            - xyzquat[0:3]: XYZ position coordinates
            - xyzquat[3:7]: WXYZ quaternion orientation

    Returns:
        pin.SE3: Rigid body transformation with:
            - Translation from XYZ position
            - Rotation matrix from normalized quaternion

    Example:
        >>> pos_quat = np.array([0.1, 0.2, 0.3, 1.0, 0, 0, 0])
        >>> transform = xyzquat_to_SE3(pos_quat)
    """
    xyzquat = np.array(xyzquat)
    xyzquat = xyzquat.flatten("C")
    translation = xyzquat[0:3]
    rot_matrix = pin.Quaternion(xyzquat[3:7]).normalize().toRotationMatrix()
    placement = pin.SE3(rot_matrix, translation)
    return placement


def get_rel_transform(model, data, start_frame, end_frame):
    """Get relative transformation between two frames.

    Calculates the transform from start_frame to end_frame in the kinematic chain.
    Assumes forward kinematics has been updated.

    Args:
        model (pin.Model): Robot model
        data (pin.Data): Robot data
        start_frame (str): Starting frame name
        end_frame (str): Target frame name

    Returns:
        pin.SE3: Relative transformation sMt from start to target frame

    Raises:
        AssertionError: If frame names don't exist in model
    """
    frames = [f.name for f in model.frames]
    assert start_frame in frames, "{} does not exist.".format(start_frame)
    assert end_frame in frames, "{} does not exist.".format(end_frame)
    start_frameId = model.getFrameId(start_frame)
    oMsf = data.oMf[start_frameId]
    end_frameId = model.getFrameId(end_frame)
    oMef = data.oMf[end_frameId]
    sMef = oMsf.actInv(oMef)
    return sMef


def get_sup_joints(model, start_frame, end_frame):
    """Get list of supporting joints between two frames in kinematic chain.

    Finds all joints that contribute to relative motion between start_frame and
    end_frame by analyzing their support branches in the kinematic tree.

    Args:
        model (pin.Model): Robot model
        start_frame (str): Name of starting frame
        end_frame (str): Name of ending frame

    Returns:
        list[int]: Joint IDs ordered from start to end frame, handling cases:
            1. Branch entirely contained in another branch
            2. Disjoint branches with fixed root
            3. Partially overlapping branches

    Raises:
        AssertionError: If end frame appears before start frame in chain

    Note:
        Excludes "universe" joints from returned list since they don't
        contribute to relative motion.
    """
    start_frameId = model.getFrameId(start_frame)
    end_frameId = model.getFrameId(end_frame)
    start_par = model.frames[start_frameId].parent
    end_par = model.frames[end_frameId].parent
    branch_s = model.supports[start_par].tolist()
    branch_e = model.supports[end_par].tolist()
    # remove 'universe' joint from branches
    if model.names[branch_s[0]] == "universe":
        branch_s.remove(branch_s[0])
    if model.names[branch_e[0]] == "universe":
        branch_e.remove(branch_e[0])

    # find over-lapping joints in two branches
    shared_joints = list(set(branch_s) & set(branch_e))
    # create a list of supporting joints between two frames
    list_1 = [x for x in branch_s if x not in branch_e]
    list_1.reverse()
    list_2 = [y for y in branch_e if y not in branch_s]
    # case 2: root_joint is fixed; branch_s and branch_e are separate
    if shared_joints == []:
        sup_joints = list_1 + list_2
    else:
        # case 1: branch_s is part of branch_e
        if shared_joints == branch_s:
            sup_joints = [branch_s[-1]] + list_2
        else:
            assert (
                shared_joints != branch_e
            ), "End frame should be before start frame."
            # case 3: there are overlapping joints between two branches
            sup_joints = list_1 + [shared_joints[-1]] + list_2
    return sup_joints


def get_rel_kinreg(model, data, start_frame, end_frame, q):
    """Calculate relative kinematic regressor between frames.

    Computes frame Jacobian-based regressor matrix mapping small joint displacements
    to spatial velocities.

    Args:
        model (pin.Model): Robot model
        data (pin.Data): Robot data
        start_frame (str): Starting frame name
        end_frame (str): Target frame name
        q (ndarray): Joint configuration vector

    Returns:
        ndarray: (6, 6n) regressor matrix for n joints
    """
    sup_joints = get_sup_joints(model, start_frame, end_frame)
    pin.framesForwardKinematics(model, data, q)
    pin.updateFramePlacements(model, data)
    kinreg = np.zeros((6, 6 * (model.njoints - 1)))
    frame = model.frames[model.getFrameId(end_frame)]
    oMf = data.oMi[frame.parent] * frame.placement
    for p in sup_joints:
        oMp = data.oMi[model.parents[p]] * model.jointPlacements[p]
        fMp = oMf.actInv(oMp)
        fXp = fMp.toActionMatrix()
        kinreg[:, 6 * (p - 1) : 6 * p] = fXp
    return kinreg


def get_rel_jac(model, data, start_frame, end_frame, q):
    """Calculate relative Jacobian matrix between two frames.

    Computes the difference between Jacobians of end_frame and start_frame,
    giving the differential mapping from joint velocities to relative spatial velocity.

    Args:
        model (pin.Model): Robot model
        data (pin.Data): Robot data
        start_frame (str): Starting frame name
        end_frame (str): Target frame name
        q (ndarray): Joint configuration vector

    Returns:
        ndarray: (6, n) relative Jacobian matrix where:
            - Rows represent [dx,dy,dz,wx,wy,wz] spatial velocities
            - Columns represent joint velocities
            - n is number of joints

    Note:
        Updates forward kinematics before computing Jacobians
    """
    start_frameId = model.getFrameId(start_frame)
    end_frameId = model.getFrameId(end_frame)

    # update frameForwardKinematics and updateFramePlacements
    pin.framesForwardKinematics(model, data, q)
    pin.updateFramePlacements(model, data)

    # relative Jacobian
    J_start = pin.computeFrameJacobian(
        model, data, q, start_frameId, pin.LOCAL
    )
    J_end = pin.computeFrameJacobian(model, data, q, end_frameId, pin.LOCAL)
    J_rel = J_end - J_start
    return J_rel


# LEVENBERG-MARQUARDT TOOLS


def initialize_variables(calib_config, mode=0, seed=0):
    """Initialize variables for Levenberg-Marquardt optimization.

    Creates initial parameter vector either as zeros or random values within bounds.

    Args:
        calib_config (dict): Parameter dictionary containing:
            - param_name: List of parameter names to initialize
        mode (int, optional): Initialization mode:
            - 0: Zero initialization
            - 1: Random uniform initialization. Defaults to 0.
        seed (float, optional): Range [-seed,seed] for random init. Defaults to 0.

    Returns:
        tuple:
            - var (ndarray): Initial parameter vector
            - nvar (int): Number of parameters

    Example:
        >>> var, n = initialize_variables(params, mode=1, seed=0.1)
        >>> print(var.shape)
        (42,)
    """
    # initialize all variables at zeros
    nvar = len(calib_config["param_name"])
    if mode == 0:
        var = np.zeros(nvar)
    elif mode == 1:
        var = np.random.uniform(-seed, seed, nvar)
    return var, nvar


def update_forward_kinematics(model, data, var, q, calib_config, verbose=0):
    """Update forward kinematics with calibration parameters.

    Applies geometric and kinematic error parameters to update joint placements
    and compute end-effector poses. Handles:
    1. Base/camera transformations
    2. Joint placement offsets
    3. End-effector marker frames
    4. Joint elasticity effects

    Args:
        model (pin.Model): Robot model to update
        data (pin.Data): Robot data
        var (ndarray): Parameter vector matching calib_config["param_name"]
        q (ndarray): Joint configurations matrix (n_samples, n_joints)
        calib_config (dict): Calibration parameters containing:
            - calib_model: "full_params" or "joint_offset"
            - start_frame, end_frame: Frame names
            - actJoint_idx: Active joint indices
            - measurability: Active DOFs
        verbose (int, optional): Print update info. Defaults to 0.

    Returns:
        ndarray: Flattened end-effector measurements for all samples

    Side Effects:
        - Modifies model joint placements temporarily
        - Reverts model to original state before returning
    """
    # read calib_config['param_name'] to allocate offset parameters to correct SE3
    # convert translation: add a vector of 3 to SE3.translation
    # convert orientation: convert SE3.rotation 3x3 matrix to vector rpy, add
    #  to vector rpy, convert back to to 3x3 matrix

    # name reference of calibration parameters
    if calib_config["calib_model"] == "full_params":
        axis_tpl = FULL_PARAMTPL
    elif calib_config["calib_model"] == "joint_offset":
        axis_tpl = JOINT_OFFSETTPL

    # order of joint in variables are arranged as in calib_config['actJoint_idx']
    assert len(var) == len(
        calib_config["param_name"]
    ), "Length of variables != length of params"
    param_dict = dict(zip(calib_config["param_name"], var))
    origin_model = model.copy()

    # update model.jointPlacements
    updated_params = []
    start_f = calib_config["start_frame"]
    end_f = calib_config["end_frame"]

    # define transformation for camera frame
    if calib_config["base_to_ref_frame"] is not None:
        start_f = calib_config["ref_frame"]
        # base frame to ref frame (i.e. Tiago: camera transformation)
        base_tf = np.zeros(6)
        for key in param_dict.keys():
            for base_id, base_ax in enumerate(BASE_TPL):
                if base_ax in key:
                    base_tf[base_id] = param_dict[key]
                    updated_params.append(key)
        b_to_cam = get_rel_transform(
            model, data, calib_config["start_frame"], calib_config["base_to_ref_frame"]
        )
        ref_to_cam = cartesian_to_SE3(base_tf)
        cam_to_ref = ref_to_cam.actInv(pin.SE3.Identity())
        bMo = b_to_cam * cam_to_ref
    else:
        if calib_config["calib_model"] == "joint_offset":
            base_tf = np.zeros(6)
            for key in param_dict.keys():
                for base_id, base_ax in enumerate(BASE_TPL):
                    if base_ax in key:
                        base_tf[base_id] = param_dict[key]
                        updated_params.append(key)
            bMo = cartesian_to_SE3(base_tf)

    # update model.jointPlacements with joint 'full_params'/'joint_offset'
    for j_id in calib_config["actJoint_idx"]:
        xyz_rpy = np.zeros(6)
        j_name = model.names[j_id]
        for key in param_dict.keys():
            if j_name in key:
                # update xyz_rpy with kinematic errors
                for axis_id, axis in enumerate(axis_tpl):
                    if axis in key:
                        if verbose == 1:
                            print(
                                "Updating [{}] joint placement at axis {} with [{}]".format(
                                    j_name, axis, key
                                )
                            )
                        xyz_rpy[axis_id] += param_dict[key]
                        updated_params.append(key)
        model = update_joint_placement(model, j_id, xyz_rpy)
    PEE = np.zeros((calib_config["calibration_index"], calib_config["NbSample"]))

    # update end_effector frame
    for marker_idx in range(1, calib_config["NbMarkers"] + 1):
        pee = np.zeros(6)
        ee_name = "EE"
        for key in param_dict.keys():
            if ee_name in key and str(marker_idx) in key:
                # update xyz_rpy with kinematic errors
                for axis_pee_id, axis_pee in enumerate(EE_TPL):
                    if axis_pee in key:
                        if verbose == 1:
                            print(
                                "Updating [{}_{}] joint placement at axis {} with [{}]".format(
                                    ee_name, str(marker_idx), axis_pee, key
                                )
                            )
                        pee[axis_pee_id] += param_dict[key]
                        # updated_params.append(key)

        eeMf = cartesian_to_SE3(pee)

    # get transform
    q_ = np.copy(q)
    for i in range(calib_config["NbSample"]):
        pin.framesForwardKinematics(model, data, q_[i, :])
        pin.updateFramePlacements(model, data)
        # update model.jointPlacements with joint elastic error
        if calib_config["non_geom"]:
            tau = pin.computeGeneralizedGravity(
                model, data, q_[i, :]
            )  # vector size of 32 = nq < njoints
            # update xyz_rpy with joint elastic error
            for j_id in calib_config["actJoint_idx"]:
                xyz_rpy = np.zeros(6)
                j_name = model.names[j_id]
                tau_j = tau[j_id - 1]  # nq = njoints -1
                if j_name in key:
                    for elas_id, elas in enumerate(ELAS_TPL):
                        if elas in key:
                            param_dict[key] = param_dict[key] * tau_j
                            xyz_rpy[elas_id + 3] += param_dict[
                                key
                            ]  # +3 to add only on orienation
                            updated_params.append(key)
                model = update_joint_placement(model, j_id, xyz_rpy)
            # get relative transform with updated model
            oMee = get_rel_transform(
                model, data, calib_config["start_frame"], calib_config["end_frame"]
            )
            # revert model back to origin from added joint elastic error
            for j_id in calib_config["actJoint_idx"]:
                xyz_rpy = np.zeros(6)
                j_name = model.names[j_id]
                tau_j = tau[j_id - 1]  # nq = njoints -1
                if j_name in key:
                    for elas_id, elas in enumerate(ELAS_TPL):
                        if elas in key:
                            param_dict[key] = param_dict[key] * tau_j
                            xyz_rpy[elas_id + 3] += param_dict[
                                key
                            ]  # +3 to add only on orienation
                            updated_params.append(key)
                model = update_joint_placement(model, j_id, -xyz_rpy)

        else:
            oMee = get_rel_transform(
                model, data, calib_config["start_frame"], calib_config["end_frame"]
            )

        if len(updated_params) < len(param_dict):

            oMf = oMee * eeMf
            # final transform
            trans = oMf.translation.tolist()
            orient = pin.rpy.matrixToRpy(oMf.rotation).tolist()
            loc = trans + orient
            measure = []
            for mea_id, mea in enumerate(calib_config["measurability"]):
                if mea:
                    measure.append(loc[mea_id])
            # PEE[(marker_idx-1)*calib_config['calibration_index']:marker_idx*calib_config['calibration_index'], i] = np.array(measure)
            PEE[:, i] = np.array(measure)

            # assert len(updated_params) == len(param_dict), "Not all parameters are updated"

    PEE = PEE.flatten("C")
    # revert model back to original
    assert (
        origin_model.jointPlacements != model.jointPlacements
    ), "before revert"
    for j_id in calib_config["actJoint_idx"]:
        xyz_rpy = np.zeros(6)
        j_name = model.names[j_id]
        for key in param_dict.keys():
            if j_name in key:
                # update xyz_rpy
                for axis_id, axis in enumerate(axis_tpl):
                    if axis in key:
                        xyz_rpy[axis_id] = param_dict[key]
        model = update_joint_placement(model, j_id, -xyz_rpy)

    assert (
        origin_model.jointPlacements != model.jointPlacements
    ), "after revert"

    return PEE


def calc_updated_fkm(model, data, var, q, calib_config, verbose=0):
    """Update forward kinematics with world frame transformations.

    Specialized version that explicitly handles transformations between:
    1. World frame to start frame (wMo)
    2. Start frame to end frame (oMee)
    3. End frame to marker frame (eeMf)

    Args:
        model (pin.Model): Robot model to update
        data (pin.Data): Robot data
        var (ndarray): Parameter vector matching calib_config["param_name"]
        q (ndarray): Joint configurations matrix (n_samples, n_joints)
        calib_config (dict): Calibration parameters containing:
            - Frames and parameters from update_forward_kinematics()
            - NbMarkers=1 (only supports single marker)
        verbose (int, optional): Print update info. Defaults to 0.

    Returns:
        ndarray: Flattened marker measurements in world frame

    Notes:
        - Excludes joint elasticity effects
        - Requires base or end-effector parameters in param_name
        - Validates all parameters are properly updated
    """

    # name reference of calibration parameters
    if calib_config["calib_model"] == "full_params":
        axis_tpl = FULL_PARAMTPL

    elif calib_config["calib_model"] == "joint_offset":
        axis_tpl = JOINT_OFFSETTPL

    # order of joint in variables are arranged as in calib_config['actJoint_idx']
    assert len(var) == len(
        calib_config["param_name"]
    ), "Length of variables != length of params"
    param_dict = dict(zip(calib_config["param_name"], var))
    origin_model = model.copy()

    # store parameter updated to the model
    updated_params = []

    # check if baseframe and end--effector frame are known
    for key in param_dict.keys():
        if "base" in key:
            base_param_incl = True
            break
        else:
            base_param_incl = False
    for key in param_dict.keys():
        if "EE" in key:
            ee_param_incl = True
            break
        else:
            ee_param_incl = False

    # kinematic chain
    start_f = calib_config["start_frame"]
    end_f = calib_config["end_frame"]

    # if world frame (measurement ref frame) to the start frame is not known,
    # base_tpl needs to be used to define the first 6 parameters

    # 1/ calc transformation from the world frame to start frame: wMo
    if base_param_incl:
        base_tf = np.zeros(6)
        for key in param_dict.keys():
            for base_id, base_ax in enumerate(BASE_TPL):
                if base_ax in key:
                    base_tf[base_id] = param_dict[key]
                    updated_params.append(key)

        wMo = cartesian_to_SE3(base_tf)
    else:
        wMo = pin.SE3.Identity()

    # 2/ calculate transformation from the end frame to the end-effector frame,
    # if not known: eeMf
    if ee_param_incl and calib_config["NbMarkers"] == 1:
        for marker_idx in range(1, calib_config["NbMarkers"] + 1):
            pee = np.zeros(6)
            ee_name = "EE"
            for key in param_dict.keys():
                if ee_name in key and str(marker_idx) in key:
                    # update xyz_rpy with kinematic errors
                    for axis_pee_id, axis_pee in enumerate(EE_TPL):
                        if axis_pee in key:
                            if verbose == 1:
                                print(
                                    "Updating [{}_{}] joint placement at axis {} with [{}]".format(
                                        ee_name, str(marker_idx), axis_pee, key
                                    )
                                )
                            pee[axis_pee_id] += param_dict[key]
                            updated_params.append(key)

            eeMf = cartesian_to_SE3(pee)
    else:
        if calib_config["NbMarkers"] > 1:
            print("Multiple markers are not supported.")
        else:
            eeMf = pin.SE3.Identity()

    # 3/ calculate transformation from start frame to end frame of kinematic chain using updated model: oMee

    # update model.jointPlacements with kinematic error parameter
    for j_id in calib_config["actJoint_idx"]:
        xyz_rpy = np.zeros(6)
        j_name = model.names[j_id]

        # check joint name in param dict
        for key in param_dict.keys():
            if j_name in key:

                # update xyz_rpy with kinematic errors based on identifiable axis
                for axis_id, axis in enumerate(axis_tpl):
                    if axis in key:
                        if verbose == 1:
                            print(
                                "Updating [{}] joint placement at axis {} with [{}]".format(
                                    j_name, axis, key
                                )
                            )
                        xyz_rpy[axis_id] += param_dict[key]
                        updated_params.append(key)

        # updaet joint placement
        model = update_joint_placement(model, j_id, xyz_rpy)

    # check if all parameters are updated to the model
    assert len(updated_params) == len(
        list(param_dict.keys())
    ), "Not all parameters are updated {} and {}".format(
        updated_params, list(param_dict.keys())
    )

    # pose vector of the end-effector
    PEE = np.zeros((calib_config["calibration_index"], calib_config["NbSample"]))

    q_ = np.copy(q)
    for i in range(calib_config["NbSample"]):

        pin.framesForwardKinematics(model, data, q_[i, :])
        pin.updateFramePlacements(model, data)

        # NOTE: joint elastic error is not considered in this version

        oMee = get_rel_transform(model, data, start_f, end_f)

        # calculate transformation from world frame to end-effector frame
        wMee = wMo * oMee
        wMf = wMee * eeMf

        # final transform
        trans = wMf.translation.tolist()
        orient = pin.rpy.matrixToRpy(wMf.rotation).tolist()
        loc = trans + orient
        measure = []
        for mea_id, mea in enumerate(calib_config["measurability"]):
            if mea:
                measure.append(loc[mea_id])
        PEE[:, i] = np.array(measure)

    # final result of updated fkm
    PEE = PEE.flatten("C")

    # revert model back to original
    assert (
        origin_model.jointPlacements != model.jointPlacements
    ), "before revert"
    for j_id in calib_config["actJoint_idx"]:
        xyz_rpy = np.zeros(6)
        j_name = model.names[j_id]
        for key in param_dict.keys():
            if j_name in key:
                # update xyz_rpy
                for axis_id, axis in enumerate(axis_tpl):
                    if axis in key:
                        xyz_rpy[axis_id] = param_dict[key]
        model = update_joint_placement(model, j_id, -xyz_rpy)

    assert (
        origin_model.jointPlacements != model.jointPlacements
    ), "after revert"

    return PEE


def update_joint_placement(model, joint_idx, xyz_rpy):
    """Update joint placement with offset parameters.

    Modifies a joint's placement transform by adding position and orientation offsets.

    Args:
        model (pin.Model): Robot model to modify
        joint_idx (int): Index of joint to update
        xyz_rpy (ndarray): (6,) array of offsets:
            - xyz_rpy[0:3]: Translation offsets (x,y,z)
            - xyz_rpy[3:6]: Rotation offsets (roll,pitch,yaw)

    Returns:
        pin.Model: Updated robot model

    Side Effects:
        Modifies model.jointPlacements[joint_idx] in place
    """
    tpl_translation = model.jointPlacements[joint_idx].translation
    tpl_rotation = model.jointPlacements[joint_idx].rotation
    tpl_orientation = pin.rpy.matrixToRpy(tpl_rotation)
    # update axes
    updt_translation = tpl_translation + xyz_rpy[0:3]
    updt_orientation = tpl_orientation + xyz_rpy[3:6]
    updt_rotation = pin.rpy.rpyToMatrix(updt_orientation)
    # update placements
    model.jointPlacements[joint_idx].translation = updt_translation
    model.jointPlacements[joint_idx].rotation = updt_rotation
    return model


# BASE REGRESSOR TOOLS


def calculate_kinematics_model(q_i, model, data, calib_config):
    """Calculate Jacobian and kinematic regressor for single configuration.

    Computes frame Jacobian and kinematic regressor matrices for tool frame
    at given joint configuration.

    Args:
        q_i (ndarray): Joint configuration vector
        model (pin.Model): Robot model
        data (pin.Data): Robot data
        calib_config (dict): Parameters containing "IDX_TOOL" frame index

    Returns:
        tuple:
            - model (pin.Model): Updated model
            - data (pin.Data): Updated data
            - R (ndarray): (6,6n) Kinematic regressor matrix
            - J (ndarray): (6,n) Frame Jacobian matrix
    """
    pin.forwardKinematics(model, data, q_i)
    pin.updateFramePlacements(model, data)

    J = pin.computeFrameJacobian(
        model, data, q_i, calib_config["IDX_TOOL"], pin.LOCAL
    )
    R = pin.computeFrameKinematicRegressor(
        model, data, calib_config["IDX_TOOL"], pin.LOCAL
    )
    return model, data, R, J


def calculate_identifiable_kinematics_model(q, model, data, calib_config):
    """Calculate identifiable Jacobian and regressor matrices.

    Builds aggregated Jacobian and regressor matrices from either:
    1. Given set of configurations, or
    2. Random configurations if none provided

    Args:
        q (ndarray, optional): Joint configurations matrix. If empty, uses random configs.
        model (pin.Model): Robot model
        data (pin.Data): Robot data
        calib_config (dict): Parameters containing:
            - NbSample: Number of configurations
            - calibration_index: Number of active DOFs
            - start_frame, end_frame: Frame names
            - calib_model: Model type

    Returns:
        ndarray: Either:
            - Joint offset case: Frame Jacobian matrix
            - Full params case: Kinematic regressor matrix

    Note:
        Removes rows corresponding to inactive DOFs and zero elements
    """
    q_temp = np.copy(q)
    # Note if no q id given then use random generation of q to determine the
    # minimal kinematics model
    if np.any(q):
        MIN_MODEL = 0
    else:
        MIN_MODEL = 1

    # obtain aggreated Jacobian matrix J and kinematic regressor R
    calib_idx = calib_config["calibration_index"]
    R = np.zeros([6 * calib_config["NbSample"], 6 * (model.njoints - 1)])
    J = np.zeros([6 * calib_config["NbSample"], model.njoints - 1])
    for i in range(calib_config["NbSample"]):
        if MIN_MODEL == 1:
            q_rand = pin.randomConfiguration(model)
            q_i = calib_config["q0"]
            q_i[calib_config["config_idx"]] = q_rand[calib_config["config_idx"]]
        else:
            q_i = q_temp[i, :]
        if calib_config["start_frame"] == "universe":
            model, data, Ri, Ji = calculate_kinematics_model(
                q_i, model, data, calib_config
            )
        else:
            Ri = get_rel_kinreg(
                model, data, calib_config["start_frame"], calib_config["end_frame"], q_i
            )
            # Ji = np.zeros([6, model.njoints-1]) ## TODO: get_rel_jac
            Ji = get_rel_jac(
                model, data, calib_config["start_frame"], calib_config["end_frame"], q_i
            )
        for j, state in enumerate(calib_config["measurability"]):
            if state:
                R[calib_config["NbSample"] * j + i, :] = Ri[j, :]
                J[calib_config["NbSample"] * j + i, :] = Ji[j, :]
    # remove zero rows
    zero_rows = []
    for r_idx in range(R.shape[0]):
        if np.linalg.norm(R[r_idx, :]) < 1e-6:
            zero_rows.append(r_idx)
    R = np.delete(R, zero_rows, axis=0)
    zero_rows = []
    for r_idx in range(J.shape[0]):
        if np.linalg.norm(J[r_idx, :]) < 1e-6:
            zero_rows.append(r_idx)
    J = np.delete(J, zero_rows, axis=0)
    
    # select regressor matrix based on calibration model
    if calib_config["calib_model"] == "joint_offset":
        return J
    elif calib_config["calib_model"] == "full_params":
        return R


def calculate_base_kinematics_regressor(q, model, data, calib_config, tol_qr=TOL_QR):
    """Calculate base regressor matrix for calibration parameters.

    Identifies base (identifiable) parameters by:
    1. Computing regressors with random/given configurations
    2. Eliminating unidentifiable parameters
    3. Finding independent regressor columns

    Args:
        q (ndarray): Joint configurations matrix
        model (pin.Model): Robot model
        data (pin.Data): Robot data
        calib_config (dict): Contains calibration settings:
            - free_flyer: Whether base is floating
            - calib_model: Either "joint_offset" or "full_params"
        tol_qr (float, optional): QR decomposition tolerance. Defaults to TOL_QR.

    Returns:
        tuple:
            - Rrand_b (ndarray): Base regressor from random configs
            - R_b (ndarray): Base regressor from given configs
            - R_e (ndarray): Full regressor after eliminating unidentifiable params
            - paramsrand_base (list): Names of base parameters from random configs
            - paramsrand_e (list): Names of identifiable parameters

    Side Effects:
        - Updates calib_config["param_name"] with identified base parameters
        - Prints regressor matrix shapes
    """
    # obtain joint names
    joint_names = [name for i, name in enumerate(model.names[1:])]
    geo_params = get_fullparam_offset(joint_names)
    joint_offsets = get_joint_offset(model, joint_names)

    # calculate kinematic regressor with random configs
    if not calib_config["free_flyer"]:
        Rrand = calculate_identifiable_kinematics_model([], model, data, calib_config)
    else:
        Rrand = calculate_identifiable_kinematics_model(q, model, data, calib_config)
    # calculate kinematic regressor with input configs
    if np.any(np.array(q)):
        R = calculate_identifiable_kinematics_model(q, model, data, calib_config)
    else:
        R = Rrand

    # only joint offset parameters
    if calib_config["calib_model"] == "joint_offset":
        geo_params_sel = joint_offsets

        # select columns corresponding to joint_idx
        Rrand_sel = Rrand

        # select columns corresponding to joint_idx
        R_sel = R

    # full 6 parameters
    elif calib_config["calib_model"] == "full_params":
        geo_params_sel = geo_params
        Rrand_sel = Rrand
        R_sel = R

    # remove non affect columns from random data => reduced regressor
    Rrand_e, paramsrand_e = eliminate_non_dynaffect(
        Rrand_sel, geo_params_sel, tol_e=1e-6
    )

    # indices of independent columns (base param) w.r.t to reduced regressor
    idx_base = get_baseIndex(Rrand_e, paramsrand_e, tol_qr=tol_qr)

    # get base regressor and base params from random data
    Rrand_b, paramsrand_base, _ = get_baseParams(
        Rrand_e, paramsrand_e, tol_qr=tol_qr
    )

    # remove non affect columns from GIVEN data
    R_e, params_e = eliminate_non_dynaffect(R_sel, geo_params_sel, tol_e=1e-6)

    # get base param from given data
    # idx_gbase = get_baseIndex(R_e, params_e, tol_qr=tol_qr)
    R_gb, params_gbase, _ = get_baseParams(R_e, params_e, tol_qr=tol_qr)

    # get base regressor from GIVEN data
    R_b = build_baseRegressor(R_e, idx_base)

    # update calibrating calib_config['param_name']/calibrating parameters
    for j in idx_base:
        calib_config["param_name"].append(paramsrand_e[j])

    return Rrand_b, R_b, R_e, paramsrand_base, paramsrand_e


# Backward compatibility wrapper for get_param_from_yaml
def get_param_from_yaml_legacy(robot, calib_data) -> dict:
    """Legacy calibration parameter parser - kept for backward compatibility.
    
    This is the original implementation. New code should use the unified
    config parser from figaroh.utils.config_parser.
    
    Args:
        robot: Robot instance
        calib_data: Calibration data dictionary
        
    Returns:
        Calibration configuration dictionary
    """
    # Keep the original implementation here for compatibility
    return get_param_from_yaml(robot, calib_data)


# Import the new unified parser as the default
try:
    from ..utils.config_parser import get_param_from_yaml as unified_get_param_from_yaml
    
    # Replace the function with unified version while maintaining signature
    def get_param_from_yaml_unified(robot, calib_data) -> dict:
        """Enhanced parameter parser using unified configuration system.
        
        This function provides backward compatibility while using the new
        unified configuration parser when possible.
        
        Args:
            robot: Robot instance
            calib_data: Configuration data (dict or file path)
            
        Returns:
            Calibration configuration dictionary
        """
        try:
            return unified_get_param_from_yaml(robot, calib_data, "calibration")
        except Exception as e:
            # Fall back to legacy parser if unified parser fails
            import warnings
            warnings.warn(
                f"Unified parser failed ({e}), falling back to legacy parser. "
                "Consider updating your configuration format.",
                UserWarning
            )
            return get_param_from_yaml_legacy(robot, calib_data)
    
    # Keep the old function available but with warning
    def get_param_from_yaml_with_warning(robot, calib_data) -> dict:
        """Original function with deprecation notice.""" 
        import warnings
        warnings.warn(
            "Direct use of get_param_from_yaml is deprecated. "
            "Consider using the unified config parser from figaroh.utils.config_parser",
            DeprecationWarning,
            stacklevel=2
        )
        return get_param_from_yaml_unified(robot, calib_data)
        
except ImportError:
    # If unified parser is not available, keep using original function
    pass
