"""Input shaping utilities for trajectory position, velocity, and acceleration."""

from collections import deque
from typing import Deque

import numpy as np


class BaseShaper:
    """Shape scalar commands with a time-varying ZVD impulse response.

    Args:
        Ts: Controller sampling period [s].
    Side Effects:
        Initializes and mutates internal signal/history buffers used by subsequent calls.
    Raises:
        ValueError: If `Ts <= 0`.
    Preconditions:
        None.
    """

    _kMinDerivativeWindow = 3

    def __init__(self, Ts: float) -> None:
        """Initialize the shaper state and history buffers.

        Args:
            Ts: Controller sampling period [s].
        Returns:
            `None`.
        Side Effects:
            Allocates internal buffers and stores shaper configuration.
        Raises:
            ValueError: If `Ts <= 0`.
        Preconditions:
            None.
        """
        if Ts <= 0.0:
            raise ValueError("Sampling time must be positive")
        self.Ts: float = Ts

        self.M: int = 0
        self.impulse: np.ndarray = np.array([1.0])
        self._x_buf: list[float] = [0.0]
        self._x_head: int = 0

        self._y_buf: Deque[float] = deque(
            [0.0] * self._kMinDerivativeWindow,
            maxlen=self._kMinDerivativeWindow,
        )
        self._v_buf: Deque[float] = deque(
            [0.0] * self._kMinDerivativeWindow,
            maxlen=self._kMinDerivativeWindow,
        )
        self._a_buf: Deque[float] = deque(
            [0.0] * self._kMinDerivativeWindow,
            maxlen=self._kMinDerivativeWindow,
        )

        self._buffers_seeded: bool = False

    def shape_sample(
        self,
        x_i: float,
        frf_params: np.ndarray,
    ) -> tuple[float, float, float]:
        """Shape one sample and estimate shaped velocity and acceleration.

        Args:
            x_i: Raw command sample [joint unit].
            frf_params: Modal parameters shaped `(m, 2)` as `[wn, zeta]` where
                `wn` is in [rad/s] and `zeta` is unitless.
        Returns:
            `tuple[float, float, float]` containing `(y, y_dot, y_ddot)` in
            `[joint unit]`, `[joint unit/s]`, and `[joint unit/s^2]`.
        Side Effects:
            Updates impulse response, input ring buffer, and derivative histories.
        Raises:
            ValueError: If `frf_params` is empty or not shaped `(m, 2)`.
        Preconditions:
            `frf_params` rows contain physically valid modal parameters.
        """
        if (
            not isinstance(frf_params, np.ndarray)
            or frf_params.ndim != 2
            or frf_params.shape[1] != 2
        ):
            raise ValueError("frf_params must have shape (m, 2)")
        if frf_params.shape[0] == 0:
            raise ValueError("frf_params cannot be empty")

        # fmt: off
        impulse_new, M_new = self.compute_zvd_shaper(frf_params)
        self._ensure_input_buffer_size_on_change(M_new, seed=x_i)
        self.impulse = impulse_new
        self.M = M_new
        # fmt: on

        if not self._buffers_seeded:
            self._x_buf = [x_i] * (self.M + 1)
            self._x_head = 0
        else:
            n = len(self._x_buf)
            if n > 0:
                self._x_head = (self._x_head + n - 1) % n
                self._x_buf[self._x_head] = x_i

        x_arr = np.fromiter(
            (self._x_at_offset(k) for k in range(self.M + 1)),
            dtype=float,
            count=self.M + 1,
        )
        y = float(self.impulse.dot(x_arr))

        return self._update_output_history(y)

    def shape_trajectory(
        self,
        x: np.ndarray,
        varying_params: list[np.ndarray] | np.ndarray,
        *,
        reset: bool = True,
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Shape a full 1D trajectory using per-sample modal parameters.

        Args:
            x: Raw command trajectory of shape `(N,)` [joint unit].
            varying_params: Either one `(m, 2)` array, one list entry `(m, 2)`, or
                an `(N, m, 2)` array providing `[wn, zeta]` at each sample.
            reset: Whether to reset internal shaper state before processing.
        Returns:
            `tuple[np.ndarray, np.ndarray, np.ndarray]` as `(y, y_dot, y_ddot)`,
            each shape `(N,)` in `[joint unit]`, `[joint unit/s]`, and `[joint unit/s^2]`.
        Side Effects:
            Mutates internal buffers unless `reset=False` continues an existing stream.
        Raises:
            ValueError: If trajectory or modal parameters are malformed.
        Preconditions:
            `x` is finite and modal parameters are physically valid.
        """
        if not isinstance(x, np.ndarray) or x.ndim != 1 or x.size <= 0:
            raise ValueError("x must be a non-empty 1D numpy array")
        N = x.size

        vp_list = self._normalize_params_sequence(varying_params, N)

        if reset:
            self.reset_state()

        y = np.zeros(N, dtype=float)
        v = np.zeros(N, dtype=float)
        a = np.zeros(N, dtype=float)

        for i in range(N):
            y[i], v[i], a[i] = self.shape_sample(float(x[i]), frf_params=vp_list[i])

        return y, v, a

    def finalize(self) -> list[tuple[float, float, float]]:
        """Flush delayed shaper output by holding the last raw command.

        Args:
            None.
        Returns:
            `list[tuple[float, float, float]]` delayed `(y, y_dot, y_ddot)` samples.
        Side Effects:
            Advances input/output history buffers.
        Raises:
            None.
        Preconditions:
            None.
        """
        if (
            not self._buffers_seeded
            or self.M <= 0
            or len(self._x_buf) == 0
            or self.impulse.size == 0
        ):
            return []
        hold_input = float(self._x_at_offset(0))

        tail = []
        for _ in range(self.M):
            n = len(self._x_buf)
            if n > 0:
                self._x_head = (self._x_head + n - 1) % n
                self._x_buf[self._x_head] = hold_input

            x_arr = np.fromiter(
                (self._x_at_offset(k) for k in range(self.M + 1)),
                dtype=float,
                count=self.M + 1,
            )
            y = float(self.impulse.dot(x_arr))
            tail.append(self._update_output_history(y))

        return tail

    def compute_zvd_shaper(self, params_array: np.ndarray) -> tuple[np.ndarray, int]:
        """Build a combined ZVD impulse response for one or more vibration modes.

        Args:
            params_array: Array shaped `(m, 2)` containing `[wn, zeta]` rows.
        Returns:
            `tuple[np.ndarray, int]` containing impulse amplitudes `impulse` and filter order `M`.
        Side Effects:
            None.
        Raises:
            ValueError: If `params_array` is empty or not shaped `(m, 2)`.
        Preconditions:
            Every row has `wn > 0` and `0 < zeta < 1`.
        """
        if (
            not isinstance(params_array, np.ndarray)
            or params_array.ndim != 2
            or params_array.shape[1] != 2
        ):
            raise ValueError("params_array must have shape (m, 2)")
        if params_array.shape[0] == 0:
            raise ValueError("params_array cannot be empty")

        impulse, _ = self._compute_single_mode(
            float(params_array[0, 0]), float(params_array[0, 1])
        )
        for r in range(1, params_array.shape[0]):
            impulse_new, _ = self._compute_single_mode(
                float(params_array[r, 0]), float(params_array[r, 1])
            )
            impulse = np.convolve(impulse, impulse_new)
        M = len(impulse) - 1

        return impulse, M

    def reset_state(self) -> None:
        """Reset shaper state to a passthrough configuration.

        Args:
            None.
        Returns:
            `None`.
        Side Effects:
            Reinitializes impulse response and all signal/history buffers.
        Raises:
            None.
        Preconditions:
            None.
        """
        self.M = 0
        self.impulse = np.array([1.0])
        self._x_buf = [0.0]
        self._x_head = 0
        k = self._kMinDerivativeWindow
        self._y_buf = deque([0.0] * k, maxlen=k)
        self._v_buf = deque([0.0] * k, maxlen=k)
        self._a_buf = deque([0.0] * k, maxlen=k)
        self._buffers_seeded = False

    def _ensure_input_buffer_size_on_change(self, M_new: int, seed: float) -> None:
        """Resize internal buffers when the impulse response length changes.

        Args:
            M_new: New shaper order.
            seed: Fill value used for uninitialized history.
        Returns:
            `None`.
        Side Effects:
            Rebuilds input and derivative history buffers.
        Raises:
            None.
        Preconditions:
            `M_new >= 0`.
        """
        if M_new == self.M:
            return

        new_len = M_new + 1
        old_len = len(self._x_buf)

        if old_len == 0:
            self._x_buf = [seed] * new_len
            self._x_head = 0
        else:
            resized = [0.0] * new_len
            copy_len = min(old_len, new_len)
            # Copy newest-first order from current head
            for i in range(copy_len):
                resized[i] = self._x_buf[(self._x_head + i) % old_len]
            # Tail-fill with last copied value (or current head value)
            tail = (
                resized[copy_len - 1]
                if copy_len > 0
                else self._x_buf[self._x_head % old_len]
            )
            for i in range(copy_len, new_len):
                resized[i] = tail
            self._x_buf = resized
            self._x_head = 0

        history_len = max(new_len, self._kMinDerivativeWindow)

        def rebuild_hist(buf: Deque[float], fill: float) -> Deque[float]:
            """Return a history buffer resized to `history_len`.

            Args:
                buf: Existing buffer to copy from.
                fill: Fill value for uninitialized entries.
            Returns:
                `Deque[float]` resized history.
            Side Effects:
                None.
            Raises:
                None.
            Preconditions:
                None.
            """
            if not self._buffers_seeded:
                return deque([fill] * history_len, maxlen=history_len)

            data = list(buf)
            if len(data) >= history_len:
                data = data[:history_len]
            else:
                pad = [data[-1] if data else fill] * (history_len - len(data))
                data = data + pad
            return deque(data, maxlen=history_len)

        self._y_buf = rebuild_hist(self._y_buf, seed)
        self._v_buf = rebuild_hist(self._v_buf, 0.0)
        self._a_buf = rebuild_hist(self._a_buf, 0.0)

        self.M = M_new

    def _x_at_offset(self, offset: int) -> float:
        """Read the input ring buffer at `offset` samples from current head.

        Args:
            offset: Ring-buffer offset from newest sample.
        Returns:
            `float` buffered value.
        Side Effects:
            None.
        Raises:
            None.
        Preconditions:
            None.
        """
        n = len(self._x_buf)
        if n == 0:
            return 0.0
        return self._x_buf[(self._x_head + offset) % n]

    def _assign_history_len(self, buf: Deque[float], length: int, fill: float) -> None:
        """Rebuild one history deque to a fixed length with fill values.

        Args:
            buf: Target history buffer.
            length: Required history length.
            fill: Fill value for all positions.
        Returns:
            `None`.
        Side Effects:
            Clears and repopulates `buf`.
        Raises:
            None.
        Preconditions:
            `length > 0`.
        """
        new_buf = deque([fill] * length, maxlen=length)
        buf.clear()
        buf.extend(new_buf)

    def _push_front_with_len(
        self, buf: Deque[float], value: float, length: int
    ) -> None:
        """Push one value to history while enforcing buffer length.

        Args:
            buf: Target history buffer.
            value: New sample to place at index `0`.
            length: Required buffer length.
        Returns:
            `None`.
        Side Effects:
            Mutates `buf` ordering and contents.
        Raises:
            None.
        Preconditions:
            `length > 0`.
        """
        data = list(buf)
        new_buf = deque(data, maxlen=length)
        new_buf.appendleft(value)
        while len(new_buf) < length:
            new_buf.append(new_buf[-1] if len(new_buf) > 0 else value)
        buf.clear()
        buf.extend(new_buf)

    def _update_output_history(self, y: float) -> tuple[float, float, float]:
        """Update output history and compute discrete derivatives.

        Args:
            y: New shaped output sample [joint unit].
        Returns:
            `tuple[float, float, float]` as `(y, y_dot, y_ddot)`.
        Side Effects:
            Mutates position/velocity/acceleration history buffers.
        Raises:
            None.
        Preconditions:
            None.
        """
        history_len = max(self.M + 1, self._kMinDerivativeWindow)

        if not self._buffers_seeded:
            self._assign_history_len(self._y_buf, history_len, fill=y)
            self._assign_history_len(self._v_buf, history_len, fill=0.0)
            self._assign_history_len(self._a_buf, history_len, fill=0.0)
            self._buffers_seeded = True
            return y, 0.0, 0.0

        self._push_front_with_len(self._y_buf, y, history_len)
        self._push_front_with_len(self._v_buf, 0.0, history_len)
        self._push_front_with_len(self._a_buf, 0.0, history_len)

        if len(self._y_buf) >= 2:
            self._v_buf[0] = (self._y_buf[0] - self._y_buf[1]) / self.Ts
        else:
            self._v_buf[0] = 0.0

        if len(self._y_buf) >= 3:
            self._v_buf[1] = (self._y_buf[0] - self._y_buf[2]) / (2.0 * self.Ts)
            self._a_buf[1] = (
                self._y_buf[0] - 2.0 * self._y_buf[1] + self._y_buf[2]
            ) / (self.Ts * self.Ts)

        if len(self._v_buf) >= 2:
            self._a_buf[0] = (self._v_buf[0] - self._v_buf[1]) / self.Ts
        else:
            self._a_buf[0] = 0.0

        return y, self._v_buf[0], self._a_buf[0]

    def _normalize_params_sequence(
        self,
        varying_params: list[np.ndarray] | np.ndarray,
        N: int,
    ) -> list[np.ndarray]:
        """Normalize per-step modal parameters into a list of length `N`.

        Args:
            varying_params: One `(m, 2)` array, one-entry list, list of length `N`,
                or an `(N, m, 2)` array.
            N: Number of trajectory samples.
        Returns:
            `list[np.ndarray]` where each entry is shaped `(m_i, 2)`.
        Side Effects:
            None.
        Raises:
            ValueError: If `varying_params` cannot be broadcast to `N`.
        Preconditions:
            `N > 0`.
        """
        if isinstance(varying_params, np.ndarray):
            if varying_params.ndim == 2 and varying_params.shape[1] == 2:
                vp_list = [varying_params] * N
            elif (
                varying_params.ndim == 3
                and varying_params.shape[-1] == 2
                and varying_params.shape[0] == N
            ):
                vp_list = [varying_params[i] for i in range(N)]
            else:
                raise ValueError("varying_params must be (m,2) or (N,m,2) ndarray")
        elif isinstance(varying_params, list):
            if len(varying_params) == 1:
                vp_list = varying_params * N
            elif len(varying_params) == N:
                vp_list = varying_params
            else:
                raise ValueError("varying_params list must have length 1 or N")
        else:
            raise ValueError("varying_params must be a list or numpy array")

        first = vp_list[0]
        if (
            not isinstance(first, np.ndarray)
            or first.ndim != 2
            or first.shape[1] != 2
            or first.shape[0] == 0
        ):
            raise ValueError("each params item must be a non-empty (m,2) numpy array")
        return vp_list

    def _compute_single_mode(self, wn: float, zeta: float) -> tuple[np.ndarray, int]:
        """Build a single-mode ZVD impulse train using continuous-time impulses and
        linear impulse splitting.

        Args:
            wn: Natural frequency [rad/s].
            zeta: Damping ratio [unitless].
        Returns:
            `tuple[np.ndarray, int]` containing impulse vector and delay index `d`.
        Side Effects:
            None.
        Raises:
            ValueError: If `wn <= 0` or `zeta` is outside `(0, 1)`.
        Preconditions:
            `self.Ts > 0`.
        """
        if wn <= 0.0:
            raise ValueError("Natural frequency must be positive")
        if not (0.0 < zeta < 1.0):
            raise ValueError("Damping ratio must be between 0 and 1")

        wd = wn * np.sqrt(1.0 - zeta * zeta)

        k = np.exp(-zeta * np.pi / np.sqrt(1.0 - zeta * zeta))
        norm = 1.0 + 2.0 * k + k * k
        # Impulse amplitudes
        a0 = 1.0 / norm
        a1 = 2.0 * k / norm
        a2 = (k * k) / norm

        # Continuous-time impulse times
        t0 = 0.0
        t1 = np.pi / wd
        t2 = 2.0 * np.pi / wd

        # Real-valued sample indices
        n0 = t0 / self.Ts
        n1 = t1 / self.Ts
        n2 = t2 / self.Ts

        # Determine needed length (account for split into k+1)
        max_n = max(n0, n1, n2)
        m = int(np.ceil(max_n))  # last nonzero may land at floor(n)+1
        impulse = np.zeros(m + 1, dtype=float)

        def add_split(n: float, A: float) -> None:
            k_i = int(np.floor(n))
            d = float(n - k_i)  # in [0,1)

            # Ensure array is long enough for k_i+1
            nonlocal impulse
            nonlocal m
            need_M = k_i + (1 if d > 0.0 else 0)
            if need_M > m:
                # grow
                new_imp = np.zeros(need_M + 1, dtype=float)
                new_imp[: impulse.size] = impulse
                impulse = new_imp
                m = need_M

            # Construct impulse train
            impulse[k_i] += (1.0 - d) * A
            if d > 0.0:
                impulse[k_i + 1] += d * A

        add_split(n0, a0)
        add_split(n1, a1)
        add_split(n2, a2)

        return impulse, m
