"""Canonical model contract for the heterogeneous teleoperation prototype.

The repository contains several historical URDF/config combinations.  This
module defines the only combination used by the new simulation:

* master: ``config/master_7dof.urdf``
* slave: ``config/real_slave_7dof.urdf``

The physical slave URDF currently ends at the final wrist body and does not
contain a calibrated tool/force-sensor frame.  For simulation only, an
operational frame named ``R_EE_SIM`` is added at a documented fixed offset.
Hardware experiments must replace that offset with the measured TCP transform.
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, Sequence

import numpy as np
import pinocchio as pin


CODE_ROOT = Path(__file__).resolve().parents[1]
CONFIG_ROOT = CODE_ROOT / "config"
MASTER_URDF = CONFIG_ROOT / "master_7dof.urdf"
SLAVE_URDF = CONFIG_ROOT / "real_slave_7dof.urdf"

MASTER_JOINT_NAMES = (
    "master_shoulder_pitch_joint",
    "master_shoulder_yaw_joint",
    "master_shoulder_roll_joint",
    "master_elbow_flex_joint",
    "master_wrist_roll_joint",
    "master_wrist_yaw_joint",
    "master_wrist_pitch_joint",
)

SLAVE_JOINT_NAMES = (
    "R_SHOULDER_P",
    "R_SHOULDER_R",
    "R_SHOULDER_Y",
    "R_ELBOW_R",
    "R_WRIST_P",
    "R_WRIST_Y",
    "R_WRIST_R",
)

MASTER_FRAMES = {
    "base": "master_base",
    "shoulder": "master_shoulder",
    "elbow": "master_forearm",
    "wrist": "master_wrist",
    "ee": "master_ee",
}

SLAVE_FRAMES = {
    "base": "PELVIS_S",
    "shoulder": "R_SHOULDER_R_S",
    "elbow": "R_ELBOW_R_S",
    "wrist": "R_WRIST_R_S",
    "ee": "R_EE_SIM",
}

# Approximate simulated TCP taken from the terminal marker in right_arm.mjcf.
# This is deliberately not presented as a physical calibration.
SLAVE_SIM_TCP_OFFSET = np.array([-0.01212, -0.17655, 0.07506], dtype=float)


@dataclass(frozen=True)
class TeleoperationModels:
    master: pin.Model
    slave: pin.Model


def require_frame(model: pin.Model, name: str) -> int:
    """Resolve a frame and reject Pinocchio's not-found sentinel."""
    fid = int(model.getFrameId(name))
    if fid < 0 or fid >= model.nframes:
        raise ValueError(
            f"Frame {name!r} not found in model {model.name!r}; "
            f"available={[frame.name for frame in model.frames]}"
        )
    return fid


def require_joint(model: pin.Model, name: str) -> int:
    """Resolve an actuated joint and reject universe/not-found sentinels."""
    jid = int(model.getJointId(name))
    if jid <= 0 or jid >= model.njoints:
        raise ValueError(
            f"Joint {name!r} not found in model {model.name!r}; "
            f"available={list(model.names)[1:]}"
        )
    return jid


def joint_q_indices(model: pin.Model, names: Sequence[str]) -> np.ndarray:
    indices = []
    for name in names:
        jid = require_joint(model, name)
        joint = model.joints[jid]
        if joint.nq != 1 or joint.nv != 1:
            raise ValueError(f"Expected a scalar revolute joint, got {name!r}")
        indices.append(int(joint.idx_q))
    return np.asarray(indices, dtype=int)


def joint_v_indices(model: pin.Model, names: Sequence[str]) -> np.ndarray:
    indices = []
    for name in names:
        jid = require_joint(model, name)
        joint = model.joints[jid]
        if joint.nq != 1 or joint.nv != 1:
            raise ValueError(f"Expected a scalar revolute joint, got {name!r}")
        indices.append(int(joint.idx_v))
    return np.asarray(indices, dtype=int)


