/**
 * @file invicta_policy_runner.hpp
 * @brief ONNX 策略推理。关节数量、顺序、动作尺度、默认姿态、增益全部来自 invicta.toml
 *
 * 观测（M20 同款配方，维度 = 9 + 3N）：
 *   [角速度(3)·omega_scale, 投影重力(3), 速度指令(3), 关节角(N), 关节角速度(N)·dof_vel_scale, 上一步动作(N)]
 * N = 关节数量（本构型 8）。轮关节位置是否进观测由 zero_obs_wheels 控制。
 *
 * 动作 → 指令：非轮关节给位置目标（action·scale + default），轮关节给速度目标。
 * 配置里的 action_scale / default_joint_pos / kp / kd 按「策略顺序」（policy.joint_order）给出，
 * 这里自动置换到机器人顺序。
 */
#pragma once

#include <onnxruntime_c_api.h>
#include <onnxruntime_cxx_api.h>

#include <algorithm>
#include <cmath>
#include <string>
#include <unordered_map>
#include <vector>

#include "config/deploy_config.h"
#include "policy_runner_base.hpp"

class InvictaPolicyRunner : public PolicyRunnerBase {
public:
    InvictaPolicyRunner(const std::string& policy_name, const std::string& policy_path)
        : PolicyRunnerBase(policy_name),
          cfg_(invicta::Config()),
          policy_path_(policy_path.empty() ? cfg_.policy_path : policy_path),
          env_(ORT_LOGGING_LEVEL_WARNING, "InvictaPolicyRunner"),
          session_options_{},
          session_{nullptr},
          memory_info(Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault)) {
        action_dim_ = cfg_.action_dim;
        obs_dim_ = cfg_.obs_dim;
        if (action_dim_ <= 0 || obs_dim_ <= 0) {
            throw std::runtime_error("policy dimension is 0: config not loaded or [policy] section incomplete");
        }

        robot2policy_idx_ = Permutation(cfg_.joint_names, cfg_.policy_joint_order);
        policy2robot_idx_ = Permutation(cfg_.policy_joint_order, cfg_.joint_names);

        ToRobotOrder(cfg_.action_scale, action_scale_robot_);
        ToRobotOrder(cfg_.default_joint_pos, default_robot_);
        ToRobotOrder(cfg_.policy_kp, kp_robot_);
        ToRobotOrder(cfg_.policy_kd, kd_robot_);
        default_policy_.resize(action_dim_);
        for (int i = 0; i < action_dim_; ++i) default_policy_(i) = float(cfg_.default_joint_pos[size_t(i)]);
        wheel_robot_idx_ = cfg_.WheelIndices();

        SetDecimation(cfg_.policy_decimation);
        session_options_.SetIntraOpNumThreads(4);
        session_options_.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);

        if (access(policy_path_.c_str(), F_OK) != 0) {
            std::cerr << "Model file not found: " << policy_path_ << std::endl;
            throw std::runtime_error("Model file missing");
        }
        session_ = Ort::Session(env_, policy_path_.c_str(), session_options_);

        current_observation_.setZero(obs_dim_);
        last_action_eigen_.setZero(action_dim_);
        tmp_action_eigen_.setZero(action_dim_);
        current_action_eigen_.setZero(action_dim_);
        joint_pos_rl_.setZero(action_dim_);
        joint_vel_rl_.setZero(action_dim_);

        robot_action_.kp = kp_robot_;
        robot_action_.kd = kd_robot_;
        robot_action_.tau_ff = VecXf::Zero(action_dim_);
        robot_action_.goal_joint_pos = VecXf::Zero(action_dim_);
        robot_action_.goal_joint_vel = VecXf::Zero(action_dim_);

        input_observation_shape_ = {1, obs_dim_};
    }

    ~InvictaPolicyRunner() override = default;

    void DisplayPolicyInfo() {
        std::cout << "[policy] " << policy_path_ << " obs=" << obs_dim_ << " action=" << action_dim_
                  << " decimation=" << decimation_ << std::endl;
    }

    void OnEnter() {
        run_cnt_ = 0;
        cmd_vel_input_.setZero();
        last_action_eigen_.setZero(action_dim_);
        tmp_action_eigen_.setZero(action_dim_);
    }

    RobotAction getRobotAction(const RobotBasicState& ro, const UserCommand& uc) {
        const Vec3f base_omega = ro.base_omega * float(cfg_.omega_scale);
        const Vec3f projected_gravity = ro.base_rot_mat.inverse() * gravity_direction_;
        const Vec3f cmd(uc.forward_vel_scale, uc.side_vel_scale, uc.turnning_vel_scale);

        for (int i = 0; i < action_dim_; ++i) {
            const int r = robot2policy_idx_[size_t(i)];  // 策略第 i 位对应机器人第 r 个关节
            joint_pos_rl_(i) = ro.joint_pos(r);
            joint_vel_rl_(i) = ro.joint_vel(r) * float(cfg_.dof_vel_scale);
        }
        if (cfg_.zero_obs_wheels) {
            for (int w : wheel_robot_idx_) joint_pos_rl_(robot2policy_idx_[size_t(w)]) = 0.f;
        }
        joint_pos_rl_ -= default_policy_;

        current_observation_ << base_omega, projected_gravity, cmd, joint_pos_rl_, joint_vel_rl_, last_action_eigen_;
        current_action_eigen_ = OnnxInfer(current_observation_);
        last_action_eigen_ = current_action_eigen_;

        for (int i = 0; i < action_dim_; ++i) {
            tmp_action_eigen_(i) = current_action_eigen_(policy2robot_idx_[size_t(i)]) * action_scale_robot_(i);
        }
        tmp_action_eigen_ += default_robot_;

        for (int i = 0; i < action_dim_; ++i) {
            if (cfg_.IsWheel(i)) {
                robot_action_.goal_joint_pos(i) = 0.f;
                robot_action_.goal_joint_vel(i) = tmp_action_eigen_(i);
            } else {
                robot_action_.goal_joint_pos(i) = tmp_action_eigen_(i);
                robot_action_.goal_joint_vel(i) = 0.f;
            }
        }

        ++run_cnt_;
        return robot_action_;
    }

