#include <filesystem>
#include <iomanip>
#include <iostream>
#include <memory>
#include <sstream>
#include <string>
#include <string_view>
#include <vector>

#include <onnxruntime_cxx_api.h>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>

#include "armor_yolo_detect/armor_yolo_detector.hpp"
#include "armor_yolo_detect/types.hpp"

namespace fs = std::filesystem;

namespace {

inline uint16_t floatToHalf(float f) {
  uint32_t x;
  std::memcpy(&x, &f, sizeof(float));
  uint32_t sign = (x >> 16) & 0x8000;
  uint32_t exp = (x & 0x7f800000) >> 23;
  uint32_t mant = x & 0x007fffff;
  if (exp == 0) return sign;
  if (exp == 255) return sign | 0x7c00 | (mant ? 1 : 0);
  int new_exp = static_cast<int>(exp) - 127 + 15;
  if (new_exp <= 0) return sign;
  if (new_exp >= 31) return sign | 0x7c00;
  return sign | (static_cast<uint32_t>(new_exp) << 10) | (mant >> 13);
}

float sigmoid(float x) {
  if (x > 0) {
    return 1.0f / (1.0f + std::exp(-x));
  }
  return std::exp(x) / (1.0f + std::exp(x));
}

std::vector<float> preprocessToCHW(const cv::Mat& bgr_image) {
  cv::Mat resized;
  cv::resize(bgr_image, resized, cv::Size(640, 640));

  cv::Mat rgb;
  // Disable BGR2RGB. The model natively tracks blue on channel 0, meaning it implies a BGR layout. 
  // cv::cvtColor(resized, rgb, cv::COLOR_BGR2RGB);
  rgb = resized;

  cv::Mat normalized;
  rgb.convertTo(normalized, CV_32FC3, 1.0 / 255.0);

  std::vector<float> chw(3 * 640 * 640);
  std::array<cv::Mat, 3> channels = {
      cv::Mat(640, 640, CV_32F, chw.data() + 0 * 640 * 640),
      cv::Mat(640, 640, CV_32F, chw.data() + 1 * 640 * 640),
      cv::Mat(640, 640, CV_32F, chw.data() + 2 * 640 * 640)};
  cv::split(normalized, channels.data());
  return chw;
}

std::vector<armor_yolo_detect::Armor> detectByOnnx(
    const std::string& onnx_path,
    const cv::Mat& image,
    int detect_color,
    float conf_threshold) {
  static Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "trt_inference_test");

