"""Dimensionally scaled residual-to-wrench solvers for H2.

Twists and Jacobians use ``[linear; angular]`` order and wrenches use
``[force; moment]`` order.  For characteristic length ``ell`` this module
implements the manuscript convention

``S = diag(1/ell I3, I3)``, ``J_tilde = S J`` and
``F_tilde = S**(-T) F``.

Consequently ``J_tilde.T @ F_tilde == J.T @ F`` and every element of
``F_tilde`` has torque units.  Both the damped and undamped methods use the
same scaling and frozen relative rank tolerance so their H2 comparison is
well-defined.
"""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import Protocol, runtime_checkable

import numpy as np


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


def _jacobian(value) -> np.ndarray:
    array = np.asarray(value, dtype=float)
    if array.ndim != 2 or array.shape[0] != 6 or array.shape[1] == 0:
        raise ValueError(
            f"jacobian must have shape (6, n_joints), got {array.shape}"
        )
    if not np.all(np.isfinite(array)):
        raise ValueError("jacobian must contain only finite values")
    return array


def _readonly(value: np.ndarray) -> np.ndarray:
    result = np.asarray(value, dtype=float).copy()
    result.setflags(write=False)
    return result


class WrenchSolveStatus(str, Enum):
    """Numerical rank status after the frozen scaled-Jacobian test."""

    FULL_RANK = "full_rank"
    RANK_DEFICIENT = "rank_deficient"


@dataclass(frozen=True)
class WrenchSolveResult:
    """Solver result and all quantities needed for H2 stratification."""

    method: str
    wrench: np.ndarray
    scaled_wrench: np.ndarray
    reconstructed_residual: np.ndarray
    residual_error_norm: float
    singular_values: np.ndarray
    rank: int
    rank_threshold: float
    condition_number: float
    status: WrenchSolveStatus
    characteristic_length_m: float
    damping: float

    def __post_init__(self) -> None:
        wrench = _vector(self.wrench, 6, "wrench")
        scaled_wrench = _vector(self.scaled_wrench, 6, "scaled_wrench")
        reconstructed = np.asarray(
            self.reconstructed_residual, dtype=float
        ).reshape(-1)
        singular_values = np.asarray(self.singular_values, dtype=float).reshape(-1)
        if reconstructed.size == 0 or not np.all(np.isfinite(reconstructed)):
            raise ValueError("reconstructed_residual must be a finite vector")
        if singular_values.size == 0 or not np.all(np.isfinite(singular_values)):
            raise ValueError("singular_values must be a non-empty finite vector")
        if self.rank < 0 or self.rank > 6:
            raise ValueError("rank must lie in [0, 6]")
        finite_nonnegative = (
            self.residual_error_norm,
            self.rank_threshold,
            self.characteristic_length_m,
            self.damping,
        )
        if any(not np.isfinite(value) or value < 0.0 for value in finite_nonnegative):
            raise ValueError(
                "solver scalar diagnostics must be finite and non-negative"
            )
        if not (
            np.isfinite(self.condition_number)
            or np.isinf(self.condition_number)
        ):
            raise ValueError("condition_number must be finite or infinity")
        object.__setattr__(self, "wrench", _readonly(wrench))
        object.__setattr__(self, "scaled_wrench", _readonly(scaled_wrench))
        object.__setattr__(
            self, "reconstructed_residual", _readonly(reconstructed)
        )
        object.__setattr__(self, "singular_values", _readonly(singular_values))

    @property
    def force(self) -> np.ndarray:
        return self.wrench[:3].copy()

    @property
    def moment(self) -> np.ndarray:
        return self.wrench[3:].copy()


@runtime_checkable
class WrenchSolver(Protocol):
    """Common H2 solver interface."""

    name: str
    characteristic_length_m: float
    relative_rank_tolerance: float

    def solve(
        self, jacobian: np.ndarray, joint_residual: np.ndarray
    ) -> WrenchSolveResult:
        ...


