"""Single-arm IK via mink + PostureTask (smith300/PARA defaults).

One function: :func:`ee_chain_to_joints`. Given a sequence of target EE
poses (world frame, 4×4) plus a seed proprio, returns a (T, 7) chain
of joint commands. Seed is updated per timestep so the elbow basin
stays consistent across the chunk.

NB: model is the combined yam-4310-linear xml (i2rt). Posture cost is
5e-2 per the old pre-refactor deploy code — empirically needed to
prevent basin jumps mid-chunk.
"""
from __future__ import annotations

from pathlib import Path
from typing import List

import numpy as np


_KIN_CONFIG = None
_NQ_FULL = None


def _get_config():
    """Lazy-init mink Configuration on the right-arm xml.

    Matches the pre-refactor `_get_mink_config` in deploy_yam_2view.py:
    ``mujoco.MjModel.from_xml_path → mink.Configuration``. (``mink.utils``
    doesn't have a model loader; build the MjModel via mujoco directly.)
    """
    global _KIN_CONFIG, _NQ_FULL
    if _KIN_CONFIG is None:
        import sys
        rf = Path(__file__).resolve().parents[1] / "raiden_fork"
        for p in (rf, rf / "third_party" / "i2rt"):
            if str(p) not in sys.path:
                sys.path.insert(0, str(p))
        import mink
        import mujoco
        from raiden._xml_paths import get_yam_4310_linear_xml_path
        mj_model = mujoco.MjModel.from_xml_path(get_yam_4310_linear_xml_path())
        _KIN_CONFIG = mink.Configuration(mj_model)
        _NQ_FULL = int(_KIN_CONFIG.model.nq)
    return _KIN_CONFIG, _NQ_FULL


def _pad_q(q: np.ndarray) -> np.ndarray:
    _, nq = _get_config()
    out = np.zeros(nq, dtype=np.float64)
    out[: min(len(q), nq)] = q[: min(len(q), nq)]
    return out


def ik_solve(target_pose: np.ndarray, init_q: np.ndarray,
              pos_cost: float = 8.0, ori_cost: float = 0.1,
              posture_cost: float = 5e-2, max_iters: int = 60,
              pos_thresh: float = 1e-4, ori_thresh: float = 1e-4,
              dt: float = 1e-2, damping: float = 1e-4,
              solver: str = "daqp") -> tuple[bool, np.ndarray, dict]:
    """Solve for joints reaching ``target_pose`` (4×4 world).

    Returns ``(ok, q, info)`` — ``info`` has ``iters``, ``pos_err_m``,
    ``ori_err_rad``. Pos/ori weight ratio is 80:1 (pos_cost=8.0,
    ori_cost=0.1) — Cameron prefers tighter position even at the cost
    of orientation, since the gripper's grasp tolerates ~5° of yaw drift
    but not mm-scale position drift. Bumped from 40:1 → 80:1 on
    2026-06-14 after observing ~1-inch position errors on the
    egg_carton deploy (we'd rather have worse rotation than be off in
    position). Posture cost 5e-2 matches the pre-refactor
    ``ik_with_posture`` (keeps elbow basin consistent across a chunk).
    """
    import mink
    config, _ = _get_config()
    config.update(_pad_q(init_q))
    ee = mink.FrameTask(frame_name="grasp_site", frame_type="site",
                         position_cost=pos_cost, orientation_cost=ori_cost,
                         lm_damping=1.0)
    ee.set_target(mink.SE3.from_matrix(target_pose))
    posture = mink.PostureTask(model=config.model, cost=posture_cost)
    posture.set_target(_pad_q(init_q))
    tasks = [ee, posture]
    limits = [mink.ConfigurationLimit(config.model)]
    err = np.zeros(6)
    for i in range(max_iters):
        vel = mink.solve_ik(config, tasks, dt, solver, damping=damping,
                              limits=limits)
        config.integrate_inplace(vel, dt)
        err = ee.compute_error(config)
        if (np.linalg.norm(err[:3]) <= pos_thresh
                and np.linalg.norm(err[3:]) <= ori_thresh):
            return True, np.asarray(config.q).copy(), {
                "iters": i + 1,
                "pos_err_m": float(np.linalg.norm(err[:3])),
                "ori_err_rad": float(np.linalg.norm(err[3:])),
            }
    return False, np.asarray(config.q).copy(), {
        "iters": max_iters,
        "pos_err_m": float(np.linalg.norm(err[:3])),
        "ori_err_rad": float(np.linalg.norm(err[3:])),
    }


def _quat_to_R(quat: np.ndarray) -> np.ndarray:
    """Quaternion (w, x, y, z) → 3×3 rotation matrix."""
    w, x, y, z = quat
    return np.array([
        [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)],
        [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)],
        [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)],
    ])


def ee_chain_to_joints(
    xyz_world: np.ndarray,        # (T, 3) WORLD (= left_arm_base) EE positions
    quat: np.ndarray,             # (T, 4) WORLD EE rotations (w, x, y, z)
    grip: np.ndarray,             # (T,) gripper values
    init_q7: np.ndarray,          # (7,) current proprio (right arm: 6 arm + 1 gripper)
    T_lfr: np.ndarray = None,     # (4, 4) T_left_from_right; REQUIRED if model outputs are in world frame
    return_info: bool = False,
):
    """Solve a chain of IK problems, seeding each from the previous solution.

    The IK config uses the **right-arm-only** ``yam_4310_linear`` XML
    (single ``grasp_site``), so the target pose must be expressed in
    **right_arm_base** frame. Model predictions are in **world**
    (= left_arm_base) frame, so we apply ``inv(T_lfr) @ T_target_in_world``
    before each solve. Pass ``T_lfr=None`` only if your targets are
    already in right_arm_base.

    Returns ``List[(7,) float32]`` — one joint vector per timestep.
    With ``return_info=True``, also returns a list of per-T convergence
    dicts (``ok``, ``iters``, ``pos_err_m``, ``ori_err_rad``).
    Failed IKs fall back to the seed and emit a warning.
    """
    T_lfr_inv = np.linalg.inv(T_lfr) if T_lfr is not None else np.eye(4)
    T = len(xyz_world)
    out = []
    infos: list[dict] = []
    seed = init_q7.astype(np.float64)
    for t in range(T):
        target_w = np.eye(4)
        target_w[:3, 3] = xyz_world[t]
        target_w[:3, :3] = _quat_to_R(quat[t])
        target_rbase = T_lfr_inv @ target_w
        ok, q, info = ik_solve(target_rbase, seed)
        info["ok"] = ok
        infos.append(info)
        if not ok:
            # Use the PARTIAL solve, not the seed. Matches the pre-refactor
            # ``ee_actions_to_joint_cmds`` which ignores the ok flag — at
            # tight (1e-4) thresholds most solves fail nominally even when
            # they're physically excellent (~5-10mm). Throwing the result
            # away cascades to "every step stays at home" via the seed chain.
            print(f"  IK warn @ t={t}: didn't hit 1e-4 thresh "
                  f"({info['iters']} iters, "
                  f"pos_err={info['pos_err_m']*1000:.2f}mm, "
                  f"ori_err={np.degrees(info['ori_err_rad']):.2f}°); "
                  f"keeping partial q")
        q7 = np.zeros(7, dtype=np.float32)
        q7[:6] = q[:6]
        q7[6] = float(grip[t])
        out.append(q7)
        seed = q.copy()
    if return_info:
        return out, infos
    return out