#include "unitree_legged_sdk/unitree_legged_sdk.h"

#include <chrono>
#include <cmath>
#include <iostream>
#include <thread>

using namespace UNITREE_LEGGED_SDK;

int main() {
  constexpr float dt = 0.002f;
  UDP udp(LOWLEVEL, 8090, UDP_SERVER_IP_BASIC, UDP_SERVER_PORT);
  Safety safety(LeggedType::Go1);
  LowCmd cmd{};
  LowState state{};
  udp.InitCmdData(cmd);

  bool have_state = false;
  for (int attempt = 0; attempt < 1000 && !have_state; ++attempt) {
    udp.SetSend(cmd);
    udp.Send();
    if (udp.Recv() > 0) {
      udp.GetRecv(state);
      have_state = state.head[0] == 0xfe && state.head[1] == 0xef;
    }
    std::this_thread::sleep_for(std::chrono::duration<float>(dt));
  }
  if (!have_state) {
    std::cerr << "No LowState received; active position command was not enabled.\n";
    return 1;
  }

  const float initial = state.motorState[FR_1].q;
  for (int step = 0; step < 2500; ++step) {
    udp.Recv();
    udp.GetRecv(state);
    const float phase = static_cast<float>(step) * dt;
    cmd.motorCmd[FR_1].q = initial + 0.15f * std::sin(phase * 2.0f);
    cmd.motorCmd[FR_1].dq = 0.0f;
    cmd.motorCmd[FR_1].Kp = 5.0f;
    cmd.motorCmd[FR_1].Kd = 1.0f;
    cmd.motorCmd[FR_1].tau = 0.0f;
    safety.PositionLimit(cmd);
    if (safety.PowerProtect(cmd, state, 1) < 0) {
      udp.InitCmdData(cmd);
      udp.SetSend(cmd);
      udp.Send();
      return 2;
    }
    safety.PositionProtect(cmd, state, 0.5);
    udp.SetSend(cmd);
    udp.Send();
    std::this_thread::sleep_for(std::chrono::duration<float>(dt));
  }

  udp.InitCmdData(cmd);
  udp.SetSend(cmd);
  udp.Send();
  return 0;
}