#pragma once

#include <Eigen/Dense>

#include <utility>

#include <memory>
#include <optional>
#include <string>
#include <tuple>
#include <vector>
#include <limits>
#include <iostream>
#include <unordered_map>

#include <pybind11/pybind11.h>

#include "base_shaper.hpp"
#include "robot_dynamics.hpp"

// Forward declarations to avoid pulling heavy headers into dependents.
namespace maploader {
class MapFitter;
class ModelLoader;
} // namespace maploader

// TODO: Remove hard-coded values and read from CSV
inline constexpr int robotNumJoints = 6;
inline constexpr int robotNumShapedAxes = 3;
inline constexpr double robotSideLength = 0.0007;  // [m]
inline constexpr double robotBaseHeight = 0.364;    // [m]
inline constexpr double defaultProbThresh = 0.5;

namespace covalentcore {

namespace kinematics {

namespace rd = reforge::dynamics;

struct BaseRotationResult {
    Eigen::Vector3d rotated_xyz;
    bool used_fallback = false;
};

BaseRotationResult rotateXyzToZeroBase(const rd::RobotDynamics& dynamics,
                                       const Eigen::VectorXd& joint_angles,
                                       const Eigen::Vector3d& xyz);

std::pair<double, double> cartesianToPolar(const Eigen::Vector3d& xyz,
                                           double side_length,
                                           double base_height);

} // namespace kinematics


struct RobotState {
    Eigen::VectorXd joint_angles;                              ///< Current joint configuration [rad]
    std::optional<Eigen::Vector3d> tcp_position = std::nullopt; ///< Optional TCP position [m]
};

struct ShapedSample {
    Eigen::VectorXd positions;
    Eigen::VectorXd velocities;
    Eigen::VectorXd accelerations;
};

struct ShapedTrajectory {
    Eigen::MatrixXd positions;
    Eigen::MatrixXd velocities;
    Eigen::MatrixXd accelerations;
    Eigen::VectorXd time;
};

struct ShaperRow {
    int axis{0};
    double order{1.0};  // >1 means consider second mode
    double wn1{std::numeric_limits<double>::quiet_NaN()};   // [Hz]
    double wn2{std::numeric_limits<double>::quiet_NaN()};   // [Hz] optional
    double zeta1{std::numeric_limits<double>::quiet_NaN()};
    double zeta2{std::numeric_limits<double>::quiet_NaN()};
};

class ShaperInterface {
public:
    ShaperInterface(double sample_time,
                    std::string model_directory,
                    const std::string& python_src_root,
                    const std::string& urdf_filepath,
                    double side_length = std::numeric_limits<double>::quiet_NaN(),
                    double base_height = std::numeric_limits<double>::quiet_NaN(),
                    int num_axes = 0,
                    int num_joints = 0,
                    double prob_thresh = defaultProbThresh);

    /**
     * @brief Process one command sample with optional vibration shaping.
     *
     * When `enable_vibration_control` is false, the command is passed through
     * unchanged and derivatives are returned from provided values or finite
     * differences against the previous sample. This keeps one stable API for
     * both shaped and unshaped execution paths.
     */
    ShapedSample processSample(
        const Eigen::VectorXd& command,
        const RobotState& state,
        std::optional<Eigen::VectorXd> command_dot = std::nullopt,
        std::optional<Eigen::VectorXd> command_ddot = std::nullopt,
        bool enable_vibration_control = true);

    ShapedSample shapeSample(
        const Eigen::VectorXd& command,
        const RobotState& state,
        std::optional<Eigen::VectorXd> command_dot = std::nullopt,
        std::optional<Eigen::VectorXd> command_ddot = std::nullopt);

    /**
     * @brief Process a full trajectory with optional vibration shaping.
     *
     * When `enable_vibration_control` is false, the input position trajectory
     * is returned unchanged and velocity/acceleration are computed (or reused
     * when provided), with no shaper tail extension.
     */
    ShapedTrajectory processTrajectory(
        const Eigen::MatrixXd& command,
        const std::vector<RobotState>& states,
        std::optional<Eigen::MatrixXd> command_dot = std::nullopt,
        std::optional<Eigen::MatrixXd> command_ddot = std::nullopt,
        std::optional<std::vector<double>> time_vector = std::nullopt,
        bool enable_vibration_control = true);

    ShapedTrajectory shapeTrajectory(
        const Eigen::MatrixXd& command,
        const std::vector<RobotState>& states,
        std::optional<Eigen::MatrixXd> command_dot = std::nullopt,
        std::optional<Eigen::MatrixXd> command_ddot = std::nullopt,
        std::optional<std::vector<double>> time_vector = std::nullopt);

    /**
     * @brief Flush the per-axis shaper delay using the most recent commands.
     *
     * Runs BaseShaper::finalize() for each shaped axis and stitches the
     * resulting tails into a sequence of shaped samples. If a particular axis
     * reports a shorter tail than the maximum delay, the last known shaped
     * values are held constant for the remainder of the sequence.
     */
    std::vector<ShapedSample> finalize();

    void reset();
    
    double sampleTime() const { return Ts_; }
    int axes() const { return num_axes_; }
    int joints() const { return num_joints_; }

private:
    struct NNInputs {
        double v_rad = 0.0;
        double r_m = 0.0;
        Eigen::VectorXd inertia;
    };

    struct FiniteDifferenceCache {
        Eigen::MatrixXd positions;
        Eigen::MatrixXd velocities;
        Eigen::MatrixXd accelerations;
        bool velocities_valid = false;
        bool accelerations_valid = false;
    };

    NNInputs computeNNInputs(const RobotState& state) const;
    Eigen::Vector3d computeForwardKinematics(const Eigen::VectorXd& joint_angles) const;
    Eigen::VectorXd computeInertia(const Eigen::VectorXd& joint_angles) const;
    std::pair<double, double> cartesianToPolar(const Eigen::Vector3d& xyz) const;
    Eigen::MatrixXd finiteDifferenceFirst(const Eigen::MatrixXd& samples) const;
    Eigen::MatrixXd finiteDifferenceSecond(const Eigen::MatrixXd& first) const;
    void ensurePythonPath(const std::string& python_src_root) const;
    bool handleVibrationModeTransition(bool enable_vibration_control);

    double Ts_;
    int num_axes_;
    int num_joints_;
    double side_length_;
    double base_height_;
    double prob_thresh_;
    std::unordered_map<int, ShaperRow> axis_params_;

    std::string model_directory_;
    std::shared_ptr<maploader::MapFitter> map_fitter_;
    std::unique_ptr<reforge::dynamics::RobotDynamics> robot_dynamics_;
    std::vector<control::BaseShaper> shapers_;
    std::optional<FiniteDifferenceCache> fd_cache_;
    std::optional<Eigen::VectorXd> last_input_position_;
    std::optional<Eigen::VectorXd> last_input_velocity_;
    std::optional<Eigen::VectorXd> last_input_acceleration_;
    std::optional<bool> last_vibration_control_enabled_;
};

} // namespace covalentcore
