"""Module A: ODE simulation and trajectory data collection.

Implements inverted pendulum dynamics and persistent excitation signal
generation for data-driven adaptive optimal control.
"""

import numpy as np
from scipy.integrate import solve_ivp


class DataCollector:
    """Inverted pendulum simulator with data collection capabilities.

    The pendulum dynamics are given by Eq. (27):
        dx₁/dt = x₂
        dx₂/dt = (mgl/J)·sin(x₁) + (1/J)·u
    """

    def __init__(self, J: float = 1.0, mgl: float = 1.0, n_freqs: int = 100):
        """
        Args:
            J: Moment of inertia (default 1.0).
            mgl: Mass × gravity × length product (default 1.0).
            n_freqs: Number of frequency components in PE signal.
        """
        self.J = float(J)
        self.mgl = float(mgl)
        self.n_freqs = n_freqs

        # Pre-compute PE signal parameters
        self._rng = np.random.RandomState(42)
        self._omegas = self._rng.uniform(0.5, 50.0, size=n_freqs)
        self._a = self._rng.uniform(-1.0, 1.0, size=n_freqs)
        self._b = self._rng.uniform(-1.0, 1.0, size=n_freqs)
        # Scale amplitudes so that max |pe_signal| ≈ 0.1
        raw_max = np.sum(np.abs(self._a) + np.abs(self._b))
        self._scale = 0.1 / raw_max if raw_max > 0 else 1.0

    def pe_signal(self, t):
        """Persistent excitation signal.

        u_e(t) = scale · Σ (a_k·sin(ω_k·t) + b_k·cos(ω_k·t))

        Args:
            t: Time value (scalar).

        Returns:
            float: PE signal value at time t.
        """
        t = float(t)
        signals = self._a * np.sin(self._omegas * t) + self._b * np.cos(self._omegas * t)
        return self._scale * np.sum(signals)

    def f(self, x):
        """Drift dynamics f(x) — Eq. (27).

        f(x) = [x₂, (mgl/J)·sin(x₁)]ᵀ

        Args:
            x: State vector (2,) or (N, 2).

        Returns:
            np.ndarray: Drift dynamics values.
        """
        x = np.asarray(x, dtype=np.float64)
        if x.ndim == 1:
            return np.array([x[1], (self.mgl / self.J) * np.sin(x[0])])
        return np.column_stack([x[:, 1], (self.mgl / self.J) * np.sin(x[:, 0])])

    def g(self, x):
        """Input matrix g(x) — Eq. (27).

        g(x) = [0, 1/J]ᵀ  (constant)

        Args:
            x: State vector (2,) or (N, 2) — ignored, kept for interface consistency.

        Returns:
            np.ndarray: Input matrix [0, 1/J]ᵀ.
        """
        x = np.asarray(x, dtype=np.float64)
        if x.ndim == 1:
            return np.array([0.0, 1.0 / self.J])
        N = x.shape[0]
        result = np.zeros((N, 2))
        result[:, 1] = 1.0 / self.J
        return result

    def dynamics(self, t, x, u_func):
        """ODE right-hand side: ẋ = f(x) + g(x)·u(t).

        Args:
            t: Current time.
            x: State vector (2,).
            u_func: Callable u_func(t, x) → control input.

        Returns:
            np.ndarray: State derivative (2,).
        """
        u_val = u_func(t, x)
        return self.f(x) + self.g(x) * u_val

    def collect_trajectory(self, x0, T, dt, u_func):
        """Collect a trajectory by simulating the system.

        Args:
            x0: Initial state (2,).
            T: Simulation duration.
            dt: Time step for output sampling.
            u_func: Control function u_func(t, x) → float.

        Returns:
            tuple: (t, X, U) where
                t: Time points (M,).
                X: State trajectory (M, 2).
                U: Control inputs (M,).
        """
        t_eval = np.arange(0, T, dt)
        t_span = (0.0, T)

        def ode_rhs(t, x):
            return self.dynamics(t, x, u_func)

        sol = solve_ivp(ode_rhs, t_span, x0, method="RK45",
                        t_eval=t_eval, rtol=1e-8, atol=1e-10)

        t = sol.t
        X = sol.y.T  # (M, 2)
        U = np.array([u_func(ti, xi) for ti, xi in zip(t, X)])

        return t, X, U