# CelerisLab/body/action_smoother.py
"""
ActionSmoother — ramp object actions (velocity/force) to avoid
instantaneous jumps that cause lattice instability.

Modes:
  NONE           — no smoothing (default)
  LINEAR_RAMP    — linear ramp from 0 to target over [n_min, n_max] steps
  EXPONENTIAL    — exponential approach: a(t) = target * (1 - exp(-t/tau))

Usage::

    smoother = ActionSmoother(mode="LINEAR_RAMP", n_min=0, n_max=500)
    for step in range(total):
        scaled = smoother.scale(raw_action, step)
        # pass scaled action to kernel
"""

import numpy as np
from dataclasses import dataclass
from enum import Enum


class SmoothMode(Enum):
    NONE = "NONE"
    LINEAR_RAMP = "LINEAR_RAMP"
    EXPONENTIAL = "EXPONENTIAL"


@dataclass
class ActionSmoother:
    """Scale actions with a ramp profile."""

    mode: str = "NONE"
    n_min: int = 0       # step at which ramp begins
    n_max: int = 500     # step at which ramp reaches 1.0
    tau: float = 100.0   # exponential time constant (steps)

    def scale(self, action: np.ndarray, step: int) -> np.ndarray:
        """Return action * alpha(step)."""
        alpha = self._alpha(step)
        if alpha >= 1.0:
            return action
        return action * np.float32(alpha)

    def _alpha(self, step: int) -> float:
        if self.mode == "NONE" or self.mode == SmoothMode.NONE.value:
            return 1.0

        if self.mode == "LINEAR_RAMP" or self.mode == SmoothMode.LINEAR_RAMP.value:
            if step <= self.n_min:
                return 0.0
            if step >= self.n_max:
                return 1.0
            return float(step - self.n_min) / float(self.n_max - self.n_min)

        if self.mode == "EXPONENTIAL" or self.mode == SmoothMode.EXPONENTIAL.value:
            if step <= self.n_min:
                return 0.0
            t = float(step - self.n_min)
            return 1.0 - float(np.exp(-t / max(self.tau, 1.0)))

        return 1.0