private:
    static std::vector<int> Permutation(const std::vector<std::string>& from, const std::vector<std::string>& to) {
        std::unordered_map<std::string, int> idx;
        for (size_t i = 0; i < from.size(); ++i) idx[from[i]] = int(i);
        std::vector<int> perm;
        for (const auto& n : to) {
            auto it = idx.find(n);
            perm.push_back(it == idx.end() ? 0 : it->second);
        }
        return perm;
    }

    /// 把「策略顺序」的数组搬到「机器人顺序」
    void ToRobotOrder(const std::vector<double>& policy_order_values, VecXf& out) {
        out.resize(action_dim_);
        for (int i = 0; i < action_dim_; ++i) {
            out(i) = float(policy_order_values[size_t(robot2policy_idx_[size_t(i)])]);
        }
    }

    VecXf OnnxInfer(VecXf observation) {
        Ort::Value input_tensor =
            Ort::Value::CreateTensor<float>(memory_info, observation.data(), observation.size(),
                                            input_observation_shape_.data(), input_observation_shape_.size());
        std::vector<Ort::Value> inputs;
        inputs.emplace_back(std::move(input_tensor));
        auto outputs = session_.Run(Ort::RunOptions{nullptr}, input_names_, inputs.data(), 1, output_names_, 1);
        float* data = outputs[0].GetTensorMutableData<float>();
        Eigen::Map<Eigen::VectorXf> mapped(data, action_dim_);
        return VecXf(mapped);
    }

    const invicta::DeployConfig& cfg_;
    const std::string policy_path_;
    int action_dim_ = 0, obs_dim_ = 0;
    const Vec3f gravity_direction_ = Vec3f(0., 0., -1.);

    std::vector<int> robot2policy_idx_, policy2robot_idx_, wheel_robot_idx_;
    VecXf action_scale_robot_, default_robot_, kp_robot_, kd_robot_, default_policy_;

    VecXf joint_pos_rl_, joint_vel_rl_;
    VecXf current_observation_, last_action_eigen_, tmp_action_eigen_, current_action_eigen_;
    RobotAction robot_action_;

    Ort::SessionOptions session_options_;
    Ort::Session session_;
    Ort::Env env_;
    const char* input_names_[1] = {"obs"};  // 必须与模型导出时一致
    const char* output_names_[1] = {"actions"};
    Ort::MemoryInfo memory_info;
    std::vector<int64_t> input_observation_shape_;
};