#include "yolo_detector.h"
#include <iostream>
#include <fstream>
#include <algorithm>
#include <cstring>
#include <cmath>

using namespace std;

static const vector<string> COCO_NAMES = {
    "person","bicycle","car","motorcycle","airplane","bus","train","truck",
    "boat","traffic light","fire hydrant","stop sign","parking meter","bench",
    "bird","cat","dog","horse","sheep","cow","elephant","bear","zebra",
    "giraffe","backpack","umbrella","handbag","tie","suitcase","frisbee",
    "skis","snowboard","sports ball","kite","baseball bat","baseball glove",
    "skateboard","surfboard","tennis racket","bottle","wine glass","cup",
    "fork","knife","spoon","bowl","banana","apple","sandwich","orange",
    "broccoli","carrot","hot dog","pizza","donut","cake","chair","couch",
    "potted plant","bed","dining table","toilet","tv","laptop","mouse",
    "remote","keyboard","cell phone","microwave","oven","toaster","sink",
    "refrigerator","book","clock","vase","scissors","teddy bear",
    "hair drier","toothbrush"
};

YoloDetector::~YoloDetector() {
    if (ctx_) {
        rknn_destroy(ctx_);
        ctx_ = 0;
    }
}

bool YoloDetector::load(const string& model_path, int input_size,
                         float conf_thresh, float nms_thresh) {
    input_size_ = input_size;
    conf_thresh_ = conf_thresh;
    nms_thresh_ = nms_thresh;

    ifstream ifs(model_path, ios::binary | ios::ate);
    if (!ifs.is_open()) {
        cerr << "[RKNN] Failed to open model: " << model_path << endl;
        return false;
    }
    size_t model_size = ifs.tellg();
    ifs.seekg(0, ios::beg);
    vector<char> model_buf(model_size);
    ifs.read(model_buf.data(), model_size);
    ifs.close();

    int ret = rknn_init(&ctx_, model_buf.data(), model_buf.size(), 0, nullptr);
    if (ret < 0) {
        cerr << "[RKNN] rknn_init failed: " << ret << endl;
        return false;
    }

    rknn_input_output_num io_num;
    ret = rknn_query(ctx_, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
    if (ret < 0) {
        cerr << "[RKNN] query IN_OUT_NUM failed: " << ret << endl;
        return false;
    }
    n_input_ = io_num.n_input;
    n_output_ = io_num.n_output;

    input_attrs_.resize(n_input_);
    for (uint32_t i = 0; i < n_input_; i++) {
        memset(&input_attrs_[i], 0, sizeof(rknn_tensor_attr));
        input_attrs_[i].index = i;
        rknn_query(ctx_, RKNN_QUERY_INPUT_ATTR, &input_attrs_[i], sizeof(rknn_tensor_attr));
    }

    output_attrs_.resize(n_output_);
    for (uint32_t i = 0; i < n_output_; i++) {
        memset(&output_attrs_[i], 0, sizeof(rknn_tensor_attr));
        output_attrs_[i].index = i;
        rknn_query(ctx_, RKNN_QUERY_OUTPUT_ATTR, &output_attrs_[i], sizeof(rknn_tensor_attr));
    }

    cout << "[RKNN] Model loaded: " << model_path << endl;
    for (uint32_t i = 0; i < n_input_; i++) {
        auto& a = input_attrs_[i];
        cout << "  Input  [" << i << "] " << a.name
             << "  fmt=" << (a.fmt == RKNN_TENSOR_NCHW ? "NCHW" : "NHWC")
             << "  dims=(";
        for (uint32_t d = 0; d < a.n_dims; d++) {
            if (d) cout << "x";
            cout << a.dims[d];
        }
        cout << ")" << endl;
    }
    cout << "[RKNN] n_output=" << n_output_ << endl;
    for (uint32_t i = 0; i < n_output_; i++) {
        auto& a = output_attrs_[i];
        cout << "  Output [" << i << "] " << a.name
             << "  fmt=" << (a.fmt == RKNN_TENSOR_NCHW ? "NCHW" : "NHWC")
             << "  dims=(";
        for (uint32_t d = 0; d < a.n_dims; d++) {
            if (d) cout << "x";
            cout << a.dims[d];
        }
        cout << ")" << endl;
    }

    loaded_ = true;
    return true;
}

static float sigmoid(float x) {
    return 1.0f / (1.0f + expf(-x));
}

vector<Detection> YoloDetector::detect(const float* rgb_data, int img_w, int img_h,
                                        float scale, int pad_x, int pad_y) {
    if (!loaded_ || !rgb_data) return {};

    int num_pixels = input_size_ * input_size_ * 3;

    vector<int8_t> input_i8(num_pixels);
    for (int i = 0; i < num_pixels; i++) {
        int v = (int)(rgb_data[i] * 255.0f) - 128;
        input_i8[i] = (int8_t)(v < -128 ? -128 : (v > 127 ? 127 : v));
    }

    rknn_input inputs[1];
    memset(&inputs[0], 0, sizeof(rknn_input));
    inputs[0].index = 0;
    inputs[0].type = RKNN_TENSOR_INT8;
    inputs[0].fmt = RKNN_TENSOR_NHWC;
    inputs[0].size = input_i8.size() * sizeof(int8_t);
    inputs[0].buf = input_i8.data();
    inputs[0].pass_through = 1;

    int ret = rknn_inputs_set(ctx_, 1, inputs);
    if (ret < 0) {
        cerr << "[RKNN] rknn_inputs_set failed: " << ret << endl;
        return {};
    }

    ret = rknn_run(ctx_, nullptr);
    if (ret < 0) {
        cerr << "[RKNN] rknn_run failed: " << ret << endl;
        return {};
    }

    vector<rknn_output> outputs(n_output_);
    for (uint32_t i = 0; i < n_output_; i++) {
        memset(&outputs[i], 0, sizeof(rknn_output));
        outputs[i].want_float = 1;
    }
    ret = rknn_outputs_get(ctx_, n_output_, outputs.data(), nullptr);
    if (ret < 0) {
        cerr << "[RKNN] rknn_outputs_get failed: " << ret << endl;
        return {};
    }

    vector<Detection> dets;

    if (n_output_ == 1) {
        auto& attr = output_attrs_[0];
        int num_preds = attr.dims[1];
        int num_attrs = attr.dims[2];
        float* data = static_cast<float*>(outputs[0].buf);

        if (num_attrs == 85) {
            for (int i = 0; i < num_preds; i++) {
                float* pred = data + i * num_attrs;

                float obj = pred[4];
                if (obj < conf_thresh_) continue;

                float max_score = 0;
                int max_id = -1;
                for (int c = 5; c < 85; c++) {
                    if (pred[c] > max_score) {
                        max_score = pred[c];
                        max_id = c - 5;
                    }
                }

                float final_score = obj * max_score;
                if (final_score < conf_thresh_) continue;

                float cx = pred[0];
                float cy = pred[1];
                float w = pred[2];
                float h = pred[3];

                float x1 = (cx - w * 0.5f - pad_x) / scale;
                float y1 = (cy - h * 0.5f - pad_y) / scale;
                float x2 = (cx + w * 0.5f - pad_x) / scale;
                float y2 = (cy + h * 0.5f - pad_y) / scale;

                int ix1 = max(0, (int)x1);
                int iy1 = max(0, (int)y1);
                int ix2 = min(img_w, (int)x2);
                int iy2 = min(img_h, (int)y2);
                int bw = ix2 - ix1;
                int bh = iy2 - iy1;
                if (bw <= 0 || bh <= 0) continue;

                Detection det;
                det.bbox = cv::Rect(ix1, iy1, bw, bh);
                det.confidence = final_score;
                det.class_id = max_id;
                det.class_name = (max_id >= 0 && max_id < (int)COCO_NAMES.size())
                                 ? COCO_NAMES[max_id] : "unknown";
                dets.push_back(det);
            }
        }
    } else if (n_output_ == 3) {
        static const int ANCHORS[3][6] = {
            {10,13, 16,30, 33,23},
            {30,61, 62,45, 59,119},
            {116,90, 156,198, 373,326}
        };
        int grid_sizes[3] = {80, 40, 20};
        int strides[3] = {8, 16, 32};

        for (int g = 0; g < 3; g++) {
            float* data = static_cast<float*>(outputs[g].buf);
            auto& attr = output_attrs_[g];
            int grid_h = attr.dims[2];
            int grid_w = attr.dims[3];
            int grid_len = grid_h * grid_w;
            int stride = strides[g];

            for (int a = 0; a < 3; a++) {
                for (int i = 0; i < grid_h; i++) {
                    for (int j = 0; j < grid_w; j++) {
                        // YOLOv5 输出 255 通道 = 3锚点 x 85(tx,ty,tw,th,obj,80类),
                        // 锚点步长为 85; 模型已融合 sigmoid, 输出值域 [0,1]
                        int offset = (85 * a + 4) * grid_len + i * grid_w + j;
                        float obj = data[offset];
                        if (obj < conf_thresh_) continue;

                        float* in_ptr = data + (85 * a) * grid_len + i * grid_w + j;
                        float box_x = in_ptr[0] * 2.0f - 0.5f;
                        float box_y = in_ptr[grid_len] * 2.0f - 0.5f;
                        float box_w = in_ptr[2 * grid_len] * 2.0f;
                        float box_h = in_ptr[3 * grid_len] * 2.0f;
                        box_x = (box_x + j) * stride;
                        box_y = (box_y + i) * stride;
                        // 每个尺度的第a个锚点: ANCHORS[g][a*2], ANCHORS[g][a*2+1]
                        box_w = box_w * box_w * ANCHORS[g][a * 2];
                        box_h = box_h * box_h * ANCHORS[g][a * 2 + 1];

                        float max_score = in_ptr[5 * grid_len];
                        int max_id = 0;
                        for (int k = 1; k < 80; k++) {
                            float prob = in_ptr[(5 + k) * grid_len];
                            if (prob > max_score) {
                                max_score = prob;
                                max_id = k;
                            }
                        }

                        float final_score = obj * max_score;
                        if (final_score < conf_thresh_) continue;

                        float x1 = (box_x - box_w * 0.5f - pad_x) / scale;
                        float y1 = (box_y - box_h * 0.5f - pad_y) / scale;
                        float x2 = (box_x + box_w * 0.5f - pad_x) / scale;
                        float y2 = (box_y + box_h * 0.5f - pad_y) / scale;

                        int ix1 = max(0, (int)x1);
                        int iy1 = max(0, (int)y1);
                        int ix2 = min(img_w, (int)x2);
                        int iy2 = min(img_h, (int)y2);
                        int bw = ix2 - ix1;
                        int bh = iy2 - iy1;
                        if (bw <= 0 || bh <= 0) continue;

                        Detection det;
                        det.bbox = cv::Rect(ix1, iy1, bw, bh);
                        det.confidence = final_score;
                        det.class_id = max_id;
                        det.class_name = (max_id >= 0 && max_id < (int)COCO_NAMES.size())
                                         ? COCO_NAMES[max_id] : "unknown";
                        dets.push_back(det);
                    }
                }
            }
        }
    }

    rknn_outputs_release(ctx_, n_output_, outputs.data());
    nms(dets, nms_thresh_);
    return dets;
}

void YoloDetector::nms(vector<Detection>& dets, float threshold) {
    if (dets.empty()) return;

    sort(dets.begin(), dets.end(), [](const Detection& a, const Detection& b) {
        return a.confidence > b.confidence;
    });

    vector<bool> suppressed(dets.size(), false);
    for (size_t i = 0; i < dets.size(); i++) {
        if (suppressed[i]) continue;
        for (size_t j = i + 1; j < dets.size(); j++) {
            if (suppressed[j]) continue;
            float ix1 = dets[i].bbox.x, iy1 = dets[i].bbox.y;
            float ix2 = ix1 + dets[i].bbox.width, iy2 = iy1 + dets[i].bbox.height;
            float jx1 = dets[j].bbox.x, jy1 = dets[j].bbox.y;
            float jx2 = jx1 + dets[j].bbox.width, jy2 = jy1 + dets[j].bbox.height;

            float xx1 = max(ix1, jx1), yy1 = max(iy1, jy1);
            float xx2 = min(ix2, jx2), yy2 = min(iy2, jy2);
            float inter = max(0.f, xx2 - xx1) * max(0.f, yy2 - yy1);
            float area_i = (ix2 - ix1) * (iy2 - iy1);
            float area_j = (jx2 - jx1) * (jy2 - jy1);
            float iou = inter / (area_i + area_j - inter + 1e-6f);
            if (iou > threshold) suppressed[j] = true;
        }
    }

    vector<Detection> result;
    for (size_t i = 0; i < dets.size(); i++)
        if (!suppressed[i]) result.push_back(dets[i]);
    dets = result;
}