#include "covalent_wrapper.hpp"

#include <array>
#include <cmath>
#include <cstring>
#include <fstream>
#include <limits>
#include <sstream>
#include <stdexcept>
#include <string>
#include <tuple>
#include <utility>
#include <vector>

#include <Eigen/Geometry>

#include <pybind11/embed.h>
#include <pybind11/pybind11.h>
#include <pybind11/eigen.h>
#include <pybind11/numpy.h>
#include <pybind11/stl.h>

#include <torch/torch.h>

#include "map_loader.hpp"

namespace py = pybind11;
namespace covalentcore {

namespace {
namespace rd = reforge::dynamics;
constexpr double kDegPerRad = 180.0 / M_PI;
constexpr double kMmPerMeter = 1000.0;
constexpr const char* kIdentificationParamsFilename = "identification_parameters.csv";
constexpr const char* kDefaultModelDirectory = "src/robot/models/current";

/*
Summary:
  Initialize the embedded Python interpreter if it is not running.
Args:
  None.
Returns:
  None.
Side Effects:
  Starts a process-wide Python interpreter.
Raises:
  pybind11::error_already_set: If interpreter initialization fails.
Preconditions:
  Called before Python module imports in this translation unit.
*/
void ensurePythonInterpreter() {
    if (!Py_IsInitialized()) {
        py::initialize_interpreter();
    }
}

/*
Summary:
  Convert a Python sequence/array value to `Eigen::Vector3d`.
Args:
  obj: Python object expected to represent exactly three numeric values.
Returns:
  `Eigen::Vector3d` with Cartesian components.
Side Effects:
  None.
Raises:
  std::runtime_error: If `obj` is not length-3 data.
Preconditions:
  `obj` is convertible to numeric doubles.
*/
Eigen::Vector3d toVector3(py::handle obj) {
    if (py::isinstance<py::array>(obj)) {
        py::array arr = py::reinterpret_borrow<py::array>(obj);
        py::array_t<double, py::array::c_style | py::array::forcecast> arr_d(arr);
        py::buffer_info info = arr_d.request();
        if (info.ndim != 1 || info.shape[0] != 3) {
            throw std::runtime_error("Expected 1D array of length 3 for Cartesian position");
        }
        Eigen::Vector3d vec;
        std::memcpy(vec.data(), info.ptr, 3 * sizeof(double));
        return vec;
    }

    py::sequence seq = py::reinterpret_borrow<py::sequence>(obj);
    if (py::len(seq) != 3) {
        throw std::runtime_error("Expected sequence of length 3 for Cartesian position");
    }

    Eigen::Vector3d vec;
    for (ssize_t i = 0; i < 3; ++i) {
        vec[static_cast<Eigen::Index>(i)] = static_cast<double>(py::float_(seq[i]));
    }
    return vec;
}

/*
Summary:
  Convert a Python matrix-like object to a dense 3x3 matrix.
Args:
  obj: Python object expected to be a numeric 3x3 array.
Returns:
  `Eigen::Matrix3d` containing copied row-major values.
Side Effects:
  None.
Raises:
  std::runtime_error: If `obj` is not a 3x3 array.
Preconditions:
  `obj` supports NumPy buffer conversion.
*/
Eigen::Matrix3d toMatrix3(py::handle obj) {
    py::array arr = py::reinterpret_borrow<py::array>(obj);
    py::array_t<double, py::array::c_style | py::array::forcecast> arr_d(arr);
    py::buffer_info info = arr_d.request();
    if (info.ndim != 2 || info.shape[0] != 3 || info.shape[1] != 3) {
        throw std::runtime_error("Expected 3x3 matrix");
    }
    const double* data = static_cast<const double*>(info.ptr);
    Eigen::Map<const Eigen::Matrix<double, 3, 3, Eigen::RowMajor>> map(data);
    return map;
}

/*
Summary:
  Convert a Python matrix-like object to a 4x4 homogeneous transform matrix.
Args:
  obj: Python object expected to be a numeric 4x4 array.
Returns:
  `Eigen::Matrix4d` containing copied row-major values.
Side Effects:
  None.
Raises:
  std::runtime_error: If `obj` is not a 4x4 array.
Preconditions:
  `obj` supports NumPy buffer conversion.
*/
Eigen::Matrix4d toMatrix4(py::handle obj) {
    py::array arr = py::reinterpret_borrow<py::array>(obj);
    py::array_t<double, py::array::c_style | py::array::forcecast> arr_d(arr);
    py::buffer_info info = arr_d.request();
    if (info.ndim != 2 || info.shape[0] != 4 || info.shape[1] != 4) {
        throw std::runtime_error("Expected 4x4 matrix");
    }
    const double* data = static_cast<const double*>(info.ptr);
    Eigen::Map<const Eigen::Matrix<double, 4, 4, Eigen::RowMajor>> map(data);
    return map;
}

/*
Summary:
  Compute first derivative samples using NumPy-style edge/center differences.
Args:
  x: Input samples.
  Ts: Sample period [s].
Returns:
  `Eigen::VectorXd` derivative estimate with same length as `x`.
Side Effects:
  None.
Raises:
  None.
Preconditions:
  `Ts > 0`.
*/
Eigen::VectorXd gradient(const Eigen::VectorXd& x, double Ts) {
    const Eigen::Index N = x.size();
    Eigen::VectorXd g = Eigen::VectorXd::Zero(N);

    if (N <= 1) {
        return g;
    }

    if (N == 2) {
        const double diff = (x(1) - x(0)) / Ts;
        g(0) = diff;
        g(1) = diff;
        return g;
    }

    g(0) = (x(1) - x(0)) / Ts;
    for (Eigen::Index i = 1; i < N - 1; ++i) {
        g(i) = (x(i + 1) - x(i - 1)) / (2.0 * Ts);
    }
    g(N - 1) = (x(N - 1) - x(N - 2)) / Ts;
    return g;
}

/*
Summary:
  Convert a `std::vector<double>` to Eigen while enforcing exact length.
Args:
  data: Source vector values.
  expected_size: Required vector length.
  field_name: Human-readable field identifier for error reporting.
Returns:
  `Eigen::VectorXd` with `expected_size` elements.
Side Effects:
  None.
Raises:
  std::invalid_argument: If `data.size() != expected_size`.
Preconditions:
  `expected_size >= 0`.
*/
Eigen::VectorXd toEigenVectorExact(const std::vector<double>& data, int expected_size, const char* field_name) {
    if (static_cast<int>(data.size()) != expected_size) {
        throw std::invalid_argument(std::string(field_name) + " must contain exactly " +
                                    std::to_string(expected_size) + " elements");
    }
    Eigen::VectorXd vec(expected_size);
    for (int i = 0; i < expected_size; ++i) {
        vec[i] = data[static_cast<std::size_t>(i)];
    }
    return vec;
}

/*
Summary:
  Remove leading/trailing ASCII whitespace from a CSV token.
Args:
  input: Raw token string.
Returns:
  `std::string` trimmed token.
Side Effects:
  None.
Raises:
  None.
Preconditions:
  None.
*/
std::string trim(const std::string& input) {
    const auto first = input.find_first_not_of(" \t\r\n");
    if (first == std::string::npos) {
        return "";
    }
    const auto last = input.find_last_not_of(" \t\r\n");
    return input.substr(first, last - first + 1);
}

/*
Summary:
  Split one CSV line into comma-delimited, trimmed fields.
Args:
  line: CSV line to parse.
Returns:
  `std::vector<std::string>` field list.
Side Effects:
  None.
Raises:
  None.
Preconditions:
  Input line does not rely on quoted commas.
*/
std::vector<std::string> splitCsvLine(const std::string& line) {
    std::vector<std::string> fields;
    std::stringstream ss(line);
    std::string item;
    while (std::getline(ss, item, ',')) {
        fields.emplace_back(trim(item));
    }
    return fields;
}

struct IdentificationDefaults {
    double side_length = robotSideLength;
    double base_height = robotBaseHeight;
    int num_joints = robotNumJoints;
    int num_axes = robotNumShapedAxes;
};

std::string buildIdentificationParamsPath(const std::string& model_directory) {
    std::string base = model_directory.empty() ? kDefaultModelDirectory : model_directory;
    if (!base.empty() && base.back() != '/') {
        base.push_back('/');
    }
    base += kIdentificationParamsFilename;
    return base;
}

/*
Summary:
  Resolve geometry and joint-count defaults from CSV with fallback locations.
Args:
  model_directory: Preferred model directory for `identification_parameters.csv`.
Returns:
  `IdentificationDefaults` populated from CSV or compile-time constants.
Side Effects:
  Reads CSV files and writes warnings to stderr on malformed or missing data.
Raises:
  None.
Preconditions:
  None.
Fallback chain: user directory -> project default -> compile-time constants.
*/
IdentificationDefaults resolveIdentificationDefaults(const std::string& model_directory) {
    IdentificationDefaults defaults;

    std::string csv_path = buildIdentificationParamsPath(model_directory);
    std::ifstream file(csv_path);
    if (!file && model_directory != kDefaultModelDirectory) {
        csv_path = buildIdentificationParamsPath(kDefaultModelDirectory);
        file.open(csv_path);
    }

    if (!file) {
        std::cerr << "[ShaperInterface] Warning: identification_parameters.csv not found in "
                  << model_directory << " or " << kDefaultModelDirectory
                  << "; using built-in defaults." << std::endl;
        return defaults;
    }

    std::string header_line;
    std::string value_line;
    if (!std::getline(file, header_line) || !std::getline(file, value_line)) {
        std::cerr << "[ShaperInterface] Warning: identification_parameters.csv is missing header/value rows."
                  << std::endl;
        return defaults;
    }

    const auto headers = splitCsvLine(header_line);
    const auto values = splitCsvLine(value_line);
    if (headers.size() != values.size()) {
        std::cerr << "[ShaperInterface] Warning: identification_parameters.csv header/value mismatch."
                  << std::endl;
        return defaults;
    }

    std::unordered_map<std::string, std::string> row;
    row.reserve(headers.size());
    for (std::size_t i = 0; i < headers.size(); ++i) {
        row[headers[i]] = values[i];
    }

    auto read_double = [&](const char* key, double& out_value) {
        const auto it = row.find(key);
        if (it == row.end()) {
            return;
        }
        try {
            const double parsed = std::stod(it->second);
            if (parsed > 0.0) {
                out_value = parsed;
            } else {
                std::cerr << "[ShaperInterface] Warning: " << key
                          << " must be positive in identification_parameters.csv."
                          << std::endl;
            }
        } catch (const std::exception&) {
            std::cerr << "[ShaperInterface] Warning: failed to parse " << key
                      << " from identification_parameters.csv."
                      << std::endl;
        }
    };

    auto read_int = [&](const char* key, int& out_value) {
        const auto it = row.find(key);
        if (it == row.end()) {
            return;
        }
        try {
            const int parsed = std::stoi(it->second);
            if (parsed > 0) {
                out_value = parsed;
            } else {
                std::cerr << "[ShaperInterface] Warning: " << key
                          << " must be positive in identification_parameters.csv."
                          << std::endl;
            }
        } catch (const std::exception&) {
            std::cerr << "[ShaperInterface] Warning: failed to parse " << key
                      << " from identification_parameters.csv."
                      << std::endl;
        }
    };

    read_double("shoulder_len", defaults.side_length);
    read_double("base_height", defaults.base_height);
    read_int("num_joints", defaults.num_joints);
    read_int("axes", defaults.num_axes);

    return defaults;
}

} // namespace

namespace kinematics {

BaseRotationResult rotateXyzToZeroBase(const rd::RobotDynamics& dynamics,
                                       const Eigen::VectorXd& joint_angles,
                                       const Eigen::Vector3d& xyz) {
    BaseRotationResult result;
    result.rotated_xyz = xyz;

    Eigen::VectorXd joint_angles_zero = joint_angles;
    if (joint_angles_zero.size() > 0) {
        joint_angles_zero(0) = 0.0;
    }

    const auto rotations_current = dynamics.individualRotations(joint_angles);
    const auto rotations_zero = dynamics.individualRotations(joint_angles_zero);
    if (rotations_current.empty() || rotations_zero.empty()) {
        result.used_fallback = true;
        if (joint_angles.size() > 0) {
            const double base_angle = joint_angles(0);
            result.rotated_xyz = Eigen::AngleAxisd(-base_angle, Eigen::Vector3d::UnitZ()) * xyz;
        }
        return result;
    }

    const Eigen::Matrix3d base_rotation =
        rotations_current.front() * rotations_zero.front().transpose();
    result.rotated_xyz = base_rotation.transpose() * xyz;
    return result;
}

std::pair<double, double> cartesianToPolar(const Eigen::Vector3d& xyz,
                                           double side_length,
                                           double base_height) {
    // Match the Python feature-engineering convention: keep the dominant
    // horizontal axis as "reach", preserve z as height, and infer lateral depth.
    Eigen::Array3d abs_xyz = xyz.array().abs();
    double max_reach;
    int reach_axis;
    max_reach = abs_xyz.maxCoeff(&reach_axis);

    const double height = xyz(2);
    const int height_axis = 2;

    int depth_axis = -1;
    const double tol = 1e-9;
    for (int axis = 0; axis < 3; ++axis) {
        const bool matches_height = std::abs(abs_xyz(axis) - std::abs(height)) < tol;
        const bool matches_reach = std::abs(abs_xyz(axis) - max_reach) < tol;
        if (!matches_height && !matches_reach) {
            depth_axis = axis;
            break;
        }
    }

    if (depth_axis == -1) {
        for (int axis = 0; axis < 3; ++axis) {
            if (axis != height_axis) {
                depth_axis = axis;
                break;
            }
        }
    }

    double long_axis_value = xyz(reach_axis);
    double width_value = xyz(depth_axis);
    double height_value = xyz(height_axis);

    double dx = std::sqrt(long_axis_value * long_axis_value + std::pow(side_length - width_value, 2));
    double dz = height_value - base_height;

    double r = std::hypot(dx, dz);
    double v = std::atan2(dz, dx);
    return {v, r};
}

} // namespace kinematics

ShaperInterface::ShaperInterface(double sample_time,
                                 std::string model_directory,
                                 const std::string& python_src_root,
                                 const std::string& urdf_filepath,
                                 double side_length,
                                 double base_height,
                                 int num_axes,
                                 int num_joints,
                                 double prob_thresh)
    : Ts_(sample_time),
      num_axes_(num_axes),
      num_joints_(num_joints),
      side_length_(side_length),
      base_height_(base_height),
      prob_thresh_(prob_thresh),
      model_directory_(std::move(model_directory)) {

    const auto defaults = resolveIdentificationDefaults(model_directory_);
    if (!std::isfinite(side_length_)) {
        side_length_ = defaults.side_length;
    }
    if (!std::isfinite(base_height_)) {
        base_height_ = defaults.base_height;
    }
    if (num_axes_ <= 0) {
        num_axes_ = defaults.num_axes;
    }
    if (num_joints_ <= 0) {
        num_joints_ = defaults.num_joints;
    }

    if (Ts_ <= 0.0) {
        throw std::invalid_argument("Sample time must be positive");
    }
    if (num_axes_ <= 0 || num_axes_ > num_joints_) {
        throw std::invalid_argument("Number of shaped axes must be in (0, num_joints]");
    }
    if (num_joints_ <= 0) {
        throw std::invalid_argument("Number of joints must be positive");
    }
    ensurePythonInterpreter();
    ensurePythonPath(python_src_root);

    std::vector<Eigen::Vector3d> origins;
    std::vector<Eigen::Vector3d> rpy;
    std::vector<Eigen::Vector3d> axes;
    std::vector<double> masses;
    std::vector<Eigen::Vector3d> coms;
    std::vector<Eigen::Matrix3d> inertias;
    Eigen::Matrix4d base_prefix = Eigen::Matrix4d::Identity();
    Eigen::Matrix4d tool_offset = Eigen::Matrix4d::Identity();

    py::gil_scoped_acquire gil;
    py::object params_cls = py::module::import("util.urdf_to_dh.generate_dh").attr("GenerateDhParams");
    py::object params_node = params_cls(py::arg("urdf_file") = urdf_filepath);
    params_node.attr("parse_urdf")();
    params_node.attr("calculate_tfs_in_world_frame")();
    params_node.attr("calculate_params")();

    py::list moving_joints = params_node.attr("moving_joints");
    py::dict urdf_joints = params_node.attr("urdf_joints");
    py::dict inertia_dict = params_node.attr("inertia_dict");

    if (py::len(moving_joints) == 0) {
        moving_joints = urdf_joints.attr("keys")();
    }

    const auto joints_in_urdf = static_cast<int>(py::len(moving_joints));
    if (joints_in_urdf != num_joints_) {
        throw std::invalid_argument("URDF joint count does not match expected num_joints");
    }

    origins.resize(static_cast<std::size_t>(num_joints_));
    rpy.resize(static_cast<std::size_t>(num_joints_));
    axes.resize(static_cast<std::size_t>(num_joints_));
    masses.resize(static_cast<std::size_t>(num_joints_), 0.0);
    coms.resize(static_cast<std::size_t>(num_joints_), Eigen::Vector3d::Zero());
    inertias.resize(static_cast<std::size_t>(num_joints_), Eigen::Matrix3d::Zero());

    py::dict mass_dict = inertia_dict[py::str("mass")].cast<py::dict>();
    py::dict com_dict = inertia_dict[py::str("com")].cast<py::dict>();
    py::dict inertia_tensor_dict = inertia_dict[py::str("inertia_tensor")].cast<py::dict>();

    for (int i = 0; i < num_joints_; ++i) {
        const std::size_t idx = static_cast<std::size_t>(i);
        const py::object joint_name = moving_joints[i];
        if (!urdf_joints.contains(joint_name)) {
            throw std::runtime_error("URDF joint entry missing for expected joint index");
        }
        py::dict joint_data = urdf_joints[joint_name].cast<py::dict>();

        origins[idx] = toVector3(joint_data[py::str("xyz")]);
        rpy[idx] = toVector3(joint_data[py::str("rpy")]);
        axes[idx] = toVector3(joint_data[py::str("axis")]);

        const py::int_ key(i);
        if (!mass_dict.contains(key) ||
            !com_dict.contains(key) ||
            !inertia_tensor_dict.contains(key)) {
            std::cerr << "[ShaperInterface] Missing inertia entry for joint " << i
                      << " (filling zeros)\n";
        }

        masses[idx] = mass_dict.contains(key) ? mass_dict[key].cast<double>() : 0.0;
        coms[idx] = com_dict.contains(key) ? toVector3(com_dict[key]) : Eigen::Vector3d::Zero();
        inertias[idx] =
            inertia_tensor_dict.contains(key) ? toMatrix3(inertia_tensor_dict[key]) : Eigen::Matrix3d::Zero();
    }

    if (!params_node.attr("base_prefix_tf").is_none()) {
        base_prefix = toMatrix4(params_node.attr("base_prefix_tf"));
    }
    if (!params_node.attr("tool_offset_tf").is_none()) {
        tool_offset = toMatrix4(params_node.attr("tool_offset_tf"));
    }

    robot_dynamics_ = std::make_unique<rd::RobotDynamics>(std::move(origins),
                                                          std::move(rpy),
                                                          std::move(axes),
                                                          std::move(masses),
                                                          std::move(coms),
                                                          std::move(inertias),
                                                          base_prefix,
                                                          tool_offset);
    if (!robot_dynamics_) {
        throw std::runtime_error("Failed to initialize C++ robot dynamics model");
    }

    map_fitter_ = maploader::ModelLoader::load(model_directory_, num_axes_, /*input_features=*/3, {64, 64});
    if (!map_fitter_) {
        throw std::runtime_error("Failed to load neural network models from directory: " + model_directory_);
    }

    shapers_.reserve(static_cast<std::size_t>(num_axes_));
    for (int axis = 0; axis < num_axes_; ++axis) {
        shapers_.emplace_back(Ts_);
    }
}

/*
Summary:
  Reset all per-axis shaper state and clear continuity markers.
Args:
  None.
Returns:
  None.
Side Effects:
  Reconstructs each `BaseShaper` instance and clears cached last input vectors.
Raises:
  None.
Preconditions:
  Interface has been constructed.
*/
void ShaperInterface::reset() {
    for (auto& shaper : shapers_) {
        shaper = control::BaseShaper(Ts_);
    }
    last_input_position_.reset();
    last_input_velocity_.reset();
    last_input_acceleration_.reset();
}

/*
Summary:
  Handle vibration-mode transitions and reset stale shaper history when needed.
Args:
  enable_vibration_control: Requested mode for the upcoming processing call.
Returns:
  `bool` true when a reset was performed due to a mode transition.
Side Effects:
  May reset per-axis shapers and updates `last_vibration_control_enabled_`.
Raises:
  None.
Preconditions:
  Interface has been constructed.
*/
bool ShaperInterface::handleVibrationModeTransition(bool enable_vibration_control) {
    const bool toggled_disabled_to_enabled =
        enable_vibration_control &&
        last_vibration_control_enabled_.has_value() &&
        !(*last_vibration_control_enabled_);

    if (toggled_disabled_to_enabled) {
        // We intentionally reset only on disabled->enabled so the next shaped
        // output does not use stale FIR history that was not updated while
        // pass-through mode was active.
        reset();
    }

    last_vibration_control_enabled_ = enable_vibration_control;
    return toggled_disabled_to_enabled;
}

/*
Summary:
  Process one command sample with optional vibration shaping.
Args:
  command: Joint command vector `[num_joints]` in radians.
  state: Robot state used for map feature computation.
  command_dot: Optional command velocity `[num_joints]` in radians/s.
  command_ddot: Optional command acceleration `[num_joints]` in radians/s^2.
  enable_vibration_control: When false, bypass shaping and return pass-through.
Returns:
  `ShapedSample` containing position/velocity/acceleration output vectors.
Side Effects:
  May update per-axis shaper states and stream continuity cache.
Raises:
  std::invalid_argument: If command or derivative vector sizes are invalid.
Preconditions:
  Interface is initialized and command dimension matches `num_joints_`.
*/
ShapedSample ShaperInterface::processSample(
    const Eigen::VectorXd& command,
    const RobotState& state,
    std::optional<Eigen::VectorXd> command_dot,
    std::optional<Eigen::VectorXd> command_ddot,
    bool enable_vibration_control) {
    handleVibrationModeTransition(enable_vibration_control);
    
    // Compute first and second derivative from finite differences?
    bool dot_from_fd = !command_dot.has_value();
    bool ddot_from_fd = !command_ddot.has_value();
    if (static_cast<int>(command.size()) == 0) {
        throw std::invalid_argument("Command must not be empty");
    }
    if (static_cast<int>(command.size()) != num_joints_) {
        throw std::invalid_argument("Command length mismatch");
    }
    if (!dot_from_fd && static_cast<int>(command_dot->size()) != num_joints_) {
         throw std::invalid_argument("Command_dot (derivative) length mismatch");
    }
    if (!ddot_from_fd && static_cast<int>(command_ddot->size()) != num_joints_) {
         throw std::invalid_argument("Command_ddot (2nd derivative) length mismatch");
    }

    if (!enable_vibration_control) {
        ShapedSample passthrough;
        passthrough.positions = command;

        if (dot_from_fd) {
            if (last_input_position_.has_value() &&
                last_input_position_->size() == num_joints_) {
                passthrough.velocities = (command - *last_input_position_) / Ts_;
            } else {
                passthrough.velocities = Eigen::VectorXd::Zero(num_joints_);
            }
        } else {
            passthrough.velocities = *command_dot;
        }

        if (ddot_from_fd) {
            if (last_input_velocity_.has_value() &&
                last_input_velocity_->size() == num_joints_) {
                passthrough.accelerations =
                    (passthrough.velocities - *last_input_velocity_) / Ts_;
            } else {
                passthrough.accelerations = Eigen::VectorXd::Zero(num_joints_);
            }
        } else {
            passthrough.accelerations = *command_ddot;
        }

        last_input_position_ = passthrough.positions;
        last_input_velocity_ = passthrough.velocities;
        last_input_acceleration_ = passthrough.accelerations;
        return passthrough;
    }

    RobotState state_copy = state;
    if (state_copy.joint_angles.size() != num_joints_) {
        state_copy.joint_angles = command;
    }
    if (!state_copy.tcp_position.has_value()) {
        state_copy.tcp_position = computeForwardKinematics(state_copy.joint_angles);
    }

    NNInputs inputs = computeNNInputs(state_copy);

    ShapedSample out;
    out.positions = command;
    out.velocities = dot_from_fd ? Eigen::VectorXd::Zero(num_joints_) : *command_dot;
    out.accelerations = ddot_from_fd ? Eigen::VectorXd::Zero(num_joints_) : *command_ddot;

    torch::NoGradGuard guard;
    torch::Tensor features = torch::empty({num_axes_, 3}, torch::dtype(torch::kFloat32));
    auto features_acc = features.accessor<float, 2>();
    const float v_deg = static_cast<float>(inputs.v_rad * kDegPerRad);
    const float r_mm = static_cast<float>(inputs.r_m * kMmPerMeter);
    for (int axis = 0; axis < num_axes_; ++axis) {
        features_acc[axis][0] = v_deg;
        features_acc[axis][1] = r_mm;
        features_acc[axis][2] = static_cast<float>(inputs.inertia(axis));
    }

    auto outputs = map_fitter_->inferBatch(features);
    torch::Tensor order_probs = outputs.first.squeeze(-1).to(torch::kDouble).contiguous();
    torch::Tensor mode_params = outputs.second.to(torch::kDouble).contiguous();

    auto order_acc = order_probs.accessor<double, 1>();
    auto modes_acc = mode_params.accessor<double, 2>();

    for (int axis = 0; axis < num_axes_; ++axis) {
        const double prob_second = order_acc[axis];
        const double wn1 = modes_acc[axis][0];
        const double z1_log = modes_acc[axis][1];
        const double wn2 = modes_acc[axis][2];
        const double z2_log = modes_acc[axis][3];

        const double z1 = std::exp(z1_log);
        const double z2 = std::exp(z2_log);
        const double wn1_rad = wn1;
        const double wn2_rad = wn2;

        const bool two_mode = prob_second > prob_thresh_;

        static std::array<bool, 6> logged_once{false, false, false, false, false, false};
        if (!logged_once[static_cast<std::size_t>(axis)]) {
            std::cout << "[ShaperInterface] Axis " << axis << " inference: wn1=" << wn1_rad
                      << " z1=" << z1 << " prob_second=" << prob_second << std::endl;
            if (two_mode) {
                std::cout << "[ShaperInterface] Axis " << axis << " second mode wn2=" << wn2_rad
                          << " z2=" << z2 << std::endl;
            }
            logged_once[static_cast<std::size_t>(axis)] = true;
        }

        Eigen::MatrixXd params(two_mode ? 2 : 1, 2);
        params(0, 0) = wn1_rad;
        params(0, 1) = z1;
        if (two_mode) {
            params(1, 0) = wn2_rad;
            params(1, 1) = z2;
        }

        try {
            auto [y, v, a] = shapers_[static_cast<std::size_t>(axis)].shapeSample(command(axis), params);
            out.positions(axis) = y;
            out.velocities(axis) = v;
            out.accelerations(axis) = a;
        } catch (const std::invalid_argument& ex) {
            std::cerr << "[ShaperInterface] Warning: axis " << axis
                      << " shaping skipped due to invalid parameters (" << ex.what() << ")" << std::endl;
        }
    }

    last_input_position_ = out.positions;
    last_input_velocity_ = out.velocities;
    last_input_acceleration_ = out.accelerations;

    return out;
}

/*
Summary:
  Process a full trajectory with optional vibration shaping.
Args:
  command: Position trajectory `[N, num_joints]` in radians.
  states: Optional per-sample states; empty means derive from command only.
  command_dot: Optional velocity trajectory `[N, num_joints]` in radians/s.
  command_ddot: Optional acceleration trajectory `[N, num_joints]` in radians/s^2.
  time_vector: Optional time vector `[N]` in seconds.
  enable_vibration_control: When false, bypass shaping and return pass-through.
Returns:
  `ShapedTrajectory` containing positions/velocities/accelerations and time.
Side Effects:
  May reset/advance shaper state and updates stream continuity cache.
Raises:
  std::invalid_argument: If trajectory matrix/vector dimensions are inconsistent.
Preconditions:
  Interface is initialized and command has `num_joints_` columns.
*/
ShapedTrajectory ShaperInterface::processTrajectory(
    const Eigen::MatrixXd& command,
    const std::vector<RobotState>& states,
    std::optional<Eigen::MatrixXd> command_dot,
    std::optional<Eigen::MatrixXd> command_ddot,
    std::optional<std::vector<double>> time_vector,
    bool enable_vibration_control) {
    const bool reset_for_mode_change = handleVibrationModeTransition(enable_vibration_control);

    if (command.rows() == 0) {
        throw std::invalid_argument("Trajectory command must contain at least one sample");
    }
    if (command.cols() != num_joints_) {
        throw std::invalid_argument("Trajectory command must have num_joints columns");
    }
    if (!states.empty() && states.size() != static_cast<std::size_t>(command.rows())) {
        throw std::invalid_argument("states length must match number of trajectory samples");
    }

    const Eigen::Index samples = command.rows();

    Eigen::MatrixXd velocity_matrix;
    if (command_dot.has_value()) {
        if (command_dot->rows() != samples || command_dot->cols() != command.cols()) {
            throw std::invalid_argument("command_dot shape mismatch");
        }
        velocity_matrix = *command_dot;
    } else {
        velocity_matrix = finiteDifferenceFirst(command);
    }

    Eigen::MatrixXd acceleration_matrix;
    if (command_ddot.has_value()) {
        if (command_ddot->rows() != samples || command_ddot->cols() != command.cols()) {
            throw std::invalid_argument("command_ddot shape mismatch");
        }
        acceleration_matrix = *command_ddot;
    } else {
        acceleration_matrix = finiteDifferenceSecond(velocity_matrix);
    }

    std::vector<double> times;
    if (time_vector.has_value()) {
        if (time_vector->size() != static_cast<std::size_t>(samples)) {
            throw std::invalid_argument("time_vector length must match number of samples");
        }
        times = *time_vector;
    } else {
        times.resize(static_cast<std::size_t>(samples));
        for (std::size_t i = 0; i < times.size(); ++i) {
            times[i] = static_cast<double>(i) * Ts_;
        }
    }

    const Eigen::VectorXd first_position = command.row(0).transpose();
    const Eigen::VectorXd first_velocity = velocity_matrix.row(0).transpose();
    const Eigen::VectorXd first_acceleration = acceleration_matrix.row(0).transpose();

    if (!enable_vibration_control) {
        ShapedTrajectory passthrough;
        passthrough.positions = command;
        passthrough.velocities = velocity_matrix;
        passthrough.accelerations = acceleration_matrix;
        passthrough.time = Eigen::Map<const Eigen::VectorXd>(times.data(), static_cast<Eigen::Index>(times.size()));

        last_input_position_ = first_position;
        last_input_velocity_ = first_velocity;
        last_input_acceleration_ = first_acceleration;
        if (samples > 1) {
            last_input_position_ = command.row(samples - 1).transpose();
            last_input_velocity_ = velocity_matrix.row(samples - 1).transpose();
            last_input_acceleration_ = acceleration_matrix.row(samples - 1).transpose();
        }
        return passthrough;
    }

    bool continue_stream = false;
    if (last_input_position_.has_value() && last_input_velocity_.has_value() &&
        last_input_acceleration_.has_value()) {
        const bool position_match = last_input_position_->size() == first_position.size() &&
                                    last_input_position_->isApprox(first_position);
        const bool velocity_match = last_input_velocity_->size() == first_velocity.size() &&
                                    last_input_velocity_->isApprox(first_velocity);
        const bool acceleration_match = last_input_acceleration_->size() == first_acceleration.size() &&
                                        last_input_acceleration_->isApprox(first_acceleration);
        continue_stream = position_match && velocity_match && acceleration_match;
    }

    if (!continue_stream && !reset_for_mode_change) {
        reset();
    }

    ShapedTrajectory traj;
    traj.positions = Eigen::MatrixXd::Zero(samples, num_joints_);
    traj.velocities = Eigen::MatrixXd::Zero(samples, num_joints_);
    traj.accelerations = Eigen::MatrixXd::Zero(samples, num_joints_);
    traj.time = Eigen::VectorXd::Zero(samples);

    for (Eigen::Index i = 0; i < samples; ++i) {
        traj.time(i) = times[static_cast<std::size_t>(i)];
        RobotState state_i = states.empty() ? RobotState{} : states[static_cast<std::size_t>(i)];
        state_i.joint_angles = command.row(i).transpose();
        if (!state_i.tcp_position.has_value()) {
            state_i.tcp_position = computeForwardKinematics(state_i.joint_angles);
        }

        Eigen::VectorXd cmd = command.row(i).transpose();
        Eigen::VectorXd vel = velocity_matrix.row(i).transpose();
        Eigen::VectorXd acc = acceleration_matrix.row(i).transpose();

        ShapedSample sample = processSample(cmd, state_i, vel, acc, /*enable_vibration_control=*/true);
        traj.positions.row(i) = sample.positions.transpose();
        traj.velocities.row(i) = sample.velocities.transpose();
        traj.accelerations.row(i) = sample.accelerations.transpose();
    }

    std::vector<std::vector<std::tuple<double, double, double>>> tails(static_cast<std::size_t>(num_axes_));
    std::size_t max_tail = 0;
    for (int axis = 0; axis < num_axes_; ++axis) {
        tails[static_cast<std::size_t>(axis)] = shapers_[static_cast<std::size_t>(axis)].finalize();
        max_tail = std::max<std::size_t>(max_tail, tails[static_cast<std::size_t>(axis)].size());
    }

    if (max_tail > 0) {
        const Eigen::Index extra = static_cast<Eigen::Index>(max_tail);
        const Eigen::VectorXd last_pos = traj.positions.row(samples - 1).transpose();

        traj.positions.conservativeResize(samples + extra, num_joints_);
        traj.velocities.conservativeResize(samples + extra, num_joints_);
        traj.accelerations.conservativeResize(samples + extra, num_joints_);
        traj.time.conservativeResize(samples + extra);

        const double last_time = times.empty()
            ? static_cast<double>(samples - 1) * Ts_
            : times.back();

        for (std::size_t t = 0; t < max_tail; ++t) {
            const Eigen::Index row = samples + static_cast<Eigen::Index>(t);
            Eigen::VectorXd pos = last_pos;
            Eigen::VectorXd vel = Eigen::VectorXd::Zero(num_joints_);
            Eigen::VectorXd acc = Eigen::VectorXd::Zero(num_joints_);

            for (int axis = 0; axis < num_axes_; ++axis) {
                if (t < tails[static_cast<std::size_t>(axis)].size()) {
                    const auto& [y, v, a] = tails[static_cast<std::size_t>(axis)][static_cast<std::size_t>(t)];
                    pos(axis) = y;
                    vel(axis) = v;
                    acc(axis) = a;
                }
            }

            traj.positions.row(row) = pos.transpose();
            traj.velocities.row(row) = vel.transpose();
            traj.accelerations.row(row) = acc.transpose();
            traj.time(row) = last_time + (static_cast<double>(t) + 1.0) * Ts_;
        }
    }

    last_input_position_ = command.row(samples - 1).transpose();
    last_input_velocity_ = velocity_matrix.row(samples - 1).transpose();
    last_input_acceleration_ = acceleration_matrix.row(samples - 1).transpose();

    return traj;
}

ShapedSample ShaperInterface::shapeSample(
    const Eigen::VectorXd& command,
    const RobotState& state,
    std::optional<Eigen::VectorXd> command_dot,
    std::optional<Eigen::VectorXd> command_ddot) {
    return processSample(command, state, std::move(command_dot), std::move(command_ddot),
                         /*enable_vibration_control=*/true);
}

ShapedTrajectory ShaperInterface::shapeTrajectory(
    const Eigen::MatrixXd& command,
    const std::vector<RobotState>& states,
    std::optional<Eigen::MatrixXd> command_dot,
    std::optional<Eigen::MatrixXd> command_ddot,
    std::optional<std::vector<double>> time_vector) {
    return processTrajectory(command, states, std::move(command_dot), std::move(command_ddot),
                             std::move(time_vector), /*enable_vibration_control=*/true);
}

/*
Summary:
  Finalize each axis shaper and return synchronized tail samples.
Args:
  None.
Returns:
  `std::vector<ShapedSample>` tail sequence after stream completion.
Side Effects:
  Consumes per-axis shaper tail state.
Raises:
  None.
Preconditions:
  Shapers have been initialized.
*/
std::vector<ShapedSample> ShaperInterface::finalize() {
    if (last_vibration_control_enabled_.has_value() && !(*last_vibration_control_enabled_)) {
        return {};
    }

    std::vector<std::vector<std::tuple<double, double, double>>> tails;
    tails.reserve(static_cast<std::size_t>(num_axes_));

    std::size_t max_tail = 0;
    for (auto& shaper : shapers_) {
        auto tail = shaper.finalize();
        max_tail = std::max<std::size_t>(max_tail, tail.size());
        tails.push_back(std::move(tail));
    }

    std::vector<ShapedSample> sequence;
    sequence.reserve(max_tail);

    for (std::size_t t = 0; t < max_tail; ++t) {
        ShapedSample sample;
        sample.positions = Eigen::VectorXd::Zero(num_axes_);
        sample.velocities = Eigen::VectorXd::Zero(num_axes_);
        sample.accelerations = Eigen::VectorXd::Zero(num_axes_);

        for (int axis = 0; axis < num_axes_; ++axis) {
            const auto& tail = tails[static_cast<std::size_t>(axis)];
            if (t < tail.size()) {
                const auto& [y, v, a] = tail[t];
                sample.positions(axis) = y;
                sample.velocities(axis) = v;
                sample.accelerations(axis) = a;
            } else {
                if (last_input_position_.has_value()) sample.positions(axis) = (*last_input_position_)(axis);
                if (last_input_velocity_.has_value()) sample.velocities(axis) = (*last_input_velocity_)(axis);
                if (last_input_acceleration_.has_value()) sample.accelerations(axis) = (*last_input_acceleration_)(axis);
            }
        }

        sequence.push_back(std::move(sample));
    }

    return sequence;
}

/*
Summary:
  Compute neural-network input features from robot state.
Args:
  state: Robot state containing joint angles and optional TCP position.
Returns:
  `NNInputs` containing inertia diagonal and `(v, r)` geometric features.
Side Effects:
  May log fallback warnings to stderr.
Raises:
  std::runtime_error: If robot dynamics model is unavailable.
Preconditions:
  `state.joint_angles` length matches configured joint count.
*/
ShaperInterface::NNInputs ShaperInterface::computeNNInputs(const RobotState& state) const {
    NNInputs inputs;
    inputs.inertia = computeInertia(state.joint_angles);

    Eigen::Vector3d xyz;
    if (state.tcp_position.has_value()) {
        xyz = *state.tcp_position;
    } else {
        xyz = computeForwardKinematics(state.joint_angles);
    }

    if (!robot_dynamics_) {
        throw std::runtime_error("Robot dynamics model not initialized");
    }

    // Cross-module contract: NN was trained with base angle normalized to zero.
    // Keep this normalization here so feature semantics match training data.
    const auto rotation = kinematics::rotateXyzToZeroBase(*robot_dynamics_, state.joint_angles, xyz);
    if (rotation.used_fallback) {
        std::cerr << "Warning: Robot dynamics model returned no joint rotations; "
                     "falling back to Z-axis base rotation assumption."
                  << std::endl;
    }

    auto [v, r] = cartesianToPolar(rotation.rotated_xyz);
    inputs.v_rad = v;
    inputs.r_m = r;
    return inputs;
}

/*
Summary:
  Compute TCP position in world coordinates from joint angles.
Args:
  joint_angles: Joint configuration vector.
Returns:
  `Eigen::Vector3d` TCP position [m].
Side Effects:
  None.
Raises:
  std::runtime_error: If robot dynamics model is unavailable.
Preconditions:
  `joint_angles` matches model DOF.
*/
Eigen::Vector3d ShaperInterface::computeForwardKinematics(const Eigen::VectorXd& joint_angles) const {
    if (!robot_dynamics_) {
        throw std::runtime_error("Robot dynamics model not initialized");
    }
    return robot_dynamics_->tcpTransform(joint_angles).block<3, 1>(0, 3);
}

/*
Summary:
  Compute joint-space inertia features from the mass-matrix diagonal.
Args:
  joint_angles: Joint configuration vector.
Returns:
  `Eigen::VectorXd` diagonal inertia terms.
Side Effects:
  None.
Raises:
  std::runtime_error: If robot dynamics model is unavailable.
Preconditions:
  `joint_angles` matches model DOF.
*/
Eigen::VectorXd ShaperInterface::computeInertia(const Eigen::VectorXd& joint_angles) const {
    if (!robot_dynamics_) {
        throw std::runtime_error("Robot dynamics model not initialized");
    }
    return robot_dynamics_->massMatrix(joint_angles).diagonal();
}

/*
Summary:
  Convert Cartesian position to model-specific polar features.
Args:
  xyz: Cartesian position [m].
Returns:
  `std::pair<double,double>` containing `(v [rad], r [m])`.
Side Effects:
  None.
Raises:
  None.
Preconditions:
  Geometry parameters (`side_length_`, `base_height_`) are valid.
*/
std::pair<double, double> ShaperInterface::cartesianToPolar(const Eigen::Vector3d& xyz) const {
    return kinematics::cartesianToPolar(xyz, side_length_, base_height_);
}

/*
Summary:
  Compute first finite-difference derivative per joint column.
Args:
  samples: Time-series matrix `[N x joints]`.
Returns:
  `Eigen::MatrixXd` first derivative matrix with same shape.
Side Effects:
  None.
Raises:
  None.
Preconditions:
  Rows correspond to uniform sampling at `Ts_`.
*/
Eigen::MatrixXd ShaperInterface::finiteDifferenceFirst(const Eigen::MatrixXd& samples) const {
    Eigen::Index rows = samples.rows();
    Eigen::Index cols = samples.cols();
    Eigen::MatrixXd diff = Eigen::MatrixXd::Zero(rows, cols);
    if (rows <= 1) return diff;

    // Differentiate each joint column independently at the same sample time.
    for (Eigen::Index col = 0; col < cols; ++col) {
        diff.col(col) = gradient(samples.col(col), Ts_);
    }
    return diff;
}

/*
Summary:
  Compute second finite-difference derivative per joint column.
Args:
  first: First-derivative matrix `[N x joints]`.
Returns:
  `Eigen::MatrixXd` second derivative matrix with same shape.
Side Effects:
  None.
Raises:
  None.
Preconditions:
  Input rows correspond to uniform sampling at `Ts_`.
*/
Eigen::MatrixXd ShaperInterface::finiteDifferenceSecond(const Eigen::MatrixXd& first) const {
    Eigen::Index rows = first.rows();
    Eigen::Index cols = first.cols();
    Eigen::MatrixXd diff = Eigen::MatrixXd::Zero(rows, cols);
    if (rows <= 1) return diff;

    // Apply the same gradient operator to first derivative columns.
    for (Eigen::Index col = 0; col < cols; ++col) {
        diff.col(col) = gradient(first.col(col), Ts_);
    }
    return diff;
}

/*
Summary:
  Ensure Python can import project modules by prepending source root to `sys.path`.
Args:
  python_src_root: Source root path to add.
Returns:
  None.
Side Effects:
  Mutates Python `sys.path`.
Raises:
  pybind11::error_already_set: If Python interactions fail.
Preconditions:
  Python interpreter has been initialized.
*/
void ShaperInterface::ensurePythonPath(const std::string& python_src_root) const {
    py::gil_scoped_acquire gil;
    py::module sys = py::module::import("sys");
    py::list path = sys.attr("path");
    const std::string normalized = python_src_root;
    for (auto item : path) {
        if (std::string(py::str(item)) == normalized) {
            return;
        }
    }
    path.insert(0, normalized);
}

} // namespace covalentcore
