#include "origincar_base/wall_fit_core.hpp"

#include <nlohmann/json.hpp>

#include <algorithm>
#include <array>
#include <chrono>
#include <cmath>
#include <cstdio>
#include <fstream>
#include <iomanip>
#include <limits>
#include <sstream>
#include <stdexcept>
#include <unistd.h>

namespace origincar_wall
{
namespace
{
constexpr double kPi = 3.14159265358979323846;
constexpr double kYawSearchRange = 8.0 * kPi / 180.0;
constexpr double kYawSearchStep = 1.0 * kPi / 180.0;
constexpr double kWallGroupMaxDistance = 0.45;
constexpr int kMinScanPoints = 8;
constexpr int kMinWallPoints = 3;

double mean(const std::vector<double> &values)
{
  if (values.empty())
  {
    throw std::runtime_error("mean requires values");
  }
  double total = 0.0;
  for (const auto value : values)
  {
    total += value;
  }
  return total / static_cast<double>(values.size());
}

double median(std::vector<double> values)
{
  if (values.empty())
  {
    throw std::runtime_error("median requires values");
  }
  std::sort(values.begin(), values.end());
  const auto middle = values.size() / 2;
  if (values.size() % 2 == 1)
  {
    return values[middle];
  }
  return (values[middle - 1] + values[middle]) / 2.0;
}

double percentile(std::vector<double> values, double fraction)
{
  if (values.empty())
  {
    throw std::runtime_error("percentile requires values");
  }
  std::sort(values.begin(), values.end());
  const auto raw_index = static_cast<int>(std::round((values.size() - 1) * fraction));
  const auto index = std::max(0, std::min(raw_index, static_cast<int>(values.size() - 1)));
  return values[static_cast<std::size_t>(index)];
}

WallLine wallByName(const WallFitConfig &config, const std::string &name)
{
  for (const auto &wall : config.walls)
  {
    if (wall.name == name)
    {
      return wall;
    }
  }
  throw std::runtime_error("missing wall " + name);
}

double verticalWallX(const WallLine &wall)
{
  if (std::fabs(wall.a) < 1e-9 || std::fabs(wall.b) > 1e-6)
  {
    throw std::runtime_error("wall " + wall.name + " must be vertical");
  }
  return -wall.c / wall.a;
}

double horizontalWallY(const WallLine &wall)
{
  if (std::fabs(wall.b) < 1e-9 || std::fabs(wall.a) > 1e-6)
  {
    throw std::runtime_error("wall " + wall.name + " must be horizontal");
  }
  return -wall.c / wall.b;
}

std::pair<std::string, double> nearestWallResidual(const WallPoint &point, const WallBounds &bounds)
{
  std::array<std::pair<std::string, double>, 4> candidates{{
      {"left", point.x - bounds.x_min},
      {"right", point.x - bounds.x_max},
      {"bottom", point.y - bounds.y_min},
      {"top", point.y - bounds.y_max},
  }};
  return *std::min_element(
      candidates.begin(), candidates.end(),
      [](const auto &left, const auto &right) { return std::fabs(left.second) < std::fabs(right.second); });
}

std::map<std::string, std::vector<WallPoint>> groupScanPointsByWall(
    const std::vector<WallPoint> &scan_map, const WallBounds &bounds)
{
  std::map<std::string, std::vector<WallPoint>> groups{
      {"left", {}}, {"right", {}}, {"bottom", {}}, {"top", {}}};
  for (const auto &point : scan_map)
  {
    const auto nearest = nearestWallResidual(point, bounds);
    if (std::fabs(nearest.second) <= kWallGroupMaxDistance)
    {
      groups[nearest.first].push_back(point);
    }
  }
  return groups;
}

std::vector<double> yawCandidates(double center_theta)
{
  std::vector<double> candidates;
  const auto steps = static_cast<int>(std::round(kYawSearchRange / kYawSearchStep));
  for (int index = -steps; index <= steps; ++index)
  {
    candidates.push_back(normalizeAngle(center_theta + index * kYawSearchStep));
  }
  return candidates;
}

LocalizationResult solveWallOffsetsForPose(const std::vector<WallPoint> &scan_points,
                                           const WallBounds &bounds,
                                           const Pose2D &pose)
{
  const auto scan_map = transformScanPointsToMap(scan_points, pose);
  const auto groups = groupScanPointsByWall(scan_map, bounds);
  std::vector<double> x_corrections;
  std::vector<double> y_corrections;
  std::vector<MatchedWall> matched_walls;

  for (const auto &entry : groups)
  {
    const auto &name = entry.first;
    const auto &group_points = entry.second;
    if (static_cast<int>(group_points.size()) < kMinWallPoints)
    {
      continue;
    }
    const bool x_axis = (name == "left" || name == "right");
    const double target = name == "left"    ? bounds.x_min
                          : name == "right" ? bounds.x_max
                          : name == "bottom" ? bounds.y_min
                                             : bounds.y_max;
    std::vector<double> observed_values;
    observed_values.reserve(group_points.size());
    for (const auto &point : group_points)
    {
      observed_values.push_back(x_axis ? point.x : point.y);
    }
    const auto observed = median(observed_values);
    const auto correction = target - observed;
    std::vector<double> errors;
    errors.reserve(group_points.size());
    for (const auto &value : observed_values)
    {
      errors.push_back(std::fabs((value + correction) - target));
    }
    matched_walls.push_back(
        {name, static_cast<int>(group_points.size()), observed, target, correction, mean(errors)});
    if (x_axis)
    {
      x_corrections.push_back(correction);
    }
    else
    {
      y_corrections.push_back(correction);
    }
  }

  LocalizationQuality quality;
  quality.wall_count = static_cast<int>(matched_walls.size());
  quality.has_mean_error = !matched_walls.empty();
  quality.mean_error = 0.0;
  if (quality.has_mean_error)
  {
    std::vector<double> wall_errors;
    for (const auto &wall : matched_walls)
    {
      wall_errors.push_back(wall.mean_error);
    }
    quality.mean_error = mean(wall_errors);
  }

  const bool valid = matched_walls.size() >= 2 && !x_corrections.empty() && !y_corrections.empty();
  if (!valid)
  {
    return {false, pose, scan_map, matched_walls, quality, "fewer than two usable perpendicular wall constraints"};
  }

  Pose2D corrected_pose{
      pose.x + mean(x_corrections),
      pose.y + mean(y_corrections),
      normalizeAngle(pose.theta),
  };
  return {
      true,
      corrected_pose,
      transformScanPointsToMap(scan_points, corrected_pose),
      matched_walls,
      quality,
      "localized from odom prediction and calibrated walls"};
}

nlohmann::json pointToJson(const WallPoint &point)
{
  return nlohmann::json{{"x", point.x}, {"y", point.y}};
}

WallPoint pointFromJson(const nlohmann::json &json)
{
  return {json.at("x").get<double>(), json.at("y").get<double>()};
}

nlohmann::json wallToJson(const WallLine &wall)
{
  return nlohmann::json{{"name", wall.name}, {"a", wall.a}, {"b", wall.b}, {"c", wall.c}, {"color", wall.color}};
}

WallLine wallFromJson(const nlohmann::json &json)
{
  return {
      json.at("name").get<std::string>(),
      json.at("a").get<double>(),
      json.at("b").get<double>(),
      json.at("c").get<double>(),
      json.value("color", std::string("#64748b"))};
}

nlohmann::json configToJsonObject(const WallFitConfig &config)
{
  nlohmann::json walls = nlohmann::json::array();
  for (const auto &wall : config.walls)
  {
    walls.push_back(wallToJson(wall));
  }
  nlohmann::json refs = nlohmann::json::object();
  for (const auto &entry : config.reference_points)
  {
    refs[entry.first] = {
        {"label", entry.second.label},
        {"image", pointToJson(entry.second.image)},
        {"map", pointToJson(entry.second.map)}};
  }
  nlohmann::json frame = {
      {"origin", config.coordinate_frame.origin},
      {"origin_note", config.coordinate_frame.origin_note},
      {"width_m", config.coordinate_frame.width_m},
      {"height_m", config.coordinate_frame.height_m},
      {"image_width_px", config.coordinate_frame.image_width_px},
      {"image_height_px", config.coordinate_frame.image_height_px},
      {"pixels_per_meter", config.coordinate_frame.pixels_per_meter},
      {"p_image_px", pointToJson(config.coordinate_frame.p_image_px)},
      {"odom_zero_reference", config.coordinate_frame.odom_zero_reference},
      {"odom_zero_map", pointToJson(config.coordinate_frame.odom_zero_map)}};

  nlohmann::json result = {{"coordinate_frame", frame}, {"walls", walls}, {"reference_points", refs}};
  if (!config.generated_at.empty())
  {
    result["generated_at"] = config.generated_at;
  }
  return result;
}

WallFitConfig configFromJsonObject(const nlohmann::json &json)
{
  auto config = defaultWallFitConfig();
  const auto &frame = json.at("coordinate_frame");
  config.coordinate_frame.origin = frame.at("origin").get<std::string>();
  config.coordinate_frame.origin_note = frame.value("origin_note", config.coordinate_frame.origin_note);
  config.coordinate_frame.width_m = frame.value("width_m", config.coordinate_frame.width_m);
  config.coordinate_frame.height_m = frame.value("height_m", config.coordinate_frame.height_m);
  config.coordinate_frame.image_width_px = frame.value("image_width_px", config.coordinate_frame.image_width_px);
  config.coordinate_frame.image_height_px = frame.value("image_height_px", config.coordinate_frame.image_height_px);
  config.coordinate_frame.pixels_per_meter = frame.value("pixels_per_meter", config.coordinate_frame.pixels_per_meter);
  if (frame.contains("p_image_px"))
  {
    config.coordinate_frame.p_image_px = pointFromJson(frame.at("p_image_px"));
  }
  config.coordinate_frame.odom_zero_reference =
      frame.value("odom_zero_reference", config.coordinate_frame.odom_zero_reference);
  if (frame.contains("odom_zero_map"))
  {
    config.coordinate_frame.odom_zero_map = pointFromJson(frame.at("odom_zero_map"));
  }

  config.walls.clear();
  for (const auto &wall_json : json.at("walls"))
  {
    config.walls.push_back(wallFromJson(wall_json));
  }

  if (json.contains("reference_points"))
  {
    config.reference_points.clear();
    for (auto it = json.at("reference_points").begin(); it != json.at("reference_points").end(); ++it)
    {
      config.reference_points[it.key()] = {
          it.value().value("label", it.key()),
          pointFromJson(it.value().at("image")),
          pointFromJson(it.value().at("map"))};
    }
  }
  config.generated_at = json.value("generated_at", std::string());
  if (config.coordinate_frame.origin != "map_lower_left")
  {
    throw std::runtime_error("wall_fit coordinate_frame.origin must be map_lower_left");
  }
  wallBounds(config);
  return config;
}

std::string utcNow()
{
  const auto now = std::chrono::system_clock::now();
  const auto time = std::chrono::system_clock::to_time_t(now);
  std::tm tm{};
  gmtime_r(&time, &tm);
  std::ostringstream stream;
  stream << std::put_time(&tm, "%Y-%m-%dT%H:%M:%SZ");
  return stream.str();
}
}  // namespace

WallFitConfig defaultWallFitConfig()
{
  WallFitConfig config;
  config.coordinate_frame = {
      "map_lower_left",
      "map (0,0) is the lower-left map corner; odom zero is placed at the P point",
      5.0,
      5.0,
      500,
      500,
      100,
      {56.0, 482.0},
      "P",
      {0.56, 0.18},
  };
  config.walls = {
      {"top", 0.0, -1.0, 5.0, "#d62728"},
      {"bottom", 0.0, 1.0, 0.0, "#1f77b4"},
      {"left", 1.0, 0.0, 0.0, "#2ca02c"},
      {"right", -1.0, 0.0, 5.0, "#ff7f0e"},
  };
  config.reference_points = {
      {"P", {"P", {56.0, 482.0}, {0.56, 0.18}}},
      {"task_release", {"task_release_left_midpoint", {450.0, 319.0}, {4.5, 1.81}}},
      {"passage", {"passage", {250.0, 240.0}, {2.5, 2.6}}},
      {"clinic", {"clinic_room", {280.0, 123.0}, {2.8, 3.77}}},
  };
  return config;
}

WallBounds wallBounds(const WallFitConfig &config)
{
  const WallBounds bounds{
      verticalWallX(wallByName(config, "left")),
      verticalWallX(wallByName(config, "right")),
      horizontalWallY(wallByName(config, "bottom")),
      horizontalWallY(wallByName(config, "top")),
  };
  if (bounds.x_min >= bounds.x_max || bounds.y_min >= bounds.y_max)
  {
    throw std::runtime_error("wall bounds are invalid");
  }
  return bounds;
}

double normalizeAngle(double angle)
{
  return std::fmod(angle + kPi, 2.0 * kPi) < 0.0
             ? std::fmod(angle + kPi, 2.0 * kPi) + kPi
             : std::fmod(angle + kPi, 2.0 * kPi) - kPi;
}

Pose2D predictPoseFromOdomDelta(const Pose2D &previous_wall_pose,
                                const Pose2D &previous_odom_pose,
                                const Pose2D &current_odom_pose)
{
  return {
      previous_wall_pose.x + current_odom_pose.x - previous_odom_pose.x,
      previous_wall_pose.y + current_odom_pose.y - previous_odom_pose.y,
      normalizeAngle(previous_wall_pose.theta + normalizeAngle(current_odom_pose.theta - previous_odom_pose.theta)),
  };
}

std::vector<WallPoint> transformScanPointsToMap(const std::vector<WallPoint> &scan_points,
                                                const Pose2D &pose)
{
  const auto cos_t = std::cos(pose.theta);
  const auto sin_t = std::sin(pose.theta);
  std::vector<WallPoint> result;
  result.reserve(scan_points.size());
  for (const auto &point : scan_points)
  {
    result.push_back({pose.x + point.x * cos_t - point.y * sin_t,
                      pose.y + point.x * sin_t + point.y * cos_t});
  }
  return result;
}

std::vector<WallPoint> transformScanPointsToPoseFrame(const std::vector<WallPoint> &scan_points,
                                                      const Pose2D &laser_pose)
{
  return transformScanPointsToMap(scan_points, laser_pose);
}

std::vector<WallPoint> mapPointsToScan(const std::vector<WallPoint> &map_points,
                                       const Pose2D &pose)
{
  const auto cos_t = std::cos(pose.theta);
  const auto sin_t = std::sin(pose.theta);
  std::vector<WallPoint> result;
  result.reserve(map_points.size());
  for (const auto &point : map_points)
  {
    const auto dx = point.x - pose.x;
    const auto dy = point.y - pose.y;
    result.push_back({dx * cos_t + dy * sin_t, -dx * sin_t + dy * cos_t});
  }
  return result;
}

LocalizationResult localizeFromScan(const std::vector<WallPoint> &scan_points,
                                    const WallFitConfig &config,
                                    const Pose2D &predicted_pose)
{
  if (static_cast<int>(scan_points.size()) < kMinScanPoints)
  {
    return {
        false,
        predicted_pose,
        transformScanPointsToMap(scan_points, predicted_pose),
        {},
        {0, 0.0, false},
        "not enough scan points for wall localization"};
  }
  const auto bounds = wallBounds(config);
  const auto center_theta = predicted_pose.theta;
  std::vector<LocalizationResult> valid_results;
  for (const auto theta : yawCandidates(center_theta))
  {
    auto solution = solveWallOffsetsForPose(scan_points, bounds, {predicted_pose.x, predicted_pose.y, theta});
    if (solution.ok)
    {
      valid_results.push_back(solution);
    }
  }
  if (valid_results.empty())
  {
    return solveWallOffsetsForPose(scan_points, bounds, predicted_pose);
  }
  return *std::min_element(
      valid_results.begin(), valid_results.end(),
      [center_theta](const auto &left, const auto &right) {
        if (std::fabs(left.quality.mean_error - right.quality.mean_error) > 1e-9)
        {
          return left.quality.mean_error < right.quality.mean_error;
        }
        if (left.quality.wall_count != right.quality.wall_count)
        {
          return left.quality.wall_count > right.quality.wall_count;
        }
        return std::fabs(normalizeAngle(left.pose.theta - center_theta)) <
               std::fabs(normalizeAngle(right.pose.theta - center_theta));
      });
}

WallFitConfig fitWallConfigFromCaptures(const std::vector<WallFitCapture> &captures,
                                        const WallFitConfig &base_config)
{
  if (captures.size() < 4)
  {
    throw std::runtime_error("wall calibration requires four captures");
  }
  std::vector<WallPoint> all_points;
  for (const auto &capture : captures)
  {
    const auto scan_map = transformScanPointsToMap(capture.scan_points, capture.odom_pose);
    all_points.insert(all_points.end(), scan_map.begin(), scan_map.end());
  }
  if (all_points.size() < 16)
  {
    throw std::runtime_error("wall calibration requires enough scan points");
  }
  std::vector<double> xs;
  std::vector<double> ys;
  xs.reserve(all_points.size());
  ys.reserve(all_points.size());
  for (const auto &point : all_points)
  {
    if (std::isfinite(point.x))
    {
      xs.push_back(point.x);
    }
    if (std::isfinite(point.y))
    {
      ys.push_back(point.y);
    }
  }
  const auto left_cut = percentile(xs, 0.08);
  const auto right_cut = percentile(xs, 0.92);
  const auto bottom_cut = percentile(ys, 0.08);
  const auto top_cut = percentile(ys, 0.92);
  std::vector<double> left_values;
  std::vector<double> right_values;
  std::vector<double> bottom_values;
  std::vector<double> top_values;
  for (const auto value : xs)
  {
    if (value <= left_cut)
    {
      left_values.push_back(value);
    }
    if (value >= right_cut)
    {
      right_values.push_back(value);
    }
  }
  for (const auto value : ys)
  {
    if (value <= bottom_cut)
    {
      bottom_values.push_back(value);
    }
    if (value >= top_cut)
    {
      top_values.push_back(value);
    }
  }

  const auto x_min = median(left_values);
  const auto x_max = median(right_values);
  const auto y_min = median(bottom_values);
  const auto y_max = median(top_values);

  auto config = base_config;
  config.coordinate_frame = defaultWallFitConfig().coordinate_frame;
  config.reference_points = defaultWallFitConfig().reference_points;
  config.walls = {
      {"top", 0.0, -1.0, y_max, "#d62728"},
      {"bottom", 0.0, 1.0, -y_min, "#1f77b4"},
      {"left", 1.0, 0.0, -x_min, "#2ca02c"},
      {"right", -1.0, 0.0, x_max, "#ff7f0e"},
  };
  config.generated_at = utcNow();
  wallBounds(config);
  return config;
}

WallFitConfig loadWallFitConfig(const std::string &path)
{
  std::ifstream file(path);
  if (!file.good())
  {
    return defaultWallFitConfig();
  }
  nlohmann::json json;
  file >> json;
  return configFromJsonObject(json);
}

void saveWallFitConfig(const std::string &path, const WallFitConfig &config)
{
  const auto json = configToJsonObject(config);
  const auto temp_path = path + ".tmp";
  {
    std::ofstream file(temp_path, std::ios::trunc);
    if (!file.good())
    {
      throw std::runtime_error("failed to open temporary wall fit config");
    }
    file << json.dump(2) << "\n";
  }
  if (std::rename(temp_path.c_str(), path.c_str()) != 0)
  {
    std::remove(temp_path.c_str());
    throw std::runtime_error("failed to replace wall fit config");
  }
}

std::string wallFitConfigToJson(const WallFitConfig &config)
{
  return configToJsonObject(config).dump(2);
}

}  // namespace origincar_wall
