"""Causal signal and calibration components for interaction estimation.

These components are deliberately independent of the Pinocchio model.  A
locked experiment can therefore replay the exact same measured signals through
the nominal, no-bias, and no-friction conditions without changing its dynamics
or sample selection.
"""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import Optional

import numpy as np

from .wrench_solver import (
    ScaledDLSSolver,
    WrenchSolveResult,
    WrenchSolver,
)


def _finite_vector(value, name: str, size: Optional[int] = None) -> np.ndarray:
    vector = np.asarray(value, dtype=float).reshape(-1)
    if size is not None and vector.shape != (size,):
        raise ValueError(f"{name} must have shape ({size},), got {vector.shape}")
    if vector.size == 0 or not np.all(np.isfinite(vector)):
        raise ValueError(f"{name} must be a non-empty finite vector")
    return vector


def _immutable_vector(value, name: str, size: Optional[int] = None) -> np.ndarray:
    vector = _finite_vector(value, name, size).copy()
    vector.setflags(write=False)
    return vector


class AccelerationEstimateStatus(str, Enum):
    INITIALIZING = "initializing"
    VALID = "valid"
    INVALID_DT = "invalid_dt"


@dataclass(frozen=True)
class AccelerationEstimatorConfig:
    """Frozen one-pole differentiator configuration."""

    cutoff_hz: float
    minimum_dt_s: float
    maximum_dt_s: float

    def __post_init__(self) -> None:
        values = (self.cutoff_hz, self.minimum_dt_s, self.maximum_dt_s)
        if any(not np.isfinite(value) or value <= 0.0 for value in values):
            raise ValueError("acceleration settings must be finite and positive")
        if self.minimum_dt_s >= self.maximum_dt_s:
            raise ValueError("minimum_dt_s must be smaller than maximum_dt_s")

    @property
    def nominal_filter_delay_s(self) -> float:
        """Continuous one-pole time constant, logged as a delay descriptor."""
        return 1.0 / (2.0 * np.pi * self.cutoff_hz)


@dataclass(frozen=True)
class AccelerationEstimate:
    acceleration: np.ndarray
    status: AccelerationEstimateStatus
    dt_s: float
    raw_acceleration: np.ndarray

    def __post_init__(self) -> None:
        acceleration = _immutable_vector(self.acceleration, "acceleration")
        raw = _immutable_vector(
            self.raw_acceleration, "raw_acceleration", acceleration.size
        )
        if not np.isfinite(self.dt_s):
            raise ValueError("dt_s must be finite")
        if (
            self.status is not AccelerationEstimateStatus.INVALID_DT
            and self.dt_s < 0.0
        ):
            raise ValueError("dt_s must be non-negative for a valid estimate")
        object.__setattr__(self, "acceleration", acceleration)
        object.__setattr__(self, "raw_acceleration", raw)


class CausalAccelerationEstimator:
    """Backward difference followed by a causal one-pole low-pass filter."""

    def __init__(
        self,
        joint_count: int,
        config: AccelerationEstimatorConfig,
    ) -> None:
        if int(joint_count) <= 0:
            raise ValueError("joint_count must be positive")
        self.joint_count = int(joint_count)
        self.config = config
        self.reset()

    def reset(self) -> None:
        self._previous_velocity: Optional[np.ndarray] = None
        self._previous_time_s: Optional[float] = None
        self._filtered = np.zeros(self.joint_count, dtype=float)

    def update(
        self, joint_velocity: np.ndarray, timestamp_s: float
    ) -> AccelerationEstimate:
        velocity = _finite_vector(
            joint_velocity, "joint_velocity", self.joint_count
        )
        timestamp = float(timestamp_s)
        if not np.isfinite(timestamp):
            raise ValueError("timestamp_s must be finite")

        if self._previous_velocity is None:
            self._previous_velocity = velocity.copy()
            self._previous_time_s = timestamp
            zeros = np.zeros(self.joint_count, dtype=float)
            return AccelerationEstimate(
                acceleration=zeros,
                raw_acceleration=zeros,
                status=AccelerationEstimateStatus.INITIALIZING,
                dt_s=0.0,
            )

        dt = timestamp - float(self._previous_time_s)
        if dt < self.config.minimum_dt_s or dt > self.config.maximum_dt_s:
            # Invalid timestamps are observable and do not contaminate state.
            return AccelerationEstimate(
                acceleration=self._filtered,
                raw_acceleration=np.zeros(self.joint_count),
                status=AccelerationEstimateStatus.INVALID_DT,
                dt_s=dt,
            )

        raw = (velocity - self._previous_velocity) / dt
        alpha = 1.0 - np.exp(-2.0 * np.pi * self.config.cutoff_hz * dt)
        self._filtered = self._filtered + alpha * (raw - self._filtered)
        self._previous_velocity = velocity.copy()
        self._previous_time_s = timestamp
        return AccelerationEstimate(
            acceleration=self._filtered,
            raw_acceleration=raw,
            status=AccelerationEstimateStatus.VALID,
            dt_s=dt,
        )


