#pragma once

#include "policy_runner_base.hpp"

#include <array>
#include <cmath>
#include <onnxruntime_cxx_api.h>
#include <stdexcept>

/** Stage-2 FireAnt policy adapter (63 observations, 8 actions, 50 Hz). */
class FireAntPolicyRunner : public PolicyRunnerBase {
private:
    static constexpr int kMotorNum = 8;
    static constexpr int kObservationDim = 63;
    static constexpr int kActionDim = 8;
    static constexpr float kWheelRadius = 0.096f;

    // Policy: FL/FR/HL/HR knees, then FL/FR/HL/HR wheels.
    // Robot/DDS: FL knee/wheel, FR knee/wheel, HL knee/wheel, HR knee/wheel.
    const std::array<int, kActionDim> policy_to_robot_{{0, 2, 4, 6, 1, 3, 5, 7}};
    const std::array<float, 4> knee_default_{{0.5f, 0.5f, -0.5f, -0.5f}};

    std::string policy_path_;
    VecXf observation_ = VecXf::Zero(kObservationDim);
    VecXf last_action_ = VecXf::Zero(kActionDim);
    RobotAction robot_action_;

    Ort::Env env_{ORT_LOGGING_LEVEL_WARNING, "FireAntPolicyRunner"};
    Ort::SessionOptions session_options_;
    Ort::Session session_{nullptr};
    Ort::MemoryInfo memory_info_{Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault)};
    const char *input_names_[1] = {"obs"};
    const char *output_names_[1] = {"actions"};
    const std::array<int64_t, 2> input_shape_{{1, kObservationDim}};

    static float Clamp(float value, float low, float high) {
        return std::max(low, std::min(value, high));
    }

    void ValidateModel() {
        if (session_.GetInputCount() != 1 || session_.GetOutputCount() != 1) {
            throw std::runtime_error("FireAnt ONNX must have exactly one input and one output");
        }
        auto input_dims = session_.GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();
        auto output_dims = session_.GetOutputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();
        if (input_dims.size() != 2 || input_dims[1] != kObservationDim) {
            throw std::runtime_error("FireAnt ONNX input must be [batch, 63]");
        }
        if (output_dims.size() != 2 || output_dims[1] != kActionDim) {
            throw std::runtime_error("FireAnt ONNX output must be [batch, 8]");
        }
    }

    VecXf Infer() {
        Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
            memory_info_, observation_.data(), observation_.size(), input_shape_.data(), input_shape_.size());
        auto outputs = session_.Run(
            Ort::RunOptions{nullptr}, input_names_, &input_tensor, 1, output_names_, 1);
        const float *data = outputs[0].GetTensorData<float>();
        return Eigen::Map<const VecXf>(data, kActionDim);
    }