  Ort::SessionOptions session_options;
  session_options.SetIntraOpNumThreads(1);
  session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);

  Ort::Session session(env, onnx_path.c_str(), session_options);
  Ort::AllocatorWithDefaultOptions allocator;

  auto input_name_alloc = session.GetInputNameAllocated(0, allocator);
  auto output_name_alloc = session.GetOutputNameAllocated(0, allocator);
  const char* input_name = input_name_alloc.get();
  const char* output_name = output_name_alloc.get();

  Ort::TypeInfo input_type_info = session.GetInputTypeInfo(0);
  auto input_tensor_info = input_type_info.GetTensorTypeAndShapeInfo();
  auto input_elem_type = input_tensor_info.GetElementType();
  auto input_shape = input_tensor_info.GetShape();
  if (input_shape.size() != 4) {
    throw std::runtime_error("Unexpected ONNX input shape rank");
  }
  input_shape[0] = 1;
  input_shape[1] = 3;
  input_shape[2] = 640;
  input_shape[3] = 640;

  std::vector<float> chw = preprocessToCHW(image);

  Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);

  std::vector<Ort::Value> input_tensors;
  input_tensors.reserve(1);
    std::vector<Ort::Float16_t> input_half;

  if (input_elem_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16) {
    input_half.resize(chw.size());
    for (size_t i = 0; i < chw.size(); ++i) {
      input_half[i] = Ort::Float16_t::FromBits(floatToHalf(chw[i]));
    }
    input_tensors.emplace_back(Ort::Value::CreateTensor<Ort::Float16_t>(
        memory_info, input_half.data(), input_half.size(), input_shape.data(), input_shape.size()));
  } else {
    input_tensors.emplace_back(Ort::Value::CreateTensor<float>(
        memory_info, chw.data(), chw.size(), input_shape.data(), input_shape.size()));
  }

  const char* input_names[] = {input_name};
  const char* output_names[] = {output_name};

  auto output_tensors = session.Run(
      Ort::RunOptions{nullptr}, input_names, input_tensors.data(), 1, output_names, 1);

  if (output_tensors.empty() || !output_tensors[0].IsTensor()) {
    throw std::runtime_error("ONNX output tensor is empty");
  }

  auto output_info = output_tensors[0].GetTensorTypeAndShapeInfo();
  auto output_shape = output_info.GetShape();
  auto output_elem_type = output_info.GetElementType();
  const size_t output_count = output_info.GetElementCount();

  int rows = 0;
  int cols = 0;
  std::cout << "ONNX output shape: ";
  for (auto s : output_shape) std::cout << s << " ";
  std::cout << std::endl;
  if (output_shape.size() >= 3) {
    rows = static_cast<int>(output_shape[output_shape.size() - 2]);
    cols = static_cast<int>(output_shape[output_shape.size() - 1]);
  } else if (output_shape.size() == 2) {
    rows = static_cast<int>(output_shape[0]);
    cols = static_cast<int>(output_shape[1]);
  }
  if (rows <= 0 || cols <= 0) {
    throw std::runtime_error("Unexpected ONNX output shape");
  }

  std::vector<float> output_data(output_count);
  if (output_elem_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16) {
    const Ort::Float16_t* output_half = output_tensors[0].GetTensorData<Ort::Float16_t>();
    for (size_t i = 0; i < output_count; ++i) {
      output_data[i] = output_half[i].ToFloat();
    }
  } else {
    const float* output_float = output_tensors[0].GetTensorData<float>();
    std::copy(output_float, output_float + output_count, output_data.begin());
  }

  std::vector<armor_yolo_detect::Armor> armors;
  armors.reserve(static_cast<size_t>(rows));

  for (int i = 0; i < rows; ++i) {
    const float* row = output_data.data() + static_cast<size_t>(i) * static_cast<size_t>(cols);
    // if (cols < 23) {
    //   continue;
    // }

    const float conf = sigmoid(row[8]);
    if (conf < conf_threshold) {
      continue;
    }

    int color_id = 0;
    float max_color_score = sigmoid(row[9]);
    for (int c = 1; c < 4; ++c) {
      const float score = sigmoid(row[9 + c]);
      if (score > max_color_score) {
        max_color_score = score;
        color_id = c;
      }
    }

    if (color_id == 2 || color_id == 3) {
      continue;
    }
    if (detect_color == 0 && color_id == 1) {
      continue;
    }
    if (detect_color == 1 && color_id == 0) {
      continue;
    }

    int num_id = 0;
    float max_num_score = sigmoid(row[13]);
    for (int c = 1; c < 10; ++c) {
      const float score = sigmoid(row[13 + c]);
      if (score > max_num_score) {
        max_num_score = score;
        num_id = c;
      }
    }

    armor_yolo_detect::YoloObject obj;
    obj.label = num_id;
    obj.prob = conf;
    obj.color = static_cast<armor_yolo_detect::ModelColor>(color_id);
    for (int k = 0; k < 8; ++k) {
      obj.landmarks[k] = row[k];
    }

    const float scale_x = static_cast<float>(image.cols) / 640.0f;
    const float scale_y = static_cast<float>(image.rows) / 640.0f;
    const float max_x = static_cast<float>(std::max(0, image.cols - 1));
    const float max_y = static_cast<float>(std::max(0, image.rows - 1));

    for (int k = 0; k < 4; ++k) {
      const int x_idx = k * 2;
      const int y_idx = x_idx + 1;
      obj.landmarks[x_idx] = std::clamp(obj.landmarks[x_idx] * scale_x, 0.0f, max_x);
      obj.landmarks[y_idx] = std::clamp(obj.landmarks[y_idx] * scale_y, 0.0f, max_y);
    }

    cv::Point2f center(
        (obj.landmarks[0] + obj.landmarks[2] + obj.landmarks[4] + obj.landmarks[6]) / 4.0f,
        (obj.landmarks[1] + obj.landmarks[3] + obj.landmarks[5] + obj.landmarks[7]) / 4.0f);

    armor_yolo_detect::Armor armor(obj, center);
    armor.number_class = static_cast<armor_yolo_detect::NumberClass>(obj.label);
    armor.confidence = obj.prob;
    armors.push_back(armor);
  }
std::vector<int> keep_idx;
  std::vector<float> scores;
  std::vector<cv::Rect> bboxes;
  for (const auto& armor : armors) {
    scores.push_back(armor.confidence);
    cv::Rect bbox = cv::boundingRect(armor.landmark_points);
    bboxes.push_back(bbox);
  }
  cv::dnn::NMSBoxes(bboxes, scores, conf_threshold, 0.1f, keep_idx);

  std::vector<armor_yolo_detect::Armor> nms_armors;
  for (int idx : keep_idx) {
    nms_armors.push_back(armors[idx]);
  }
  return nms_armors;

}