class _ScaledSolverBase:
    def __init__(
        self,
        characteristic_length_m: float,
        *,
        relative_rank_tolerance: float = 1e-9,
    ) -> None:
        length = float(characteristic_length_m)
        tolerance = float(relative_rank_tolerance)
        if not np.isfinite(length) or length <= 0.0:
            raise ValueError(
                "characteristic_length_m must be finite and positive"
            )
        if not np.isfinite(tolerance) or tolerance <= 0.0 or tolerance >= 1.0:
            raise ValueError(
                "relative_rank_tolerance must be finite and lie in (0, 1)"
            )
        self.characteristic_length_m = length
        self.relative_rank_tolerance = tolerance
        self._scaling = np.diag(
            np.array([1.0 / length] * 3 + [1.0] * 3, dtype=float)
        )

    def scaled_jacobian(self, jacobian: np.ndarray) -> np.ndarray:
        """Return ``S_ell @ J`` using the frozen characteristic length."""
        return self._scaling @ _jacobian(jacobian)

    def _decompose(
        self, jacobian: np.ndarray, joint_residual: np.ndarray
    ) -> tuple[
        np.ndarray,
        np.ndarray,
        np.ndarray,
        np.ndarray,
        np.ndarray,
        float,
        int,
    ]:
        jacobian = _jacobian(jacobian)
        residual = _vector(
            joint_residual, jacobian.shape[1], "joint_residual"
        )
        scaled_jacobian = self._scaling @ jacobian
        u, singular_values, vt = np.linalg.svd(
            scaled_jacobian, full_matrices=False
        )
        largest = float(singular_values[0]) if singular_values.size else 0.0
        threshold = self.relative_rank_tolerance * largest
        rank = int(np.count_nonzero(singular_values > threshold))
        return (
            jacobian,
            residual,
            u,
            singular_values,
            vt,
            threshold,
            rank,
        )

    def _result(
        self,
        *,
        jacobian: np.ndarray,
        residual: np.ndarray,
        scaled_wrench: np.ndarray,
        singular_values: np.ndarray,
        threshold: float,
        rank: int,
        damping: float,
    ) -> WrenchSolveResult:
        wrench = self._scaling.T @ scaled_wrench
        reconstructed = jacobian.T @ wrench
        smallest = float(singular_values[-1])
        condition = (
            float(singular_values[0] / smallest)
            if rank == 6 and smallest > 0.0
            else float("inf")
        )
        return WrenchSolveResult(
            method=self.name,
            wrench=wrench,
            scaled_wrench=scaled_wrench,
            reconstructed_residual=reconstructed,
            residual_error_norm=float(np.linalg.norm(reconstructed - residual)),
            singular_values=singular_values,
            rank=rank,
            rank_threshold=threshold,
            condition_number=condition,
            status=(
                WrenchSolveStatus.FULL_RANK
                if rank == 6
                else WrenchSolveStatus.RANK_DEFICIENT
            ),
            characteristic_length_m=self.characteristic_length_m,
            damping=damping,
        )


class ScaledDLSSolver(_ScaledSolverBase):
    """Dimensionally scaled damped least-squares wrench reconstruction."""

    name = "scaled_dls"

    def __init__(
        self,
        characteristic_length_m: float,
        damping: float,
        *,
        relative_rank_tolerance: float = 1e-9,
    ) -> None:
        super().__init__(
            characteristic_length_m,
            relative_rank_tolerance=relative_rank_tolerance,
        )
        damping_value = float(damping)
        if not np.isfinite(damping_value) or damping_value <= 0.0:
            raise ValueError("damping must be finite and positive")
        self.damping = damping_value

    def solve(
        self, jacobian: np.ndarray, joint_residual: np.ndarray
    ) -> WrenchSolveResult:
        (
            jacobian,
            residual,
            u,
            singular_values,
            vt,
            threshold,
            rank,
        ) = self._decompose(jacobian, joint_residual)
        gains = singular_values / (singular_values**2 + self.damping**2)
        scaled_wrench = u @ (gains * (vt @ residual))
        return self._result(
            jacobian=jacobian,
            residual=residual,
            scaled_wrench=scaled_wrench,
            singular_values=singular_values,
            threshold=threshold,
            rank=rank,
            damping=self.damping,
        )


class UndampedSVDSolver(_ScaledSolverBase):
    """Frozen-tolerance, dimensionally scaled undamped pseudoinverse baseline."""

    name = "undamped_svd"

    def solve(
        self, jacobian: np.ndarray, joint_residual: np.ndarray
    ) -> WrenchSolveResult:
        (
            jacobian,
            residual,
            u,
            singular_values,
            vt,
            threshold,
            rank,
        ) = self._decompose(jacobian, joint_residual)
        inverse = np.zeros_like(singular_values)
        retained = singular_values > threshold
        inverse[retained] = 1.0 / singular_values[retained]
        scaled_wrench = u @ (inverse * (vt @ residual))
        return self._result(
            jacobian=jacobian,
            residual=residual,
            scaled_wrench=scaled_wrench,
            singular_values=singular_values,
            threshold=threshold,
            rank=rank,
            damping=0.0,
        )