public:
    FireAntPolicyRunner(const std::string &policy_name, const std::string &policy_path)
        : PolicyRunnerBase(policy_name), policy_path_(policy_path) {
        if (access(policy_path_.c_str(), F_OK) != 0) {
            throw std::runtime_error("FireAnt policy not found: " + policy_path_);
        }
        SetDecimation(4);  // 200 Hz state loop / 4 = the 50 Hz training policy rate.
        session_options_.SetIntraOpNumThreads(2);
        session_options_.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
        session_ = Ort::Session(env_, policy_path_.c_str(), session_options_);
        ValidateModel();

        robot_action_.goal_joint_pos = VecXf::Zero(kMotorNum);
        robot_action_.goal_joint_vel = VecXf::Zero(kMotorNum);
        robot_action_.kp = VecXf::Zero(kMotorNum);
        robot_action_.kd = VecXf::Zero(kMotorNum);
        robot_action_.tau_ff = VecXf::Zero(kMotorNum);
        for (int leg = 0; leg < 4; ++leg) {
            robot_action_.kp(2 * leg) = 80.0f;
            robot_action_.kd(2 * leg) = 2.0f;
            robot_action_.kd(2 * leg + 1) = 0.6f;
        }
    }

    void DisplayPolicyInfo() override {
        std::cout << "[FireAnt] ONNX: " << policy_path_ << "\n"
                  << "[FireAnt] observations=63, actions=8, policy_rate=50 Hz" << std::endl;
    }

    void OnEnter() override {
        run_cnt_ = 0;
        cmd_vel_input_.setZero();
        last_action_.setZero();
    }

    RobotAction getRobotAction(const RobotBasicState &ro, const UserCommand &uc) override {
        if (ro.joint_pos.size() != kMotorNum || ro.joint_vel.size() != kMotorNum ||
            ro.joint_tau.size() != kMotorNum) {
            throw std::runtime_error("FireAnt robot state must contain 8 joints");
        }

        // Keyboard commands act as a deadman switch. Stage-2 sampled only positive
        // forward commands, so zero-command inference can otherwise cause creeping.
        if (std::abs(uc.forward_vel_scale) + std::abs(uc.side_vel_scale) +
                std::abs(uc.turnning_vel_scale) < 0.05f) {
            last_action_.setZero();
            robot_action_.goal_joint_pos.setZero();
            robot_action_.goal_joint_vel.setZero();
            for (int leg = 0; leg < 4; ++leg) {
                robot_action_.goal_joint_pos(2 * leg) = knee_default_[leg];
            }
            ++run_cnt_;
            return robot_action_;
        }

        int cursor = 0;
        observation_.segment<3>(cursor) = 0.25f * ro.base_omega;
        cursor += 3;
        observation_.segment<3>(cursor) = ro.base_rot_mat.transpose() * Vec3f(0.0f, 0.0f, -1.0f);
        cursor += 3;
        observation_.segment<3>(cursor) = Vec3f(
            uc.forward_vel_scale, uc.side_vel_scale, uc.turnning_vel_scale);
        cursor += 3;

        // Relative joint position in policy order. Continuous wheel positions are intentionally zero.
        for (int leg = 0; leg < 4; ++leg) {
            observation_(cursor + leg) = ro.joint_pos(policy_to_robot_[leg]) - knee_default_[leg];
            observation_(cursor + 4 + leg) = 0.0f;
        }
        cursor += 8;
        for (int i = 0; i < kActionDim; ++i) {
            observation_(cursor + i) = 0.05f * ro.joint_vel(policy_to_robot_[i]);
        }
        cursor += 8;
        observation_.segment(cursor, kActionDim) = last_action_;
        cursor += 8;

        VecXf contact = ro.contact_force.size() == 9 ? ro.contact_force : VecXf::Zero(9);
        for (int i = 0; i < 4; ++i) observation_(cursor + i) = contact(i) > 2.0f ? 1.0f : 0.0f;
        cursor += 4;
        for (int i = 0; i < 5; ++i) observation_(cursor + i) = contact(4 + i) > 2.0f ? 1.0f : 0.0f;
        cursor += 5;
        for (int i = 0; i < 4; ++i) observation_(cursor + i) = Clamp(contact(i) / 100.0f, 0.0f, 2.0f);
        cursor += 4;
        for (int i = 0; i < 5; ++i) observation_(cursor + i) = Clamp(contact(4 + i) / 100.0f, 0.0f, 2.0f);
        cursor += 5;

        const float forward_speed = std::abs(ro.base_lin_vel(0));
        for (int leg = 0; leg < 4; ++leg) {
            const float wheel_speed = std::abs(ro.joint_vel(2 * leg + 1) * kWheelRadius);
            observation_(cursor + leg) = 0.5f * Clamp(std::abs(wheel_speed - forward_speed), 0.0f, 3.0f);
        }
        cursor += 4;
        for (int i = 0; i < kActionDim; ++i) {
            observation_(cursor + i) = Clamp(ro.joint_tau(policy_to_robot_[i]) / 80.0f, -2.0f, 2.0f);
        }
        cursor += 8;

        if (cursor != kObservationDim || !observation_.allFinite()) {
            throw std::runtime_error("Invalid FireAnt policy observation");
        }

        VecXf action = Infer();
        if (!action.allFinite()) throw std::runtime_error("FireAnt policy produced non-finite actions");
        for (int i = 0; i < kActionDim; ++i) action(i) = Clamp(action(i), -100.0f, 100.0f);
        last_action_ = action;

        robot_action_.goal_joint_pos.setZero();
        robot_action_.goal_joint_vel.setZero();
        for (int leg = 0; leg < 4; ++leg) {
            robot_action_.goal_joint_pos(2 * leg) = knee_default_[leg] + 0.25f * action(leg);
            robot_action_.goal_joint_vel(2 * leg + 1) = 5.0f * action(4 + leg);
        }
        ++run_cnt_;
        return robot_action_;
    }
};