#pragma once

#include <vector>
#include <deque>
#include <tuple>
#include <unordered_map>
#include <functional>
#include <Eigen/Dense>
#include <stdexcept>

/**
 * @file base_shaper.hpp
 * @brief C++ implementation of input signal shaping for vibration suppression.
 *
 * This class implements Zero-Vibration-Derivative (ZVD) input shaping, supports
 * multiple vibration modes, and can operate online (sample-by-sample) or on full
 * trajectories. It now exposes shaped position, velocity, and acceleration to
 * mirror the Python version.
 */

namespace control {

/**
 * @brief Base signal shaper for vibration suppression using ZVD filters.
 *
 * Maintains internal input/history buffers to apply an FIR impulse train I[n].
 * Velocity and acceleration are tracked incrementally from the shaped stream.
 * A small rolling history allows us to back-fill centred finite differences for
 * interior samples while keeping the per-call update cost \(\mathcal{O}(1)\).
 */
class BaseShaper {
public:
    /**
     * @brief Construct a BaseShaper with specified sampling time.
     * @param Ts Sampling time in seconds (Ts>0).
     */
    explicit BaseShaper(double Ts);

    /**
     * @brief Shape a single input sample using current dynamics parameters.
     *
     * Recomputes the (possibly multi-mode) ZVD impulse vector each call, resizes
     * the internal input buffer if the filter length changes, then performs OSA-style
     * FIR convolution. Returns shaped position, velocity, acceleration.
     *
     * @param x_i        Input sample to shape.
     * @param frf_params Dynamics parameters [num_modes x 2] rows: [wn, zeta].
     * @return (y, v, a) = (shaped position, velocity, acceleration).
     */
    std::tuple<double, double, double>
    shapeSample(double x_i, const Eigen::MatrixXd& frf_params);

    /**
     * @brief Flush the shaper delay by holding the last shaped sample.
     *
     * Feeds the last raw input sample through the FIR for M additional steps,
     * returning the appended (y,v,a) tuples so callers can extend trajectories.
     */
    std::vector<std::tuple<double, double, double>> finalize();

    /**
     * @brief Current filter delay in samples (M).
     */
    int delaySamples() const noexcept { return M_; }

    /**
     * @brief Shape a trajectory with per-step dynamics (broadcast-safe).
     *
     * @param x              Input trajectory [N].
     * @param varying_params Vector of size N with params matrices (m_i x 2).
     * @param reset          If true, clears internal buffers before the run.
     * @return (Y, V, A) shaped arrays, each length N.
     */
    std::tuple<Eigen::VectorXd, Eigen::VectorXd, Eigen::VectorXd>
    shapeTrajectory(const Eigen::VectorXd& x,
                    const std::vector<Eigen::MatrixXd>& varying_params,
                    bool reset = true);

    /**
     * @brief Shape a trajectory with a single fixed shaper (broadcast to all steps).
     *
     * @param x            Input trajectory [N].
     * @param fixed_params Single params matrix (m x 2) used for all N steps.
     * @param reset        If true, clears internal buffers before the run.
     * @return (Y, V, A) shaped arrays, each length N.
     */
    std::tuple<Eigen::VectorXd, Eigen::VectorXd, Eigen::VectorXd>
    shapeTrajectory(const Eigen::VectorXd& x,
                    const Eigen::MatrixXd& fixed_params,
                    bool reset = true);

    // ---------------------------------------------------------------------
    // Backward-compatible thin wrappers (return position only)
    // ---------------------------------------------------------------------

    /**
     * @brief Legacy single-sample API: returns shaped position only.
     * Forwards to shapeSample() and discards v,a.
     */
    double shape_sample(double x_i, const Eigen::MatrixXd& frf_params) {
        auto [y, v, a] = shapeSample(x_i, frf_params);
        (void)v; (void)a;
        return y;
    }

    /**
     * @brief Legacy trajectory API: returns shaped position only.
     * Forwards to shapeTrajectory() and discards V,A.
     */
    Eigen::VectorXd shape_trajectory(const Eigen::VectorXd& x,
                                     const std::vector<Eigen::MatrixXd>& varying_params) {
        auto [Y, V, A] = shapeTrajectory(x, varying_params, /*reset=*/true);
        (void)V; (void)A;
        return Y;
    }

    // ---------------------------------------------------------------------
    // Shaper builders
    // ---------------------------------------------------------------------

    /**
     * @brief Compute ZVD shaper impulse response for given dynamics parameters.
     *
     * Convolves single-mode ZV filters for each row [wn, zeta] (modes).
     * @return {I, M} where I has length M+1 and M is the max delay index.
     */
    std::pair<Eigen::VectorXd, int>
    compute_zvd_shaper(const Eigen::MatrixXd& params_array);

private:
    struct ModalKey {
        std::vector<double> data;
        bool operator==(const ModalKey& other) const noexcept {
            return data == other.data;
        }
    };

    struct ModalKeyHash {
        std::size_t operator()(const ModalKey& key) const noexcept {
            std::size_t seed = key.data.size();
            for (double value : key.data) {
                std::size_t element = std::hash<double>{}(value);
                seed ^= element + 0x9e3779b97f4a7c15ULL + (seed << 6) + (seed >> 2);
            }
            return seed;
        }
    };

    using CachedImpulse = std::pair<Eigen::VectorXd, int>;

    /**
     * @brief Compute single-mode three-impulse ZV(D) shaper.
     *
     * Impulses at 0, d, 2d where d ≈ 0.5*T_damped/Ts. Amplitudes normalized by
     * (1 + 2k + k^2) with k = exp( -zeta*pi/sqrt(1-zeta^2) ).
     *
     * @param wn   Natural frequency [rad/s] (wn>0).
     * @param zeta Damping ratio (0<zeta<1).
     * @return {impulse, d}.
     */
    std::pair<Eigen::VectorXd, int> compute_single_mode(double wn, double zeta);

    // Internal utilities
    void resetState();                         ///< Clear buffers and set I = [1]
    void ensureInputBufferSizeOnChange(int M_new, double seed);
    std::tuple<double, double, double> updateOutputHistory(double y);

private:
    double Ts_;                 ///< Sampling time
    int    M_;                  ///< Current filter length - 1
    Eigen::VectorXd I_;         ///< Current impulse vector (size M_+1)
    std::vector<double> x_buf_; ///< Input buffer stored in a ring (newest head)
    std::size_t         x_head_;///< Index of newest element in x_buf_
    std::deque<double> y_buf_;  ///< Shaped history (newest at front)
    std::deque<double> v_buf_;  ///< Velocity history (newest at front)
    std::deque<double> a_buf_;  ///< Acceleration history (newest at front)
    bool buffers_seeded_;       ///< True once buffers are initialised with a real sample

    std::unordered_map<ModalKey, CachedImpulse, ModalKeyHash> impulse_cache_;

    double xAtOffset(std::size_t offset) const;
};

} // namespace control