@dataclass(frozen=True)
class JointFrictionCalibration:
    """Per-joint smooth Coulomb plus viscous friction coefficients."""

    coulomb_nm: np.ndarray
    viscous_nm_per_rad_s: np.ndarray
    smoothing_velocity_rad_s: float = 0.02

    def __post_init__(self) -> None:
        coulomb = _immutable_vector(self.coulomb_nm, "coulomb_nm")
        viscous = _immutable_vector(
            self.viscous_nm_per_rad_s,
            "viscous_nm_per_rad_s",
            coulomb.size,
        )
        if np.any(coulomb < 0.0) or np.any(viscous < 0.0):
            raise ValueError("friction magnitudes must be non-negative")
        smoothing = float(self.smoothing_velocity_rad_s)
        if not np.isfinite(smoothing) or smoothing <= 0.0:
            raise ValueError(
                "smoothing_velocity_rad_s must be finite and positive"
            )
        object.__setattr__(self, "coulomb_nm", coulomb)
        object.__setattr__(self, "viscous_nm_per_rad_s", viscous)
        object.__setattr__(self, "smoothing_velocity_rad_s", smoothing)

    @property
    def joint_count(self) -> int:
        return int(self.coulomb_nm.size)

    def torque(self, joint_velocity: np.ndarray) -> np.ndarray:
        velocity = _finite_vector(
            joint_velocity, "joint_velocity", self.joint_count
        )
        return (
            self.coulomb_nm
            * np.tanh(velocity / self.smoothing_velocity_rad_s)
            + self.viscous_nm_per_rad_s * velocity
        )

    @classmethod
    def zeros(
        cls,
        joint_count: int,
        *,
        smoothing_velocity_rad_s: float = 0.02,
    ) -> "JointFrictionCalibration":
        if int(joint_count) <= 0:
            raise ValueError("joint_count must be positive")
        zeros = np.zeros(int(joint_count), dtype=float)
        return cls(zeros, zeros, smoothing_velocity_rad_s)

    @classmethod
    def fit(
        cls,
        velocity_samples: np.ndarray,
        friction_torque_samples: np.ndarray,
        *,
        smoothing_velocity_rad_s: float = 0.02,
    ) -> "JointFrictionCalibration":
        """Fit independent smooth-Coulomb/viscous models by least squares."""
        velocity = np.asarray(velocity_samples, dtype=float)
        torque = np.asarray(friction_torque_samples, dtype=float)
        if velocity.ndim != 2 or torque.shape != velocity.shape:
            raise ValueError(
                "velocity_samples and friction_torque_samples must share "
                "shape (n_samples, n_joints)"
            )
        if velocity.shape[0] < 2 or velocity.shape[1] < 1:
            raise ValueError("at least two samples and one joint are required")
        if not np.all(np.isfinite(velocity)) or not np.all(np.isfinite(torque)):
            raise ValueError("friction calibration samples must be finite")
        smoothing = float(smoothing_velocity_rad_s)
        if not np.isfinite(smoothing) or smoothing <= 0.0:
            raise ValueError(
                "smoothing_velocity_rad_s must be finite and positive"
            )

        coulomb = np.zeros(velocity.shape[1], dtype=float)
        viscous = np.zeros(velocity.shape[1], dtype=float)
        for joint in range(velocity.shape[1]):
            design = np.column_stack(
                (
                    np.tanh(velocity[:, joint] / smoothing),
                    velocity[:, joint],
                )
            )
            coefficients, _, _, _ = np.linalg.lstsq(
                design, torque[:, joint], rcond=None
            )
            # Negative coefficients are non-physical and signal insufficient
            # excitation; clipping makes that decision explicit and stable.
            coulomb[joint], viscous[joint] = np.maximum(coefficients, 0.0)
        return cls(coulomb, viscous, smoothing)


@dataclass(frozen=True)
class WrenchEstimatorCalibration:
    """Immutable parameters selected only from the calibration split."""

    joint_bias_nm: np.ndarray
    friction: JointFrictionCalibration
    characteristic_length_m: float
    damping: float
    relative_rank_tolerance: float = 1e-9
    calibration_id: str = "unversioned"

    def __post_init__(self) -> None:
        bias = _immutable_vector(
            self.joint_bias_nm, "joint_bias_nm", self.friction.joint_count
        )
        positive = (
            self.characteristic_length_m,
            self.damping,
            self.relative_rank_tolerance,
        )
        if any(not np.isfinite(value) or value <= 0.0 for value in positive):
            raise ValueError("wrench calibration scalars must be positive")
        if self.relative_rank_tolerance >= 1.0:
            raise ValueError("relative_rank_tolerance must be smaller than one")
        if not isinstance(self.calibration_id, str) or not self.calibration_id:
            raise ValueError("calibration_id must be a non-empty string")
        object.__setattr__(self, "joint_bias_nm", bias)

    def build_dls_solver(self) -> ScaledDLSSolver:
        return ScaledDLSSolver(
            self.characteristic_length_m,
            self.damping,
            relative_rank_tolerance=self.relative_rank_tolerance,
        )


