// Copyright chenyy 2026. Licensed under the MIT License.
// Modified for YOLOv5-based detection
//
// 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.

#ifndef ARMOR_YOLO_DETECT_TYPES_HPP_
#define ARMOR_YOLO_DETECT_TYPES_HPP_

// std
#include <algorithm>
#include <array>
#include <numeric>
#include <string>
#include <vector>
// 3rd party
#include <Eigen/Dense>
#include <opencv2/core.hpp>
#include <opencv2/core/types.hpp>
#include <opencv2/imgproc.hpp>
// project
#include "armor_detector/types.hpp"
#include "rm_utils/assert.hpp"
#include "rm_utils/common.hpp"

namespace armor_yolo_detect {

// Armor size, Unit: m
constexpr double SMALL_ARMOR_WIDTH = 133.0 / 1000.0;
constexpr double SMALL_ARMOR_HEIGHT = 50.0 / 1000.0;
constexpr double LARGE_ARMOR_WIDTH = 225.0 / 1000.0;
constexpr double LARGE_ARMOR_HEIGHT = 50.0 / 1000.0;

// Armor type
enum class ArmorType { SMALL, LARGE, INVALID };

inline std::string armorTypeToString(const ArmorType& type) {
  switch (type) {
    case ArmorType::SMALL:
      return "small";
    case ArmorType::LARGE:
      return "large";
    default:
      return "invalid";
  }
}

// Class labels for number classification
enum class NumberClass {
  GUARD = 0,
  ONE = 1,
  TWO = 2,
  THREE = 3,
  FOUR = 4,
  FIVE = 5,
  OUTPOST = 6,
  BASE_SMALL = 7,
  BASE_LARGE = 8,
  NEGATIVE = 9
};

inline std::string numberClassToString(NumberClass cls) {
  switch (cls) {
    case NumberClass::GUARD:
      return "G";
    case NumberClass::ONE:
      return "1";
    case NumberClass::TWO:
      return "2";
    case NumberClass::THREE:
      return "3";
    case NumberClass::FOUR:
      return "4";
    case NumberClass::FIVE:
      return "5";
    case NumberClass::OUTPOST:
      return "O";
    case NumberClass::BASE_SMALL:
      return "Bs";
    case NumberClass::BASE_LARGE:
      return "Bb";
    default:
      return "negative";
  }
}

// Color classes from model output
enum class ModelColor {
  RED = 0,
  BLUE = 1,
  GRAY = 2,
  PURPLE = 3
};

// YOLO detection object structure
struct YoloObject {
  cv::Rect_<float> rect;
  float landmarks[8];
  int label;
  float prob;
  ModelColor color;
};

// Struct to store detected armor with 6 landmarks (for PnP)
struct Armor {
  static constexpr const int N_LANDMARKS = 6;
  static constexpr const int N_LANDMARKS_2 = N_LANDMARKS * 2;

  Armor() = default;

  Armor(const YoloObject& obj, const cv::Point2f& center)
      : center(center), type(ArmorType::SMALL), confidence(obj.prob) {
    const cv::Point2f top_left(obj.landmarks[0], obj.landmarks[1]);
    const cv::Point2f bottom_left(obj.landmarks[2], obj.landmarks[3]);
    const cv::Point2f bottom_right(obj.landmarks[4], obj.landmarks[5]);
    const cv::Point2f top_right(obj.landmarks[6], obj.landmarks[7]);

    landmark_points = {
      bottom_left,
      (bottom_left + top_left) * 0.5f,
      top_left,
      top_right,
      (top_right + bottom_right) * 0.5f,
      bottom_right
    };

    float width = cv::norm(landmark_points[2] - landmark_points[3]);
    float height = cv::norm(landmark_points[1] - landmark_points[4]);
    if (height > 0) {
      float ratio = width / height;
      type = (ratio > 3.5) ? ArmorType::LARGE : ArmorType::SMALL;
    }
  }

  template <typename PointType>
  static inline std::vector<PointType> buildObjectPoints(const double& w, const double& h) noexcept {
    return {PointType(0, w / 2, -h / 2),    PointType(0, w / 2, 0),
            PointType(0, w / 2, h / 2),     PointType(0, -w / 2, h / 2),
            PointType(0, -w / 2, 0),        PointType(0, -w / 2, -h / 2)};
  }

  std::vector<cv::Point2f> landmarks() const {
    return {landmark_points[0], landmark_points[1], landmark_points[2],
            landmark_points[3], landmark_points[4], landmark_points[5]};
  }

  fyt::auto_aim::Light toLeftLight() const {
    fyt::auto_aim::Light light;
    light.center = (landmark_points[0] + landmark_points[2]) / 2;
    light.top = landmark_points[2];
    light.bottom = landmark_points[0];
    light.length = cv::norm(landmark_points[2] - landmark_points[0]);
    light.width =
        cv::norm(landmark_points[1] - (landmark_points[0] + landmark_points[2]) / 2) * 2;
    return light;
  }

  fyt::auto_aim::Light toRightLight() const {
    fyt::auto_aim::Light light;
    light.center = (landmark_points[3] + landmark_points[5]) / 2;
    light.top = landmark_points[3];
    light.bottom = landmark_points[5];
    light.length = cv::norm(landmark_points[3] - landmark_points[5]);
    light.width =
        cv::norm(landmark_points[4] - (landmark_points[3] + landmark_points[5]) / 2) * 2;
    return light;
  }

  fyt::auto_aim::Armor toAutoAimArmor() const {
    fyt::auto_aim::Armor armor;
    switch (type) {
      case ArmorType::SMALL:
        armor.type = fyt::auto_aim::ArmorType::SMALL;
        break;
      case ArmorType::LARGE:
        armor.type = fyt::auto_aim::ArmorType::LARGE;
        break;
      default:
        armor.type = fyt::auto_aim::ArmorType::INVALID;
        break;
    }
    armor.center = center;
    armor.confidence = confidence;
    armor.classfication_result = numberClassToString(number_class);
    armor.left_light = toLeftLight();
    armor.right_light = toRightLight();
    return armor;
  }

  std::array<cv::Point2f, 6> landmark_points;
  cv::Point2f center;
  ArmorType type;
  NumberClass number_class = NumberClass::NEGATIVE;
  float confidence = 0.0f;
};

}  // namespace armor_yolo_detect
#endif  // ARMOR_YOLO_DETECT_TYPES_HPP_