def validate_contract(
    model: pin.Model,
    joint_names: Sequence[str],
    frame_names: Iterable[str],
) -> None:
    if model.nq != 7 or model.nv != 7:
        raise ValueError(
            f"Expected a fixed-base 7-DoF model, got nq={model.nq}, nv={model.nv}"
        )
    joint_q_indices(model, joint_names)
    for frame_name in frame_names:
        require_frame(model, frame_name)


def add_fixed_operational_frame(
    model: pin.Model,
    *,
    name: str,
    parent_frame_name: str,
    translation: np.ndarray,
    rotation: np.ndarray | None = None,
) -> int:
    """Add an operational frame at a transform relative to an existing frame."""
    existing = int(model.getFrameId(name))
    if 0 <= existing < model.nframes:
        return existing

    parent_fid = require_frame(model, parent_frame_name)
    parent = model.frames[parent_fid]
    relative = pin.SE3(
        np.eye(3) if rotation is None else np.asarray(rotation, dtype=float),
        np.asarray(translation, dtype=float).reshape(3),
    )
    placement_in_parent_joint = parent.placement * relative
    return int(
        model.addFrame(
            pin.Frame(
                name,
                parent.parentJoint,
                parent_fid,
                placement_in_parent_joint,
                pin.FrameType.OP_FRAME,
            )
        )
    )


def load_models(*, add_simulated_tcp: bool = True) -> TeleoperationModels:
    """Load and validate the canonical master/slave dynamics models."""
    master = pin.buildModelFromUrdf(str(MASTER_URDF))
    slave = pin.buildModelFromUrdf(str(SLAVE_URDF))

    if add_simulated_tcp:
        add_fixed_operational_frame(
            slave,
            name=SLAVE_FRAMES["ee"],
            parent_frame_name=SLAVE_FRAMES["wrist"],
            translation=SLAVE_SIM_TCP_OFFSET,
        )

    validate_contract(master, MASTER_JOINT_NAMES, MASTER_FRAMES.values())
    validate_contract(slave, SLAVE_JOINT_NAMES, SLAVE_FRAMES.values())
    return TeleoperationModels(master=master, slave=slave)


def finite_joint_limits(
    model: pin.Model,
    joint_names: Sequence[str],
) -> tuple[np.ndarray, np.ndarray]:
    qidx = joint_q_indices(model, joint_names)
    lower = np.asarray(model.lowerPositionLimit[qidx], dtype=float)
    upper = np.asarray(model.upperPositionLimit[qidx], dtype=float)
    if not np.all(np.isfinite(lower)) or not np.all(np.isfinite(upper)):
        raise ValueError("All teleoperation joints must have finite position limits")
    if np.any(lower >= upper):
        raise ValueError("Invalid joint limits in teleoperation model")
    return lower, upper


def safe_configuration(
    model: pin.Model,
    joint_names: Sequence[str],
    *,
    fraction: float = 0.5,
    margin: float = 1e-3,
) -> np.ndarray:
    """Return a configuration strictly inside all declared joint limits."""
    if not 0.0 <= fraction <= 1.0:
        raise ValueError("fraction must lie in [0, 1]")
    lower, upper = finite_joint_limits(model, joint_names)
    span = upper - lower
    q7 = lower + fraction * span
    q7 = np.minimum(np.maximum(q7, lower + margin), upper - margin)
    q = pin.neutral(model)
    q[joint_q_indices(model, joint_names)] = q7
    return q


def clip_configuration(
    model: pin.Model,
    q: np.ndarray,
    joint_names: Sequence[str],
    *,
    margin: float = 1e-6,
) -> tuple[np.ndarray, bool]:
    """Clip scalar teleoperation joints and report whether clipping occurred."""
    q_clipped = np.asarray(q, dtype=float).copy()
    qidx = joint_q_indices(model, joint_names)
    lower, upper = finite_joint_limits(model, joint_names)
    clipped7 = np.minimum(np.maximum(q_clipped[qidx], lower + margin), upper - margin)
    changed = bool(np.any(np.abs(clipped7 - q_clipped[qidx]) > 1e-12))
    q_clipped[qidx] = clipped7
    return q_clipped, changed