@dataclass(frozen=True)
class ResidualAblation:
    """Explicit H2 switches; both default to the nominal estimator."""

    use_bias_correction: bool = True
    use_friction_compensation: bool = True

    @classmethod
    def no_bias(cls) -> "ResidualAblation":
        return cls(use_bias_correction=False)

    @classmethod
    def no_friction(cls) -> "ResidualAblation":
        return cls(use_friction_compensation=False)


@dataclass(frozen=True)
class ResidualTorqueBreakdown:
    residual_nm: np.ndarray
    raw_residual_nm: np.ndarray
    applied_bias_nm: np.ndarray
    applied_friction_nm: np.ndarray
    ablation: ResidualAblation

    def __post_init__(self) -> None:
        residual = _immutable_vector(self.residual_nm, "residual_nm")
        size = residual.size
        object.__setattr__(
            self,
            "raw_residual_nm",
            _immutable_vector(self.raw_residual_nm, "raw_residual_nm", size),
        )
        object.__setattr__(
            self,
            "applied_bias_nm",
            _immutable_vector(self.applied_bias_nm, "applied_bias_nm", size),
        )
        object.__setattr__(
            self,
            "applied_friction_nm",
            _immutable_vector(
                self.applied_friction_nm, "applied_friction_nm", size
            ),
        )
        object.__setattr__(self, "residual_nm", residual)


class ResidualTorqueModel:
    """Apply frozen bias/friction calibration with replayable ablations."""

    def __init__(self, calibration: WrenchEstimatorCalibration) -> None:
        self.calibration = calibration

    def compute(
        self,
        measured_torque_nm: np.ndarray,
        rigid_body_torque_nm: np.ndarray,
        joint_velocity_rad_s: np.ndarray,
        *,
        ablation: ResidualAblation = ResidualAblation(),
    ) -> ResidualTorqueBreakdown:
        size = self.calibration.friction.joint_count
        measured = _finite_vector(
            measured_torque_nm, "measured_torque_nm", size
        )
        rigid = _finite_vector(
            rigid_body_torque_nm, "rigid_body_torque_nm", size
        )
        velocity = _finite_vector(
            joint_velocity_rad_s, "joint_velocity_rad_s", size
        )
        raw = measured - rigid
        bias = (
            self.calibration.joint_bias_nm
            if ablation.use_bias_correction
            else np.zeros(size)
        )
        friction = (
            self.calibration.friction.torque(velocity)
            if ablation.use_friction_compensation
            else np.zeros(size)
        )
        return ResidualTorqueBreakdown(
            residual_nm=raw - bias - friction,
            raw_residual_nm=raw,
            applied_bias_nm=bias,
            applied_friction_nm=friction,
            ablation=ablation,
        )


@dataclass(frozen=True)
class CalibratedWrenchEstimate:
    residual: ResidualTorqueBreakdown
    solve: WrenchSolveResult


class CalibratedResidualWrenchEstimator:
    """Compose residual conditioning and a selectable H2 wrench solver."""

    def __init__(
        self,
        calibration: WrenchEstimatorCalibration,
        solver: Optional[WrenchSolver] = None,
    ) -> None:
        self.calibration = calibration
        self.residual_model = ResidualTorqueModel(calibration)
        self.solver: WrenchSolver = (
            calibration.build_dls_solver() if solver is None else solver
        )
        if not np.isclose(
            self.solver.characteristic_length_m,
            calibration.characteristic_length_m,
        ):
            raise ValueError(
                "solver and calibration must use the same characteristic length"
            )
        if not np.isclose(
            self.solver.relative_rank_tolerance,
            calibration.relative_rank_tolerance,
        ):
            raise ValueError(
                "solver and calibration must use the same rank tolerance"
            )

    def estimate(
        self,
        jacobian: np.ndarray,
        measured_torque_nm: np.ndarray,
        rigid_body_torque_nm: np.ndarray,
        joint_velocity_rad_s: np.ndarray,
        *,
        ablation: ResidualAblation = ResidualAblation(),
    ) -> CalibratedWrenchEstimate:
        residual = self.residual_model.compute(
            measured_torque_nm,
            rigid_body_torque_nm,
            joint_velocity_rad_s,
            ablation=ablation,
        )
        solve = self.solver.solve(jacobian, residual.residual_nm)
        return CalibratedWrenchEstimate(residual=residual, solve=solve)