"""Action smoothing and physical-unit conversion.

Mimics the legacy FlowField.run() built-in exponential smoothing:
    action_pinned = (1 - weight) * action_pinned + weight * action_target

Two usage modes:
  (A) DRL inference — convert normalized PPO output to omega:
        raw_action = model.predict(obs)[0]              # [-1, 1] normalized
        smoothed = smoother(raw_action)                 # smoothed normalized
        omega = norm_action_to_omega(smoothed, scale=8, bias=[0,-4,4])
        for i, body_id in enumerate(pinball_ids):
            sim.set_body(body_id, omega=omega[i])

  (B) Bias FIFO — directly specify surface_vel, bypass scale/bias mapping:
        bias_surf = np.array([0.0, -4.0, 4.0]) * U0     # surface velocity
        bias_omega = surface_vel_to_omega(bias_surf)
        ema = ActionSmoother(weight=0.1)
        ema.reset(np.zeros(3))                          # legacy: starts from zero
        for _ in range(FIFO_LEN):
            s = ema(bias_omega)
            sim.set_body(fid, omega=s[0]); ...
            sim.run(SI, zero_obs=True)

IMPORTANT: New CelerisLab kernel has Uw = -omega * ry.
  The minus sign means omega > 0 produces CW rotation
  (opposite to naive expectation). Verified against legacy.
  omega = -surface_vel / radius
"""
from __future__ import annotations

from typing import Optional

import numpy as np

U0 = 0.01
RADIUS = 10.0  # pinball cylinder radius


class ActionSmoother:
    """Exponential moving-average action smoother.

    Matches legacy ``FlowField.run()`` smoothing:
        ``pinned = (1 - weight) * pinned + weight * target``

    Stateful across calls: call ``reset()`` to clear internal state.
    Use ``reset(np.zeros(3))`` for bias FIFO (legacy starts from zero).
    For DRL inference, reset to the bias-omega value before the episode.
    """

    def __init__(self, weight: float = 0.1):
        self.weight = float(weight)
        self._smoothed: Optional[np.ndarray] = None

    def __call__(self, target: np.ndarray) -> np.ndarray:
        """Apply exponential smoothing.  Returns smoothed copy."""
        t = np.asarray(target, dtype=np.float32)
        if self._smoothed is None:
            self._smoothed = t.copy()
        else:
            self._smoothed = (1.0 - self.weight) * self._smoothed + self.weight * t
        return self._smoothed.copy()

    def reset(self, value: Optional[np.ndarray] = None) -> None:
        """Reset smoother state.  Value=None means cold-start (first call
        will initialise from its argument).  Pass np.zeros(3) for bias FIFO."""
        if value is not None:
            self._smoothed = np.asarray(value, dtype=np.float32).copy()
        else:
            self._smoothed = None


# ── Physical-unit conversion ───────────────────────────────────────────

def norm_action_to_omega(
    action_norm: np.ndarray,
    scale: float = 8.0,
    bias: np.ndarray = None,
    u0: float = U0,
    radius: float = RADIUS,
) -> np.ndarray:
    """Convert PPO normalised action [-1, 1]^3 to angular velocity [lat-units].

    surface_vel = (action_norm * scale + bias) * u0
    omega = -surface_vel / radius   (new CelerisLab sign convention)
    """
    if bias is None:
        bias = np.zeros(3, dtype=np.float32)
    b = np.asarray(bias, dtype=np.float32)
    surface_vel = (np.asarray(action_norm, dtype=np.float32) * scale + b) * u0
    return -surface_vel / radius


def surface_vel_to_omega(
    surface_vel: np.ndarray,
    radius: float = RADIUS,
) -> np.ndarray:
    """Convert surface tangential velocity directly to angular velocity.

    Use this for bias FIFO where you know the exact surface_vel (e.g.
    bias_surf = [0, -4, 4] * U0) and don't want scale/bias remapping.
    """
    return -np.asarray(surface_vel, dtype=np.float32) / radius


def omega_to_norm_action(
    omega: np.ndarray,
    scale: float = 8.0,
    bias: np.ndarray = None,
    u0: float = U0,
    radius: float = RADIUS,
) -> np.ndarray:
    """Inverse of ``norm_action_to_omega`` — angular velocity to PPO action."""
    if bias is None:
        bias = np.zeros(3, dtype=np.float32)
    b = np.asarray(bias, dtype=np.float32)
    # omega = -surface_vel / R  →  surface_vel = -omega * R
    surface_vel = -np.asarray(omega, dtype=np.float32) * radius
    return np.clip((surface_vel / u0 - b) / scale, -1.0, 1.0)
