from __future__ import annotations

from typing import Optional

import numpy as np
import pinocchio as pin

from .wrench_solver import WrenchSolveResult, WrenchSolver


def checked_frame_id(model: pin.Model, frame_name: str) -> int:
    """Resolve a Pinocchio frame name and reject its not-found sentinel."""
    if not isinstance(frame_name, str) or not frame_name:
        raise ValueError("frame_name must be a non-empty string")

    frame_id = int(model.getFrameId(frame_name))
    if frame_id < 0 or frame_id >= model.nframes:
        available = ", ".join(frame.name for frame in model.frames)
        raise ValueError(
            f"Frame not found: {frame_name!r}. "
            f"Expected an id in [0, {model.nframes}), got {frame_id}. "
            f"Available frames: {available}"
        )
    return frame_id


def checked_joint_id(model: pin.Model, joint_name: str) -> int:
    """Resolve a movable joint name and reject universe/not-found sentinels."""
    if not isinstance(joint_name, str) or not joint_name:
        raise ValueError("joint_name must be a non-empty string")

    joint_id = int(model.getJointId(joint_name))
    if joint_id <= 0 or joint_id >= model.njoints:
        available = ", ".join(model.names[1:])
        raise ValueError(
            f"Movable joint not found: {joint_name!r}. "
            f"Expected an id in [1, {model.njoints}), got {joint_id}. "
            f"Available movable joints: {available}"
        )
    return joint_id


def _as_vector(value, size: int, name: str) -> np.ndarray:
    vector = np.asarray(value, dtype=float).reshape(-1)
    if vector.shape != (size,):
        raise ValueError(f"{name} must have shape ({size},), got {vector.shape}")
    if not np.all(np.isfinite(vector)):
        raise ValueError(f"{name} must contain only finite values")
    return vector


def _as_sample_matrix(value, width: int, name: str) -> np.ndarray:
    samples = np.asarray(value, dtype=float)
    if samples.ndim == 1:
        samples = samples.reshape(1, -1)
    if samples.ndim != 2 or samples.shape[1] != width:
        raise ValueError(
            f"{name} must have shape (n_samples, {width}), got {samples.shape}"
        )
    if samples.shape[0] == 0:
        raise ValueError(f"{name} must contain at least one sample")
    if not np.all(np.isfinite(samples)):
        raise ValueError(f"{name} must contain only finite values")
    return samples


