// Copyright chenyy 2026. Licensed under the MIT License.
// YOLOv5-based armor detector with traditional vision refinement
//
// Uses YOLO to provide ROI regions, then applies traditional binary + light detection
// for precise corner localization. Input/output compatible with armor_detector.
//
// Supports multiple detection modes:
// - HYBRID_SYNC: YOLO runs every frame (original)
// - HYBRID_ROI_CACHE: YOLO runs every N frames, detection runs every frame
// - HYBRID_ASYNC: Async queue with ROI cache
//
// Copyright (C) FYT Vision Group. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "armor_yolo_detect/armor_yolo_detector.hpp"

#include <algorithm>
#include <cmath>

namespace armor_yolo_detect {

bool Detector::init(const std::string& engine_path, const Params& params) {
  this->params = params;
  tensorrt_ = std::make_unique<YoloTensorRT>();

  if (!tensorrt_->initFromEngine(engine_path)) {
    std::cerr << "Failed to initialize TensorRT engine from: " << engine_path << std::endl;
    return false;
  }

  tensorrt_->setPostprocessOptions(params.conf_threshold, params.nms_threshold, 300);

  // Initialize corner corrector for precise corner refinement
  corner_corrector_ = std::make_unique<fyt::auto_aim::LightCornerCorrector>();

  return true;
}

void Detector::setMode(DetectorMode mode) {
  mode_ = mode;
  if (mode_ != DetectorMode::HYBRID_SYNC) {
    // Clear cache when switching to cached modes
    cached_yolo_objects_.clear();
    frame_counter_.store(0, std::memory_order_relaxed);
  }
}

void Detector::setRoiUpdateInterval(int interval) {
  roi_update_interval_ = std::max(1, interval);
}

std::vector<fyt::auto_aim::Armor> Detector::detect(const cv::Mat& input) {
  armors_.clear();
  debug_rois_.clear();

  // Preprocess image once
  cv::Mat gray_img;
  cv::cvtColor(input, gray_img, cv::COLOR_BGR2GRAY);

  cv::Mat full_binary_img;
  cv::threshold(gray_img, full_binary_img, binary_thres, 255, cv::THRESH_BINARY);

  uint64_t current_frame = frame_counter_.fetch_add(1, std::memory_order_relaxed);

  // Update YOLO objects based on mode
  switch (mode_) {
    case DetectorMode::HYBRID_SYNC:
      // YOLO runs every frame
      yolo_objects_ = tensorrt_->infer(input, params.detect_color);
      break;

    case DetectorMode::HYBRID_ROI_CACHE:
    case DetectorMode::HYBRID_ASYNC:
      // YOLO runs every roi_update_interval_ frames
      if (current_frame % roi_update_interval_ == 0) {
        yolo_objects_ = tensorrt_->infer(input, params.detect_color);
        cached_yolo_objects_ = yolo_objects_;
      } else {
        // Use cached ROIs
        yolo_objects_ = cached_yolo_objects_;
      }
      break;
  }

  if (yolo_objects_.empty()) {
    return armors_;
  }

  // Process ROIs to detect armors
  armors_ = processROIs(input, gray_img, full_binary_img);

  return armors_;
}

std::vector<fyt::auto_aim::Armor> Detector::processROIs(
    const cv::Mat& input,
    const cv::Mat& gray_img,
    const cv::Mat& binary_img) {
  std::vector<fyt::auto_aim::Armor> result_armors;

  for (const auto& yolo_obj : yolo_objects_) {
    // Get bounding box from YOLO
    cv::Rect_<float> bbox = yolo_obj.rect;

    // Expand bbox to create ROI
    cv::Rect roi = expandBBox(bbox, input.size(), params.roi_expand_pixel);

    // Clip ROI to image bounds
    roi = roi & cv::Rect(0, 0, input.size().width, input.size().height);

    if (roi.area() <= 0) {
      continue;
    }

    // Store ROI for debug visualization
    debug_rois_.push_back(roi);

    // Extract ROI images
    cv::Mat roi_gray = gray_img(roi);
    cv::Mat roi_binary = binary_img(roi);

    // Find lights within this ROI
    auto lights = findLightsInROI(input, roi_gray, roi_binary, roi);

    if (lights.empty()) {
      continue;
    }

    // Match lights to form armors within this ROI
    auto roi_armors = matchLightsInROI(lights);

    // Process each armor: assign YOLO classification info
    for (auto& armor : roi_armors) {
      // Convert ROI coordinates to global coordinates
      armor.left_light.top += cv::Point2f(static_cast<float>(roi.x), static_cast<float>(roi.y));
      armor.left_light.bottom += cv::Point2f(static_cast<float>(roi.x), static_cast<float>(roi.y));
      armor.left_light.center += cv::Point2f(static_cast<float>(roi.x), static_cast<float>(roi.y));
      armor.right_light.top += cv::Point2f(static_cast<float>(roi.x), static_cast<float>(roi.y));
      armor.right_light.bottom += cv::Point2f(static_cast<float>(roi.x), static_cast<float>(roi.y));
      armor.right_light.center += cv::Point2f(static_cast<float>(roi.x), static_cast<float>(roi.y));
      armor.center = (armor.left_light.center + armor.right_light.center) / 2;

      // Use YOLO's classification result (color and number class)
      armor.classfication_result = numberClassToString(static_cast<NumberClass>(yolo_obj.label));
      armor.confidence = yolo_obj.prob;

      // Use armor type from traditional detection (more accurate)

      // Correct corners using traditional method
      correctCorners(armor, gray_img);

      result_armors.push_back(armor);
    }
  }

  return result_armors;
}

std::vector<YoloObject> Detector::detectRaw(const cv::Mat& input) {
  return tensorrt_->infer(input, params.detect_color);
}

cv::Rect Detector::expandBBox(const cv::Rect_<float>& bbox, const cv::Size& img_size, int pixel_expand) {
  int x = static_cast<int>(bbox.x - pixel_expand);
  int y = static_cast<int>(bbox.y - pixel_expand);
  int w = static_cast<int>(bbox.width + 2 * pixel_expand);
  int h = static_cast<int>(bbox.height + 2 * pixel_expand);

  return cv::Rect(x, y, w, h);
}

std::vector<fyt::auto_aim::Light> Detector::findLightsInROI(
    const cv::Mat& rgb_img,
    const cv::Mat& gray_img,
    const cv::Mat& binary_img,
    const cv::Rect& roi) {
  std::vector<fyt::auto_aim::Light> lights;

  using std::vector;
  vector<vector<cv::Point>> contours;
  vector<cv::Vec4i> hierarchy;
  cv::findContours(binary_img, contours, hierarchy, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);

  lights.reserve(contours.size());

  for (const auto& contour : contours) {
    if (contour.size() < 6) continue;

    fyt::auto_aim::Light light;
    std::vector<cv::Point> points;

    if (light_params.use_fit_line) {
      auto b_rect = cv::boundingRect(contour);
      auto r_rect = cv::minAreaRect(contour);
      cv::Mat mask = cv::Mat::zeros(b_rect.size(), CV_8UC1);
      vector<cv::Point> mask_contour;
      mask_contour.reserve(contour.size());
      for (const auto& p : contour) {
        mask_contour.emplace_back(p - cv::Point(b_rect.x, b_rect.y));
      }
      cv::fillPoly(mask, {mask_contour}, 255);
      cv::findNonZero(mask, points);

      if (points.empty()) continue;

      cv::Vec4f return_param;
      cv::fitLine(points, return_param, cv::DIST_L2, 0, 0.01, 0.01);
      cv::Point2f top, bottom;
      if (int(return_param[0] * 100) == 100 || int(return_param[1] * 100) == 0) {
        top = cv::Point2f(b_rect.x + b_rect.width / 2, b_rect.y);
        bottom = cv::Point2f(b_rect.x + b_rect.width / 2, b_rect.y + b_rect.height);
      } else {
        auto k = return_param[1] / return_param[0];
        auto b = (return_param[3] + b_rect.y) - k * (return_param[2] + b_rect.x);
        top = cv::Point2f((b_rect.y - b) / k, b_rect.y);
        bottom = cv::Point2f((b_rect.y + b_rect.height - b) / k, b_rect.y + b_rect.height);
      }

      double width = std::min(r_rect.size.width, r_rect.size.height);
      double length = cv::norm(top - bottom);
      light = fyt::auto_aim::Light(r_rect, top, bottom, length, width);
    } else {
      light = fyt::auto_aim::Light(contour);
    }

    if (isLight(light)) {
      // Color detection within ROI
      int sum_r = 0;
      int sum_b = 0;
      int sample_count = 0;

      if (light_params.use_fit_line && !points.empty()) {
        auto b_rect = cv::boundingRect(contour);
        for (const auto& point : points) {
          const int x = point.x + b_rect.x;
          const int y = point.y + b_rect.y;
          const auto& pixel = rgb_img.at<cv::Vec3b>(y, x);
          sum_r += pixel[0];
          sum_b += pixel[2];
          ++sample_count;
        }
      } else {
        for (const auto& point : contour) {
          const auto& pixel = rgb_img.at<cv::Vec3b>(point.y, point.x);
          sum_r += pixel[0];
          sum_b += pixel[2];
          ++sample_count;
        }
      }

      if (sample_count > 0 &&
          std::abs(sum_r - sum_b) / sample_count > light_params.color_diff_thresh) {
        light.color = sum_r > sum_b ? EnemyColor::RED : EnemyColor::BLUE;
      }
      lights.emplace_back(light);
    }
  }

  std::sort(lights.begin(), lights.end(), [](const fyt::auto_aim::Light& l1, const fyt::auto_aim::Light& l2) {
    return l1.center.x < l2.center.x;
  });

  return lights;
}

bool Detector::isLight(const fyt::auto_aim::Light& light) {
  float ratio = light.width / light.length;
  bool ratio_ok = light_params.min_ratio < ratio && ratio < light_params.max_ratio;
  bool angle_ok = light.tilt_angle < light_params.max_angle;
  return ratio_ok && angle_ok;
}

bool Detector::containLight(const int i, const int j, const std::vector<fyt::auto_aim::Light>& lights) {
  const fyt::auto_aim::Light& light_1 = lights.at(i);
  const fyt::auto_aim::Light& light_2 = lights.at(j);
  auto points = std::vector<cv::Point2f>{light_1.top, light_1.bottom, light_2.top, light_2.bottom};
  auto bounding_rect = cv::boundingRect(points);
  double avg_length = (light_1.length + light_2.length) / 2.0;
  double avg_width = (light_1.width + light_2.width) / 2.0;

  for (int k = i + 1; k < j; k++) {
    const fyt::auto_aim::Light& test_light = lights.at(k);

    if (test_light.width > 2 * avg_width) {
      continue;
    }
    if (test_light.length < 0.5 * avg_length) {
      continue;
    }

    if (bounding_rect.contains(test_light.top) || bounding_rect.contains(test_light.bottom) ||
        bounding_rect.contains(test_light.center)) {
      return true;
    }
  }
  return false;
}

fyt::auto_aim::ArmorType Detector::isArmor(
    const fyt::auto_aim::Light& light_1,
    const fyt::auto_aim::Light& light_2) {
  float light_length_ratio = light_1.length < light_2.length
      ? light_1.length / light_2.length
      : light_2.length / light_1.length;
  bool light_ratio_ok = light_length_ratio > armor_params.min_light_ratio;

  float avg_light_length = (light_1.length + light_2.length) / 2;
  float center_distance = cv::norm(light_1.center - light_2.center) / avg_light_length;
  bool center_distance_ok =
      (armor_params.min_small_center_distance <= center_distance &&
       center_distance < armor_params.max_small_center_distance) ||
      (armor_params.min_large_center_distance <= center_distance &&
       center_distance < armor_params.max_large_center_distance);

  cv::Point2f diff = light_1.center - light_2.center;
  float angle = std::abs(std::atan(diff.y / diff.x)) / CV_PI * 180;
  bool angle_ok = angle < armor_params.max_angle;

  bool is_armor = light_ratio_ok && center_distance_ok && angle_ok;

  if (is_armor) {
    return center_distance > 3.5 ? fyt::auto_aim::ArmorType::LARGE : fyt::auto_aim::ArmorType::SMALL;
  }
  return fyt::auto_aim::ArmorType::INVALID;
}

std::vector<fyt::auto_aim::Armor> Detector::matchLightsInROI(const std::vector<fyt::auto_aim::Light>& lights) {
  std::vector<fyt::auto_aim::Armor> armors;

  EnemyColor detect_color = params.detect_color == 0 ? EnemyColor::RED : EnemyColor::BLUE;

  for (auto light_1 = lights.begin(); light_1 != lights.end(); light_1++) {
    // Skip lights not matching target color (or accept if color not determined)
    if (light_1->color != EnemyColor::WHITE && light_1->color != detect_color) {
      continue;
    }

    double max_iter_width = light_1->length * armor_params.max_large_center_distance;

    for (auto light_2 = light_1 + 1; light_2 != lights.end(); light_2++) {
      if (light_2->color != EnemyColor::WHITE && light_2->color != detect_color) {
        continue;
      }
      if (containLight(light_1 - lights.begin(), light_2 - lights.begin(), lights)) {
        continue;
      }
      if (light_2->center.x - light_1->center.x > max_iter_width) {
        break;
      }

      auto type = isArmor(*light_1, *light_2);
      if (type != fyt::auto_aim::ArmorType::INVALID) {
        auto armor = fyt::auto_aim::Armor(*light_1, *light_2);
        armor.type = type;
        armors.emplace_back(armor);
      }
    }
  }

  return armors;
}

void Detector::correctCorners(fyt::auto_aim::Armor& armor, const cv::Mat& gray_img) {
  if (corner_corrector_ != nullptr) {
    corner_corrector_->correctCorners(armor, gray_img);
  }
}

void Detector::drawResults(cv::Mat& img) noexcept {
  // Draw ROI rectangles (expanded YOLO bboxes) in blue
  for (const auto& roi : debug_rois_) {
    cv::rectangle(img, roi, cv::Scalar(255, 0, 0), 2);
  }

  // Draw armor outlines
  for (const auto& armor : armors_) {
    // Draw armor outline
    cv::line(
        img, armor.left_light.top, armor.left_light.bottom, cv::Scalar(0, 255, 0), 1, cv::LINE_AA);
    cv::line(
        img, armor.right_light.bottom, armor.right_light.top, cv::Scalar(0, 255, 0), 1, cv::LINE_AA);
    cv::line(
        img, armor.left_light.top, armor.right_light.top, cv::Scalar(0, 255, 0), 1, cv::LINE_AA);
    cv::line(
        img, armor.right_light.bottom, armor.left_light.bottom, cv::Scalar(0, 255, 0), 1, cv::LINE_AA);

    // Draw number and confidence
    std::string text = fyt::auto_aim::armorTypeToString(armor.type) + " " + armor.classfication_result;
    cv::putText(
        img, text, armor.left_light.top, cv::FONT_HERSHEY_SIMPLEX, 0.8, cv::Scalar(0, 255, 255), 2);
  }
}

}  // namespace armor_yolo_detect
