#include "unitree_legged_sdk/safety.h"

#include <algorithm>
#include <cmath>

namespace UNITREE_LEGGED_SDK {
namespace {

constexpr double kTorqueLimit[12] = {
    23.7, 23.7, 35.55, 23.7, 23.7, 35.55,
    23.7, 23.7, 35.55, 23.7, 23.7, 35.55,
};

bool IsPositionStop(float value) {
  return std::fabs(static_cast<double>(value)) >= PosStopF * 0.1;
}

}  // namespace

Safety::Safety(LeggedType type)
    : WattLimit(0), Wcount(0),
      Hip_max(0.78), Hip_min(-0.78),
      Thigh_max(3.50), Thigh_min(-0.60),
      Calf_max(-0.95), Calf_min(-2.70) {
  (void)type;
}

Safety::~Safety() = default;

void Safety::PositionLimit(LowCmd& cmd) {
  for (std::size_t i = 0; i < 12; ++i) {
    auto& motor = cmd.motorCmd[i];
    if (IsPositionStop(motor.q)) continue;
    const std::size_t joint = i % 3;
    const double minimum = joint == 0 ? Hip_min : (joint == 1 ? Thigh_min : Calf_min);
    const double maximum = joint == 0 ? Hip_max : (joint == 1 ? Thigh_max : Calf_max);
    motor.q = static_cast<float>(std::max(minimum, std::min(maximum, static_cast<double>(motor.q))));
  }
}

int Safety::PowerProtect(LowCmd& cmd, LowState& state, int factor) {
  if (factor < 1 || factor > 10) return -1;
  WattLimit = factor;
  int limited = 0;
  for (std::size_t i = 0; i < 12; ++i) {
    const double maximum = kTorqueLimit[i];
    if (std::fabs(static_cast<double>(cmd.motorCmd[i].tau)) > maximum * 5.0 ||
        std::fabs(static_cast<double>(state.motorState[i].tauEst)) > maximum) {
      return -1;
    }
    const double limit = maximum * static_cast<double>(factor) / 10.0;
    const float bounded = static_cast<float>(
        std::max(-limit, std::min(limit, static_cast<double>(cmd.motorCmd[i].tau))));
    if (bounded != cmd.motorCmd[i].tau) {
      cmd.motorCmd[i].tau = bounded;
      ++limited;
    }
  }
  Wcount += limited;
  return limited;
}

int Safety::PositionProtect(LowCmd& cmd, LowState& state, double limit) {
  int protected_count = 0;
  for (std::size_t i = 0; i < 12; ++i) {
    auto& motor = cmd.motorCmd[i];
    if (IsPositionStop(motor.q)) continue;
    if (std::fabs(static_cast<double>(motor.q - state.motorState[i].q)) > limit) {
      motor.Kp = 0.0f;
      motor.Kd = 0.0f;
      ++protected_count;
    }
  }
  return protected_count;
}

}  // namespace UNITREE_LEGGED_SDK