#include "domain/layout_schema.hpp"

namespace lm {

nlohmann::json LayoutSchema::defaultBicycleWheels()
{
  return nlohmann::json::array(
    {nlohmann::json{{"id", "rear"},
          {"role", "drive"},
          {"x_m", 0},
          {"y_m", 0},
          {"joint_name", "rear_wheel_joint"},
          {"motor",
           {{"vendor", "moons"},
            {"model", "m2dc10a"},
            {"gear_ratio", 20},
            {"invert", false}}}},
     nlohmann::json{{"id", "front"},
          {"role", "steer"},
          {"x_m", 1.2},
          {"y_m", 0},
          {"joint_name", "front_steer_joint"},
          {"motor",
           {{"vendor", "moons"},
            {"model", "m2dc10a"},
            {"gear_ratio", 20},
            {"invert", false}}}}});
}

nlohmann::json LayoutSchema::defaultDiffWheels()
{
  return nlohmann::json::array(
    {nlohmann::json{{"id", "left"},
          {"side", "left"},
          {"joint_name", "wheel_left_joint"},
          {"y_m", 0.5},
          {"motor",
           {{"vendor", "moons"},
            {"model", "m2dc10a"},
            {"gear_ratio", 20},
            {"invert", false}}}},
     nlohmann::json{{"id", "right"},
          {"side", "right"},
          {"joint_name", "wheel_right_joint"},
          {"y_m", -0.5},
          {"motor",
           {{"vendor", "moons"},
            {"model", "m2dc10a"},
            {"gear_ratio", 20},
            {"invert", false}}}}});
}

nlohmann::json LayoutSchema::defaultLayoutObject()
{
  return nlohmann::json{
    {"robot",
       {{"x", 400},
        {"y", 300},
        {"yaw_deg", 0},
        {"frame_id", "base_footprint"},
        {"model", "diff"},
        {"diff",
         {{"wheel_separation_m", 1.0},
          {"wheel_radius_m", 0.3},
          {"wheel_separation_multiplier", 1.0},
          {"wheel_radius_multiplier", 1.0},
          {"display", {{"scale_m_per_px", 0.005}, {"b_px", 200}, {"d_px", 120}}},
          {"limits",
           {{"cmd_vel_timeout_s", 0.25},
            {"linear",
             {{"max_velocity", 1.0},
              {"min_velocity", -0.5},
              {"max_acceleration", 0.8},
              {"min_acceleration", -0.4}}},
            {"angular",
             {{"max_velocity", 1.7}, {"max_acceleration", 1.5}}}},
          {"wheels", defaultDiffWheels()}}}},
        {"footprint",
         nlohmann::json::array({nlohmann::json{{"x", 120}, {"y", 80}},
                      nlohmann::json{{"x", 120}, {"y", -80}},
                      nlohmann::json{{"x", -90}, {"y", -80}},
                      nlohmann::json{{"x", -90}, {"y", 80}}})}}},
    {"map", {{"width", 800}, {"height", 600}}},
    {"lidarPositions", nlohmann::json::object()},
    {"lidarPoses", nlohmann::json::object()},
    {"lidarPosesFrame", "robot"},
    {"imuPoses", nlohmann::json::object()},
    {"imuPosesFrame", "robot"}};
}

void LayoutSchema::ensure(nlohmann::json& layout)
{
  if (!layout.is_object())
    layout = nlohmann::json::object();
  if (!layout.contains("robot") || !layout["robot"].is_object())
    layout["robot"] = nlohmann::json::object();
  if (!layout.contains("map") || !layout["map"].is_object())
    layout["map"] = nlohmann::json::object();
  if (!layout.contains("lidarPositions") || !layout["lidarPositions"].is_object())
    layout["lidarPositions"] = nlohmann::json::object();
  if (!layout.contains("lidarPoses") || !layout["lidarPoses"].is_object())
    layout["lidarPoses"] = nlohmann::json::object();
  if (!layout.contains("lidarPosesFrame"))
    layout["lidarPosesFrame"] = "robot";
  if (!layout.contains("imuPoses") || !layout["imuPoses"].is_object())
    layout["imuPoses"] = nlohmann::json::object();
  if (!layout.contains("imuPosesFrame"))
    layout["imuPosesFrame"] = "robot";

  auto& robot = layout["robot"];
  if (!robot.contains("x"))
    robot["x"] = 400;
  if (!robot.contains("y"))
    robot["y"] = 300;
  if (!robot.contains("yaw_deg"))
    robot["yaw_deg"] = 0;
  if (!robot.contains("frame_id"))
    robot["frame_id"] = "base_footprint";
  if (!robot.contains("model"))
    robot["model"] = "diff";
  if (!robot.contains("diff") || !robot["diff"].is_object())
    robot["diff"] = nlohmann::json::object();
  auto& diff = robot["diff"];

  const double default_scale = 0.005;
  if (!diff.contains("display") || !diff["display"].is_object())
    diff["display"] = nlohmann::json::object();
  auto& display = diff["display"];
  if (!display.contains("scale_m_per_px"))
    display["scale_m_per_px"] = default_scale;
  const double scale = display["scale_m_per_px"].get<double>();

  if (!diff.contains("wheel_separation_m"))
  {
    if (diff.contains("b"))
      diff["wheel_separation_m"] = diff["b"].get<double>() * scale;
    else
      diff["wheel_separation_m"] = 1.0;
  }
  if (!diff.contains("wheel_radius_m"))
  {
    if (diff.contains("d"))
      diff["wheel_radius_m"] = diff["d"].get<double>() * scale * 0.5;
    else
      diff["wheel_radius_m"] = 0.3;
  }
  if (!diff.contains("wheel_separation_multiplier"))
    diff["wheel_separation_multiplier"] = 1.0;
  if (!diff.contains("wheel_radius_multiplier"))
    diff["wheel_radius_multiplier"] = 1.0;

  const double b_mult = diff["wheel_separation_multiplier"].get<double>();
  const double r_mult = diff["wheel_radius_multiplier"].get<double>();
  const double sep_m = diff["wheel_separation_m"].get<double>();
  const double rad_m = diff["wheel_radius_m"].get<double>();
  display["b_px"] = sep_m * b_mult / scale;
  display["d_px"] = 2.0 * rad_m * r_mult / scale;
  diff["b"] = display["b_px"];
  diff["d"] = display["d_px"];

  if (!diff.contains("limits") || !diff["limits"].is_object())
    diff["limits"] = nlohmann::json::object();
  auto& limits = diff["limits"];
  if (!limits.contains("cmd_vel_timeout_s"))
    limits["cmd_vel_timeout_s"] = 0.25;
  if (!limits.contains("linear") || !limits["linear"].is_object())
    limits["linear"] = nlohmann::json::object();
  auto& linear = limits["linear"];
  if (!linear.contains("max_velocity"))
    linear["max_velocity"] = 1.0;
  if (!linear.contains("min_velocity"))
    linear["min_velocity"] = -0.5;
  if (!linear.contains("max_acceleration"))
    linear["max_acceleration"] = 0.8;
  if (!linear.contains("min_acceleration"))
    linear["min_acceleration"] = -0.4;
  if (!limits.contains("angular") || !limits["angular"].is_object())
    limits["angular"] = nlohmann::json::object();
  auto& angular = limits["angular"];
  if (!angular.contains("max_velocity"))
    angular["max_velocity"] = 1.7;
  if (!angular.contains("max_acceleration"))
    angular["max_acceleration"] = 1.5;

  const double half_sep = sep_m / 2.0;
  if (!diff.contains("wheels") || !diff["wheels"].is_array() || diff["wheels"].empty())
  {
    diff["wheels"] = nlohmann::json::array(
      {nlohmann::json{{"id", "left"},
            {"side", "left"},
            {"joint_name", "wheel_left_joint"},
            {"y_m", half_sep},
            {"motor",
             nlohmann::json{{"vendor", "moons"},
                  {"model", "m2dc10a"},
                  {"gear_ratio", 20},
                  {"invert", false}}}},
       nlohmann::json{{"id", "right"},
            {"side", "right"},
            {"joint_name", "wheel_right_joint"},
            {"y_m", -half_sep},
            {"motor",
             nlohmann::json{{"vendor", "moons"},
                  {"model", "m2dc10a"},
                  {"gear_ratio", 20},
                  {"invert", false}}}}});
  }
  else
  {
    for (auto& w : diff["wheels"])
    {
      if (!w.is_object())
        continue;
      if (!w.contains("id"))
        w["id"] = "left";
      if (!w.contains("side"))
        w["side"] = (w["id"].get<std::string>() == "right") ? "right" : "left";
      if (!w.contains("joint_name"))
      {
        w["joint_name"] = (w["side"].get<std::string>() == "right") ? "wheel_right_joint"
                                                                    : "wheel_left_joint";
      }
      if (!w.contains("y_m"))
      {
        w["y_m"] = (w["side"].get<std::string>() == "right") ? -half_sep : half_sep;
      }
      if (!w.contains("motor") || !w["motor"].is_object())
        w["motor"] = nlohmann::json::object();
      auto& motor = w["motor"];
      if (!motor.contains("vendor"))
        motor["vendor"] = "custom";
      if (!motor.contains("model"))
        motor["model"] = "custom";
      if (!motor.contains("gear_ratio"))
        motor["gear_ratio"] = 20;
      if (!motor.contains("invert"))
        motor["invert"] = false;
    }
  }

  if (!robot.contains("bicycle") || !robot["bicycle"].is_object())
    robot["bicycle"] = nlohmann::json::object();
  auto& bicycle = robot["bicycle"];
  if (!bicycle.contains("display") || !bicycle["display"].is_object())
    bicycle["display"] = nlohmann::json::object();
  auto& bdisplay = bicycle["display"];
  if (!bdisplay.contains("scale_m_per_px"))
    bdisplay["scale_m_per_px"] = default_scale;
  const double bscale = bdisplay["scale_m_per_px"].get<double>();
  if (!bicycle.contains("wheelbase_m"))
    bicycle["wheelbase_m"] = 1.2;
  if (!bicycle.contains("wheel_radius_m"))
    bicycle["wheel_radius_m"] = 0.15;
  const double L_m = bicycle["wheelbase_m"].get<double>();
  const double b_rad = bicycle["wheel_radius_m"].get<double>();
  bdisplay["L_px"] = L_m / bscale;
  bdisplay["r_px"] = 2.0 * b_rad / bscale;
  if (!bicycle.contains("steer") || !bicycle["steer"].is_object())
    bicycle["steer"] = nlohmann::json::object();
  auto& steer = bicycle["steer"];
  if (!steer.contains("max_angle_deg"))
    steer["max_angle_deg"] = 35;
  if (!steer.contains("preview_deg"))
    steer["preview_deg"] = 15;
  if (!steer.contains("joint_name"))
    steer["joint_name"] = "front_steer_joint";
  if (!bicycle.contains("drive") || !bicycle["drive"].is_object())
    bicycle["drive"] = nlohmann::json::object();
  if (!bicycle["drive"].contains("joint_name"))
    bicycle["drive"]["joint_name"] = "rear_wheel_joint";
  if (!bicycle.contains("limits") || !bicycle["limits"].is_object())
    bicycle["limits"] = nlohmann::json::object();
  auto& blimits = bicycle["limits"];
  if (!blimits.contains("cmd_vel_timeout_s"))
    blimits["cmd_vel_timeout_s"] = 0.25;
  if (!blimits.contains("linear") || !blimits["linear"].is_object())
    blimits["linear"] = nlohmann::json::object();
  auto& blinear = blimits["linear"];
  if (!blinear.contains("max_velocity"))
    blinear["max_velocity"] = 1.0;
  if (!blinear.contains("max_acceleration"))
    blinear["max_acceleration"] = 0.8;
  if (!bicycle.contains("wheels") || !bicycle["wheels"].is_array() || bicycle["wheels"].empty())
    bicycle["wheels"] = defaultBicycleWheels();

  if (!robot.contains("footprint") || !robot["footprint"].is_array())
  {
    robot["footprint"] =
      nlohmann::json::array({nlohmann::json{{"x", 120}, {"y", 80}},
                   nlohmann::json{{"x", 120}, {"y", -80}},
                   nlohmann::json{{"x", -90}, {"y", -80}},
                   nlohmann::json{{"x", -90}, {"y", 80}}});
  }
  if (!robot.contains("footprint_shape"))
    robot["footprint_shape"] = "custom";
  if (!robot.contains("footprint_params") || !robot["footprint_params"].is_object())
  {
    robot["footprint_params"] = nlohmann::json{{"length_m", 1.4},
                                     {"width_m", 1.1},
                                     {"radius_m", 0.55},
                                     {"sides", 6},
                                     {"segments", 32}};
  }
  auto& map = layout["map"];
  if (!map.contains("width"))
    map["width"] = 800;
  if (!map.contains("height"))
    map["height"] = 600;
}
}  // namespace lm