#include "base_shaper.hpp"

#include <algorithm>
#include <deque>
#include <iostream>
#include <thread>

namespace control {

namespace {
constexpr std::size_t kMinDerivativeWindow = 3;
}

BaseShaper::BaseShaper(double Ts)
: Ts_(Ts), M_(0), I_(Eigen::VectorXd::Ones(1)), x_buf_(1, 0.0), x_head_(0),
  y_buf_(kMinDerivativeWindow, 0.0), v_buf_(kMinDerivativeWindow, 0.0),
  a_buf_(kMinDerivativeWindow, 0.0), buffers_seeded_(false) {
    if (Ts <= 0.0) throw std::invalid_argument("Sampling time must be positive");
}

/*
Summary:
  Reset internal buffers so the next call starts a fresh shaping stream.
Args:
  None.
Returns:
  None.
Side Effects:
  Clears and reinitializes `M_`, `I_`, circular input buffer, and derivative history buffers.
Raises:
  None.
Preconditions:
  The object has been constructed.
*/
void BaseShaper::resetState() {
    M_ = 0;
    I_.resize(1);
    I_(0) = 1.0;
    x_buf_.assign(1, 0.0);
    x_head_ = 0;
    y_buf_.clear();
    y_buf_.assign(kMinDerivativeWindow, 0.0);
    v_buf_.clear();
    v_buf_.assign(kMinDerivativeWindow, 0.0);
    a_buf_.clear();
    a_buf_.assign(kMinDerivativeWindow, 0.0);
    buffers_seeded_ = false;
}

/*
Summary:
  Resize input and history buffers when the active impulse length changes.
Args:
  M_new: New shaper delay in samples (impulse length is `M_new + 1`).
  seed: Fill value used when new entries must be created.
Returns:
  None.
Side Effects:
  Updates `M_`, input circular buffer layout, and derivative history buffer lengths.
  Preserves existing samples in time order to avoid a phase jump.
Raises:
  None.
Preconditions:
  `M_new` is derived from a valid impulse response for the current sample period.
*/
void BaseShaper::ensureInputBufferSizeOnChange(int M_new, double seed) {
    if (M_new == M_) return;

    const std::size_t new_len = static_cast<std::size_t>(M_new + 1);
    const std::size_t old_len = x_buf_.size();

    if (old_len == 0) {
        x_buf_.assign(new_len, seed);
        x_head_ = 0;
    } else {
        std::vector<double> resized(new_len);
        const std::size_t copy_len = std::min(old_len, new_len);
        for (std::size_t i = 0; i < copy_len; ++i) {
            resized[i] = x_buf_[(x_head_ + i) % old_len];
        }

        const double tail = copy_len ? resized[copy_len - 1]
                                     : x_buf_[(x_head_) % old_len];
        for (std::size_t i = copy_len; i < new_len; ++i) {
            resized[i] = tail;
        }

        x_buf_.swap(resized);
        x_head_ = 0;
    }

    M_ = M_new;

    const std::size_t history_len =
        std::max<std::size_t>(static_cast<std::size_t>(M_ + 1), kMinDerivativeWindow);

    auto resizeHistory = [&](std::deque<double>& buf, double fill_value) {
        if (!buffers_seeded_) {
            buf.assign(history_len, fill_value);
            return;
        }
        while (buf.size() > history_len) {
            buf.pop_back();
        }
        while (buf.size() < history_len) {
            const double tail = buf.empty() ? fill_value : buf.back();
            buf.push_back(tail);
        }
    };

    resizeHistory(y_buf_, seed);
    resizeHistory(v_buf_, 0.0);
    resizeHistory(a_buf_, 0.0);
}

/*
Summary:
  Insert the newest shaped position and compute aligned velocity/acceleration estimates.
Args:
  y: Current shaped position sample.
Returns:
  `std::tuple<double,double,double>` containing `(position, velocity, acceleration)`.
Side Effects:
  Mutates `y_buf_`, `v_buf_`, `a_buf_`, and may set `buffers_seeded_` on first call.
Raises:
  None.
Preconditions:
  `Ts_` is positive and buffer sizes are consistent with `M_`.
*/
std::tuple<double,double,double>
BaseShaper::updateOutputHistory(double y) {
    const std::size_t history_len =
        std::max<std::size_t>(static_cast<std::size_t>(M_ + 1), kMinDerivativeWindow);

    if (!buffers_seeded_) {
        y_buf_.assign(history_len, y);
        v_buf_.assign(history_len, 0.0);
        a_buf_.assign(history_len, 0.0);
        buffers_seeded_ = true;
        return {y, 0.0, 0.0};
    }

    y_buf_.push_front(y);
    if (y_buf_.size() > history_len) {
        y_buf_.pop_back();
    }
    while (y_buf_.size() < history_len) {
        const double tail = y_buf_.empty() ? y : y_buf_.back();
        y_buf_.push_back(tail);
    }

    v_buf_.push_front(0.0);
    if (v_buf_.size() > history_len) {
        v_buf_.pop_back();
    }
    while (v_buf_.size() < history_len) {
        const double tail = v_buf_.empty() ? 0.0 : v_buf_.back();
        v_buf_.push_back(tail);
    }

    a_buf_.push_front(0.0);
    if (a_buf_.size() > history_len) {
        a_buf_.pop_back();
    }
    while (a_buf_.size() < history_len) {
        const double tail = a_buf_.empty() ? 0.0 : a_buf_.back();
        a_buf_.push_back(tail);
    }

    if (y_buf_.size() >= 2) {
        v_buf_[0] = (y_buf_[0] - y_buf_[1]) / Ts_;
    } else {
        v_buf_[0] = 0.0;
    }

    if (y_buf_.size() >= 3) {
        v_buf_[1] = (y_buf_[0] - y_buf_[2]) / (2.0 * Ts_);
        a_buf_[1] = (y_buf_[0] - 2.0 * y_buf_[1] + y_buf_[2]) / (Ts_ * Ts_);
    }

    if (v_buf_.size() >= 2) {
        a_buf_[0] = (v_buf_[0] - v_buf_[1]) / Ts_;
    } else {
        a_buf_[0] = 0.0;
    }

    return {y, v_buf_[0], a_buf_[0]};
}

/*
Summary:
  Shape one command sample using the current modal parameters and return y/v/a.
Args:
  x_i: Raw command sample for one shaped joint axis.
  frf_params: Modal parameter matrix `[m x 2]` with columns `(wn [rad/s], zeta [-])`.
Returns:
  `std::tuple<double,double,double>` containing shaped `(position, velocity, acceleration)`.
Side Effects:
  Updates impulse cache usage, input/history buffers, and internal shaping state.
Raises:
  std::invalid_argument: If `frf_params` is empty or does not have exactly 2 columns.
Preconditions:
  `Ts_ > 0` and each modal row has physically valid parameters.
*/
std::tuple<double,double,double>
BaseShaper::shapeSample(double x_i, const Eigen::MatrixXd& frf_params) {
    // Validate params: (m,2), m>0
    if (frf_params.cols() != 2) {
        throw std::invalid_argument("frf_params must have shape (m,2)");
    }
    if (frf_params.rows() == 0) {
        throw std::invalid_argument("frf_params cannot be empty");
    }

    // Compute shaper & apply buffer resize if needed
    auto [I_new, M_new] = compute_zvd_shaper(frf_params);
    ensureInputBufferSizeOnChange(M_new, x_i);
    I_ = std::move(I_new);

    // Push newest raw sample to the front (newest-first)
    if (!buffers_seeded_) {
        std::fill(x_buf_.begin(), x_buf_.end(), x_i);
        x_head_ = 0;
    } else if (!x_buf_.empty()) {
        const std::size_t len = x_buf_.size();
        x_head_ = (x_head_ + len - 1) % len;
        x_buf_[x_head_] = x_i;
    }

    // OSA FIR: y = I · x_buf
    Eigen::VectorXd x_arr(M_ + 1);
    for (int i = 0; i <= M_; ++i) {
        x_arr(i) = xAtOffset(static_cast<std::size_t>(i));
    }
    double y = I_.dot(x_arr);

    return updateOutputHistory(y);
}

/*
Summary:
  Flush the FIR tail by replaying the held final input for `M_` extra samples.
Args:
  None.
Returns:
  `std::vector<std::tuple<double,double,double>>` tail samples of `(position, velocity, acceleration)`.
Side Effects:
  Advances input/history buffers as if additional samples were processed.
Raises:
  None.
Preconditions:
  Streaming state has been seeded by at least one prior `shapeSample` call.
*/
std::vector<std::tuple<double,double,double>>
BaseShaper::finalize() {
    std::vector<std::tuple<double,double,double>> tail;
    if (!buffers_seeded_ || M_ <= 0 || x_buf_.empty() || I_.size() == 0) {
        return tail;
    }

    const double hold_input = xAtOffset(0);

    tail.reserve(static_cast<std::size_t>(M_));
    for (int i = 0; i < M_; ++i) {
        if (!x_buf_.empty()) {
            const std::size_t len = x_buf_.size();
            x_head_ = (x_head_ + len - 1) % len;
            x_buf_[x_head_] = hold_input;
        }

        Eigen::VectorXd x_arr(M_ + 1);
        for (int k = 0; k <= M_; ++k) {
            x_arr(k) = xAtOffset(static_cast<std::size_t>(k));
        }
        double y = I_.dot(x_arr);
        tail.emplace_back(updateOutputHistory(y));
    }

    return tail;
}

std::tuple<Eigen::VectorXd, Eigen::VectorXd, Eigen::VectorXd>
BaseShaper::shapeTrajectory(const Eigen::VectorXd& x,
                            const std::vector<Eigen::MatrixXd>& varying_params,
                            bool reset) {
    if (x.size() <= 0) throw std::invalid_argument("x must be non-empty");
    if (varying_params.size() != static_cast<size_t>(x.size())) {
        throw std::invalid_argument("varying_params must have length N");
    }
    if (varying_params[0].cols() != 2) {
        throw std::invalid_argument("each params must have 2 columns");
    }

    if (reset) resetState();

    const int N = static_cast<int>(x.size());

    const unsigned int hw_threads = std::thread::hardware_concurrency();
    const unsigned int thread_cap = hw_threads == 0 ? 1u : hw_threads;
    constexpr unsigned int kMinChunk = 64;
    const unsigned int possible_threads =
        std::min<unsigned int>(thread_cap, static_cast<unsigned int>((N + kMinChunk - 1) / kMinChunk));
    const unsigned int thread_count = std::max(1u, possible_threads);
    if (reset && thread_count > 1) {
        std::cout << "[BaseShaper] Parallel shaping across " << thread_count
                  << " threads (chunk threshold=" << kMinChunk << " samples)."
                  << std::endl;
        std::vector<CachedImpulse> impulses;
        impulses.reserve(static_cast<std::size_t>(N));
        int max_M = 0;
        for (int i = 0; i < N; ++i) {
            CachedImpulse cached = compute_zvd_shaper(varying_params[static_cast<std::size_t>(i)]);
            max_M = std::max(max_M, cached.second);
            impulses.emplace_back(std::move(cached));
        }

        const std::size_t max_history_len =
            std::max<std::size_t>(static_cast<std::size_t>(max_M + 1), kMinDerivativeWindow);

        Eigen::VectorXd y = Eigen::VectorXd::Zero(N);
        const auto first_sample = x(0);
        const int block = (N + static_cast<int>(thread_count) - 1) / static_cast<int>(thread_count);
        std::vector<std::thread> workers;
        workers.reserve(thread_count);
        // Each worker computes an independent output block using convolution
        // against the precomputed per-sample impulse response.
        for (unsigned int t = 0; t < thread_count; ++t) {
            const int start = static_cast<int>(t) * block;
            const int end = std::min(N, start + block);
            if (start >= end) break;
            workers.emplace_back([&, start, end]() {
                for (int idx = start; idx < end; ++idx) {
                    const auto& impulse = impulses[static_cast<std::size_t>(idx)];
                    const Eigen::VectorXd& I = impulse.first;
                    const int M_local = impulse.second;
                    double y_val = 0.0;
                    for (int k = 0; k <= M_local; ++k) {
                        int sample_index = idx - k;
                        double x_val = sample_index >= 0 ? x(static_cast<Eigen::Index>(sample_index)) : first_sample;
                        y_val += I(k) * x_val;
                    }
                    y(static_cast<Eigen::Index>(idx)) = y_val;
                }
            });
        }
        for (auto& worker : workers) {
            worker.join();
        }

        Eigen::VectorXd v = Eigen::VectorXd::Zero(N);
        Eigen::VectorXd a = Eigen::VectorXd::Zero(N);

        std::deque<double> y_hist(max_history_len, y(0));
        std::deque<double> v_hist(max_history_len, 0.0);
        std::deque<double> a_hist(max_history_len, 0.0);

        if (N > 0) {
            v(0) = 0.0;
            a(0) = 0.0;
        }

        // Derivatives are evaluated in one serial pass to keep boundary
        // handling identical to the streaming path.
        for (int i = 1; i < N; ++i) {
            const double y_i = y(i);

            y_hist.push_front(y_i);
            if (y_hist.size() > max_history_len) {
                y_hist.pop_back();
            }
            while (y_hist.size() < max_history_len) {
                const double tail = y_hist.empty() ? y_i : y_hist.back();
                y_hist.push_back(tail);
            }

            v_hist.push_front(0.0);
            if (v_hist.size() > max_history_len) {
                v_hist.pop_back();
            }
            while (v_hist.size() < max_history_len) {
                const double tail = v_hist.empty() ? 0.0 : v_hist.back();
                v_hist.push_back(tail);
            }

            a_hist.push_front(0.0);
            if (a_hist.size() > max_history_len) {
                a_hist.pop_back();
            }
            while (a_hist.size() < max_history_len) {
                const double tail = a_hist.empty() ? 0.0 : a_hist.back();
                a_hist.push_back(tail);
            }

            if (y_hist.size() >= 2) {
                v_hist[0] = (y_hist[0] - y_hist[1]) / Ts_;
            } else {
                v_hist[0] = 0.0;
            }

            if (y_hist.size() >= 3) {
                v_hist[1] = (y_hist[0] - y_hist[2]) / (2.0 * Ts_);
                a_hist[1] = (y_hist[0] - 2.0 * y_hist[1] + y_hist[2]) / (Ts_ * Ts_);
            }

            if (v_hist.size() >= 2) {
                a_hist[0] = (v_hist[0] - v_hist[1]) / Ts_;
            } else {
                a_hist[0] = 0.0;
            }

            v(i) = v_hist[0];
            a(i) = a_hist[0];
        }

        if (N > 0) {
            const CachedImpulse& last_impulse = impulses.back();
            M_ = last_impulse.second;
            I_ = last_impulse.first;

            const std::size_t len = static_cast<std::size_t>(M_ + 1);
            x_buf_.assign(len, first_sample);
            x_head_ = 0;
            for (std::size_t k = 0; k < len; ++k) {
                int sample_index = N - 1 - static_cast<int>(k);
                double value = sample_index >= 0 ? x(static_cast<Eigen::Index>(sample_index)) : first_sample;
                x_buf_[k] = value;
            }

            const std::size_t final_history_len =
                std::max<std::size_t>(static_cast<std::size_t>(M_ + 1), kMinDerivativeWindow);

            while (y_hist.size() > final_history_len) {
                y_hist.pop_back();
            }
            while (y_hist.size() < final_history_len) {
                const double tail = y_hist.empty() ? y(N - 1) : y_hist.back();
                y_hist.push_back(tail);
            }

            while (v_hist.size() > final_history_len) {
                v_hist.pop_back();
            }
            while (v_hist.size() < final_history_len) {
                const double tail = v_hist.empty() ? 0.0 : v_hist.back();
                v_hist.push_back(tail);
            }

            while (a_hist.size() > final_history_len) {
                a_hist.pop_back();
            }
            while (a_hist.size() < final_history_len) {
                const double tail = a_hist.empty() ? 0.0 : a_hist.back();
                a_hist.push_back(tail);
            }

            y_buf_ = y_hist;
            v_buf_ = v_hist;
            a_buf_ = a_hist;
            buffers_seeded_ = true;
        }

        auto tail = finalize();
        if (!tail.empty()) {
            const Eigen::Index extra = static_cast<Eigen::Index>(tail.size());
            y.conservativeResize(N + extra);
            v.conservativeResize(N + extra);
            a.conservativeResize(N + extra);
            for (Eigen::Index i = 0; i < extra; ++i) {
                const auto& [y_i, v_i, a_i] = tail[static_cast<std::size_t>(i)];
                const Eigen::Index row = N + i;
                y(row) = y_i;
                v(row) = v_i;
                a(row) = a_i;
            }
        }

        return {y, v, a};
    }

    Eigen::VectorXd y(N), v(N), a(N);

    for (int i = 0; i < N; ++i) {
        auto [yi, vi, ai] = shapeSample(x(i), varying_params[static_cast<std::size_t>(i)]);
        y(i) = yi; v(i) = vi; a(i) = ai;
    }

    auto tail = finalize();
    if (!tail.empty()) {
        const Eigen::Index extra = static_cast<Eigen::Index>(tail.size());
        y.conservativeResize(N + extra);
        v.conservativeResize(N + extra);
        a.conservativeResize(N + extra);
        for (Eigen::Index i = 0; i < extra; ++i) {
            const auto& [y_i, v_i, a_i] = tail[static_cast<std::size_t>(i)];
            const Eigen::Index row = N + i;
            y(row) = y_i;
            v(row) = v_i;
            a(row) = a_i;
        }
    }

    return {y, v, a};
}

std::tuple<Eigen::VectorXd, Eigen::VectorXd, Eigen::VectorXd>
BaseShaper::shapeTrajectory(const Eigen::VectorXd& x,
                            const Eigen::MatrixXd& fixed_params,
                            bool reset) {
    if (x.size() <= 0) throw std::invalid_argument("x must be non-empty");
    if (fixed_params.cols() != 2 || fixed_params.rows() == 0) {
        throw std::invalid_argument("fixed_params must be non-empty (m,2)");
    }

    if (reset) resetState();

    const int N = static_cast<int>(x.size());
    Eigen::VectorXd y(N), v(N), a(N);

    for (int i = 0; i < N; ++i) {
        auto [yi, vi, ai] = shapeSample(x(i), fixed_params);
        y(i) = yi; v(i) = vi; a(i) = ai;
    }

    auto tail = finalize();
    if (!tail.empty()) {
        const Eigen::Index extra = static_cast<Eigen::Index>(tail.size());
        y.conservativeResize(N + extra);
        v.conservativeResize(N + extra);
        a.conservativeResize(N + extra);
        for (Eigen::Index i = 0; i < extra; ++i) {
            const auto& [y_i, v_i, a_i] = tail[static_cast<std::size_t>(i)];
            const Eigen::Index row = N + i;
            y(row) = y_i;
            v(row) = v_i;
            a(row) = a_i;
        }
    }

    return {y, v, a};
}

std::pair<Eigen::VectorXd, int>
/*
Summary:
  Build (and cache) a combined ZVD impulse response for one or more vibration modes.
Args:
  params_array: Modal parameter matrix `[m x 2]` with `(wn [rad/s], zeta [-])` rows.
Returns:
  `std::pair<Eigen::VectorXd,int>` containing `(impulse, delay_samples)`.
Side Effects:
  Inserts computed responses into `impulse_cache_`.
Raises:
  std::invalid_argument: If parameter matrix shape is invalid.
Preconditions:
  Each row has physically valid modal parameters.
*/
BaseShaper::compute_zvd_shaper(const Eigen::MatrixXd& params_array) {
    if (params_array.cols() != 2) {
        throw std::invalid_argument("params_array must have exactly 2 columns");
    }
    if (params_array.rows() == 0) {
        throw std::invalid_argument("params_array cannot be empty");
    }

    ModalKey key;
    key.data.reserve(static_cast<std::size_t>(params_array.rows() * 2));
    for (int r = 0; r < params_array.rows(); ++r) {
        key.data.push_back(params_array(r, 0));
        key.data.push_back(params_array(r, 1));
    }

    if (auto it = impulse_cache_.find(key); it != impulse_cache_.end()) {
        return it->second;
    }

    // Start with first mode
    auto [I, d0] = compute_single_mode(params_array(0,0), params_array(0,1));
    int M = static_cast<int>(I.size()) - 1;
    (void)d0;

    // Combine multi-mode shapers by convolving their impulse responses.
    for (int r = 1; r < params_array.rows(); ++r) {
        auto [I_new, d_new] = compute_single_mode(params_array(r,0), params_array(r,1));
        const int new_size = static_cast<int>(I.size() + I_new.size() - 1);
        Eigen::VectorXd I_conv = Eigen::VectorXd::Zero(new_size);
        for (int j = 0; j < I.size(); ++j) {
            for (int k = 0; k < I_new.size(); ++k) {
                I_conv(j + k) += I(j) * I_new(k);
            }
        }
        I.swap(I_conv);
        M = static_cast<int>(I.size()) - 1;
    }

    CachedImpulse cached{I, M};
    impulse_cache_.emplace(std::move(key), cached);
    return cached;
}

std::pair<Eigen::VectorXd, int>
/*
Summary:
  Compute a single-mode ZVD impulse response using continuous-time impulses and
  linear impulse splitting.
Args:
  wn: Natural frequency [rad/s].
  zeta: Damping ratio [-].
Returns:
  `std::pair<Eigen::VectorXd,int>` containing `(impulse, half-period delay in samples)`.
Side Effects:
  None.
Raises:
  std::invalid_argument: If `wn <= 0` or `zeta` is outside `(0,1)`.
Preconditions:
  Sample time `Ts_` is positive.
*/
BaseShaper::compute_single_mode(double wn, double zeta) {
    if (wn <= 0.0)  throw std::invalid_argument("Natural frequency must be positive");
    if (zeta <= 0.0 || zeta >= 1.0) throw std::invalid_argument("Damping ratio must be between 0 and 1");

    const double wd       = wn * std::sqrt(1.0 - zeta * zeta);
    const double k        = std::exp(-zeta * M_PI / std::sqrt(1.0 - zeta * zeta));

    const double norm = 1.0 + 2.0 * k + k * k;
    // Impulse amplitudes
    const double A0 = 1.0 / norm;
    const double A1 = 2.0 * k / norm;
    const double A2 = (k * k) / norm;

    // Continuous-time impulse times
    const double t0 = 0.0;
    const double t1 = M_PI / wd;
    const double t2 = 2.0 * M_PI / wd;

    // Real-valued sample indices
    const double n0 = t0 / Ts_;
    const double n1 = t1 / Ts_;
    const double n2 = t2 / Ts_;

    // Determine needed length (account for split into k+1)
    auto max_n = std::max(n0, std::max(n1, n2));
    int M = static_cast<int>(std::ceil(max_n)); // last nonzero may land at floor(n)+1

    Eigen::VectorXd impulse = Eigen::VectorXd::Zero(M + 1);

    auto add_split = [&](double n, double A) {
        const int k_i = static_cast<int>(std::floor(n));
        const double D = n - static_cast<double>(k_i); // [0,1)

        // Ensure array is long enough for k_i+1
        const int need_M = k_i + (D > 0.0 ? 1 : 0);
        if (need_M > M) {
            Eigen::VectorXd grown = Eigen::VectorXd::Zero(need_M + 1);
            grown.head(impulse.size()) = impulse;
            impulse.swap(grown);
            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};
}

/*
Summary:
  Read a sample from the circular input buffer at a fixed offset from newest.
Args:
  offset: Offset in samples where `0` is the newest input.
Returns:
  `double` buffered input value.
Side Effects:
  None.
Raises:
  None.
Preconditions:
  None.
*/
double BaseShaper::xAtOffset(std::size_t offset) const {
    if (x_buf_.empty()) return 0.0;
    const std::size_t len = x_buf_.size();
    return x_buf_[(x_head_ + offset) % len];
}

} // namespace control