class InteractionEstimator:
    """
    Estimate an external wrench at the end-effector point.

    Conventions are explicit throughout this class:

    * twists are ``[linear_velocity; angular_velocity]``;
    * wrenches are ``[force; moment]``;
    * both are expressed along the chest-frame axes, but at the EE point;
    * ``tau_int = tau_meas - tau_model - tau_bias``.

    ``estimate`` keeps the original three-value return contract. Its first
    result is the bias-corrected residual. The raw and corrected residuals are
    also available through ``last_tau_residual_raw`` and
    ``last_tau_residual_corrected``.
    """

    def __init__(
        self,
        model: pin.Model,
        chest_frame_name: str,
        ee_frame_name: str,
        lambda_damp: float = 1e-3,
        wrench_solver: Optional[WrenchSolver] = None,
    ):
        damping = float(lambda_damp)
        if not np.isfinite(damping) or damping <= 0.0:
            raise ValueError("lambda_damp must be a finite positive scalar")

        self.model = model
        self.data = model.createData()
        self.lambda_damp = damping
        self.wrench_solver = wrench_solver
        self.fid_C = checked_frame_id(model, chest_frame_name)
        self.fid_EE = checked_frame_id(model, ee_frame_name)

        self._tau_bias = np.zeros(model.nv, dtype=float)
        self._last_tau_residual_raw: Optional[np.ndarray] = None
        self._last_tau_residual_corrected: Optional[np.ndarray] = None
        self._last_wrench_solve: Optional[WrenchSolveResult] = None

    @staticmethod
    def _rotation6(rotation: np.ndarray) -> np.ndarray:
        """Apply one 3-D rotation to both linear and angular blocks."""
        rotation = np.asarray(rotation, dtype=float)
        if rotation.shape != (3, 3):
            raise ValueError(f"rotation must have shape (3, 3), got {rotation.shape}")
        rotation6 = np.zeros((6, 6), dtype=float)
        rotation6[:3, :3] = rotation
        rotation6[3:, 3:] = rotation
        return rotation6

    @property
    def tau_bias(self) -> np.ndarray:
        """Current no-contact joint-torque bias (copy)."""
        return self._tau_bias.copy()

    @property
    def last_tau_residual_raw(self) -> Optional[np.ndarray]:
        """Latest ``tau_meas - tau_model`` sample, before bias removal."""
        if self._last_tau_residual_raw is None:
            return None
        return self._last_tau_residual_raw.copy()

    @property
    def last_tau_residual_corrected(self) -> Optional[np.ndarray]:
        """Latest raw residual minus the calibrated bias."""
        if self._last_tau_residual_corrected is None:
            return None
        return self._last_tau_residual_corrected.copy()

    @property
    def last_wrench_solve(self) -> Optional[WrenchSolveResult]:
        """Latest formal scaled-solver result, if one was configured."""
        return self._last_wrench_solve

    def calibrate_bias(self, tau_residual_raw_samples) -> np.ndarray:
        """
        Calibrate a constant joint-torque bias from offline no-contact samples.

        Parameters
        ----------
        tau_residual_raw_samples:
            One sample with shape ``(nv,)`` or a batch with shape
            ``(n_samples, nv)``. Every row must already be the raw residual
            ``tau_meas - tau_model`` collected under a no-contact condition.

        Returns
        -------
        numpy.ndarray
            A copy of the calibrated mean bias.
        """
        samples = _as_sample_matrix(
            tau_residual_raw_samples, self.model.nv, "tau_residual_raw_samples"
        )
        self._tau_bias = np.mean(samples, axis=0)
        return self.tau_bias

    def calibrate_bias_from_measurements(
        self,
        q_samples,
        qd_samples,
        qdd_samples,
        tau_meas_samples,
        tau_ff_fric_samples=None,
    ) -> np.ndarray:
        """
        Compute and calibrate bias from an offline no-contact measurement batch.

        Each argument is a row-major sample matrix. ``tau_ff_fric_samples`` is
        optional and, when supplied, must have the same ``(n_samples, nv)``
        shape as the velocity/torque batches.
        """
        q_batch = _as_sample_matrix(q_samples, self.model.nq, "q_samples")
        qd_batch = _as_sample_matrix(qd_samples, self.model.nv, "qd_samples")
        qdd_batch = _as_sample_matrix(qdd_samples, self.model.nv, "qdd_samples")
        tau_batch = _as_sample_matrix(
            tau_meas_samples, self.model.nv, "tau_meas_samples"
        )

        sample_count = q_batch.shape[0]
        batches = (qd_batch, qdd_batch, tau_batch)
        if any(batch.shape[0] != sample_count for batch in batches):
            raise ValueError("all calibration batches must have the same sample count")

        friction_batch = None
        if tau_ff_fric_samples is not None:
            friction_batch = _as_sample_matrix(
                tau_ff_fric_samples, self.model.nv, "tau_ff_fric_samples"
            )
            if friction_batch.shape[0] != sample_count:
                raise ValueError(
                    "tau_ff_fric_samples must match the calibration sample count"
                )

        residuals = np.empty((sample_count, self.model.nv), dtype=float)
        for sample_index in range(sample_count):
            friction = (
                None if friction_batch is None else friction_batch[sample_index]
            )
            tau_model = self._tau_model(
                q_batch[sample_index],
                qd_batch[sample_index],
                qdd_batch[sample_index],
                friction,
            )
            residuals[sample_index] = tau_batch[sample_index] - tau_model
        return self.calibrate_bias(residuals)

    def clear_bias(self) -> None:
        """Clear the calibrated joint-torque bias."""
        self._tau_bias.fill(0.0)

    def _chest_jacobian(self, q, qd):
        """
        Return the EE-point Jacobian expressed along chest-frame axes.

        ``LOCAL_WORLD_ALIGNED`` is essential here: unlike ``WORLD``, its
        translational rows are the linear velocity of the EE origin. Rotating
        the two 3-D blocks changes only their coordinate axes; no translational
        adjoint term is used, so the wrench/twist reference point stays at EE.
        """
        q = _as_vector(q, self.model.nq, "q")
        qd = _as_vector(qd, self.model.nv, "qd")

        pin.forwardKinematics(self.model, self.data, q, qd)
        pin.updateFramePlacements(self.model, self.data)

        jacobian_lwa = pin.computeFrameJacobian(
            self.model,
            self.data,
            q,
            self.fid_EE,
            pin.ReferenceFrame.LOCAL_WORLD_ALIGNED,
        )

        rotation_world_from_chest = self.data.oMf[self.fid_C].rotation
        rotation_chest_from_world = rotation_world_from_chest.T
        return self._rotation6(rotation_chest_from_world) @ jacobian_lwa

    def _tau_model(self, q, qd, qdd, tau_ff_fric=None):
        """Return ``M qdd + C qd + g + tau_ff_fric``."""
        q = _as_vector(q, self.model.nq, "q")
        qd = _as_vector(qd, self.model.nv, "qd")
        qdd = _as_vector(qdd, self.model.nv, "qdd")

        mass_matrix = pin.crba(self.model, self.data, q)
        mass_matrix = (
            mass_matrix + mass_matrix.T - np.diag(mass_matrix.diagonal())
        )
        nonlinear = pin.nonLinearEffects(self.model, self.data, q, qd)
        tau_model = mass_matrix @ qdd + nonlinear

        if tau_ff_fric is not None:
            tau_model = tau_model + _as_vector(
                tau_ff_fric, self.model.nv, "tau_ff_fric"
            )
        return tau_model

    def estimate(self, q, qd, qdd, tau_meas, tau_ff_fric=None):
        """
        Estimate the bias-corrected joint residual and chest-axis EE wrench.

        A fixed damped least-squares solve is used:

        ``F = (J J.T + lambda**2 I)^-1 J tau_int``.
        """
        tau_meas = _as_vector(tau_meas, self.model.nv, "tau_meas")
        tau_model = self._tau_model(q, qd, qdd, tau_ff_fric)

        tau_residual_raw = tau_meas - tau_model
        tau_residual_corrected = tau_residual_raw - self._tau_bias
        self._last_tau_residual_raw = tau_residual_raw.copy()
        self._last_tau_residual_corrected = tau_residual_corrected.copy()

        chest_jacobian = self._chest_jacobian(q, qd)
        if self.wrench_solver is None:
            normal_matrix = chest_jacobian @ chest_jacobian.T
            normal_matrix = normal_matrix + (
                self.lambda_damp**2
            ) * np.eye(6, dtype=float)
            wrench_chest = np.linalg.solve(
                normal_matrix, chest_jacobian @ tau_residual_corrected
            )
            self._last_wrench_solve = None
        else:
            solve = self.wrench_solver.solve(
                chest_jacobian, tau_residual_corrected
            )
            wrench_chest = solve.wrench.copy()
            self._last_wrench_solve = solve

        return tau_residual_corrected, wrench_chest, chest_jacobian