std::vector<armor_yolo_detect::Armor> detectByTrt(
    const std::string& engine_path,
    const cv::Mat& image,
    int detect_color,
    float conf_threshold,
    float nms_threshold) {
  armor_yolo_detect::Detector detector;
  armor_yolo_detect::Detector::Params params;
  params.detect_color = detect_color;
  params.conf_threshold = conf_threshold;
  params.nms_threshold = nms_threshold;

  if (!detector.init(engine_path, params)) {
    throw std::runtime_error("Failed to initialize detector with engine: " + engine_path);
  }

  return detector.detect(image);
}

}  // namespace

int main(int argc, char** argv) {
  if (argc < 4) {
    std::cerr << "Usage: trt_inference_test <model(.trt|.onnx)> <input_image> <output_image> [detect_color] [conf] [nms]\n";
    return 1;
  }

  const std::string model_path = argv[1];
  const std::string input_image_path = argv[2];
  const std::string output_arg = argv[3];

  fs::path workspace_root;
  const char* colcon_prefix = std::getenv("COLCON_CURRENT_PREFIX");
  if (colcon_prefix != nullptr) {
    fs::path prefix_path(colcon_prefix);
    workspace_root = prefix_path.parent_path().parent_path();
  } else {
    workspace_root = fs::current_path();
  }

  fs::path default_output_dir =
      workspace_root / "src/rm_auto_aim/armor_yolo_detect/trt_inference";
  fs::create_directories(default_output_dir);

  fs::path output_path(output_arg);
  if (output_path.is_relative() && output_path.parent_path().empty()) {
    output_path = default_output_dir / output_path;
  }

  int detect_color = 1;
  float conf_threshold = 0.68f;
  float nms_threshold = 0.5f;

  if (argc >= 5) {
    detect_color = std::stoi(argv[4]);
  }
  if (argc >= 6) {
    conf_threshold = std::stof(argv[5]);
  }
  if (argc >= 7) {
    nms_threshold = std::stof(argv[6]);
  }

  if (!fs::exists(model_path)) {
    std::cerr << "Model not found: " << model_path << std::endl;
    return 2;
  }
  if (!fs::exists(input_image_path)) {
    std::cerr << "Input image not found: " << input_image_path << std::endl;
    return 3;
  }

  cv::Mat image = cv::imread(input_image_path, cv::IMREAD_COLOR);
  if (image.empty()) {
    std::cerr << "Failed to read image: " << input_image_path << std::endl;
    return 4;
  }

  std::vector<armor_yolo_detect::Armor> armors;
  try {
    std::string model_ext = fs::path(model_path).extension().string();
    std::transform(model_ext.begin(), model_ext.end(), model_ext.begin(),
                   [](unsigned char c) { return static_cast<char>(std::tolower(c)); });

    if (model_ext == ".onnx") {
      armors = detectByOnnx(model_path, image, detect_color, conf_threshold);
    } else {
      armors = detectByTrt(model_path, image, detect_color, conf_threshold, nms_threshold);
    }
  } catch (const std::exception& e) {
    std::cerr << "Inference failed: " << e.what() << std::endl;
    return 5;
  }

  cv::Mat vis = image.clone();
  cv::putText(vis,
              "det=" + std::to_string(armors.size()),
              cv::Point(20, 40),
              cv::FONT_HERSHEY_SIMPLEX,
              1.0,
              cv::Scalar(0, 255, 255),
              2);

  for (const auto& armor : armors) {
    std::vector<cv::Point> pts;
    pts.reserve(armor.landmark_points.size());
    for (const auto& p : armor.landmark_points) {
      pts.emplace_back(static_cast<int>(p.x), static_cast<int>(p.y));
      cv::circle(vis, cv::Point(static_cast<int>(p.x), static_cast<int>(p.y)), 3, cv::Scalar(255, 0, 0), -1);
    }

    const cv::Point* pts_ptr = pts.data();
    int npts = static_cast<int>(pts.size());
    cv::polylines(vis, &pts_ptr, &npts, 1, true, cv::Scalar(0, 255, 0), 2);

    std::ostringstream label;
    label << armor_yolo_detect::numberClassToString(armor.number_class) << " "
          << std::fixed << std::setprecision(2) << armor.confidence;

    cv::putText(vis,
                label.str(),
                cv::Point(static_cast<int>(armor.center.x), static_cast<int>(armor.center.y)),
                cv::FONT_HERSHEY_SIMPLEX,
                0.6,
                cv::Scalar(0, 255, 255),
                2);
  }

  if (!cv::imwrite(output_path.string(), vis)) {
    std::cerr << "Failed to write output image: " << output_path.string() << std::endl;
    return 6;
  }

  std::cout << "Inference done. model=" << model_path
            << " detections=" << armors.size()
            << " output=" << output_path.string() << std::endl;

  return 0;
}
