#include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/imgcodecs.hpp>
#include <algorithm>
#include <array>
#include <cmath>
#include <iostream>
#include <string>
#include <vector>

using namespace cv;

static std::string argVal(int argc, char** argv, const std::string& key, const std::string& def)
{
    for (int i = 1; i + 1 < argc; ++i)
        if (key == argv[i]) return argv[i + 1];
    return def;
}

struct Det { float x1, y1, x2, y2, score; int cid; };

int main(int argc, char** argv)
{
    std::string model  = argVal(argc, argv, "--model",  "efficientdet-d0_2026jul.onnx");
    std::string image  = argVal(argc, argv, "--image",  "example_outputs/input_image.png");
    std::string output = argVal(argc, argv, "--output", "example_outputs/output_image.png");
    float conf = std::stof(argVal(argc, argv, "--conf", "0.4"));

    const int sz = 512;

    Mat img = imread(image);
    if (img.empty()) { std::cerr << "could not read image: " << image << std::endl; return 1; }

    Mat rgb;
    cvtColor(img, rgb, COLOR_BGR2RGB);
    resize(rgb, rgb, Size(sz, sz));
    if (!rgb.isContinuous()) rgb = rgb.clone();

    int blobShape[] = {1, sz, sz, 3};
    Mat blob(4, blobShape, CV_8U, rgb.data);
    dnn::Net net = dnn::readNetFromONNX(model);
    net.setInput(blob);
    std::vector<Mat> outs;
    net.forward(outs, net.getUnconnectedOutLayersNames());

    const float* boxp = nullptr;
    const float* clsp = nullptr;
    int n = 0, nc = 0;
    for (size_t i = 0; i < outs.size(); ++i)
    {
        const Mat& o = outs[i];
        const float* p = (const float*)o.data;
        int last = o.size[o.dims - 1];
        if (last == 4) { boxp = p; n = o.size[o.dims - 2]; }
        else { clsp = p; nc = last; }
    }

    std::vector<std::array<float, 2>> baseWH;
    double asp[3][2] = {{1.0, 1.0}, {1.4, 0.7}, {0.7, 1.4}};
    for (int i = 0; i < 3; ++i) {
        double s = std::pow(2.0, i / 3.0);
        for (int a = 0; a < 3; ++a)
            baseWH.push_back({(float)(32.0 * s * asp[a][0]), (float)(32.0 * s * asp[a][1])});
    }
    std::vector<float> acx, acy, aw, ah;
    for (int lvl = 0; lvl < 5; ++lvl) {
        int f = sz / (8 << lvl);
        int step = 8 << lvl;
        int m = 1 << lvl;
        for (int y = 0; y < f; ++y)
            for (int x = 0; x < f; ++x) {
                float cx = (x + 0.5f) * step;
                float cy = (y + 0.5f) * step;
                for (auto& b : baseWH) {
                    acx.push_back(cx); acy.push_back(cy);
                    aw.push_back(b[0] * m); ah.push_back(b[1] * m);
                }
            }
    }

    std::vector<Det> dets;
    for (int a = 0; a < n; ++a) {
        const float* bp = boxp + (size_t)a * 4;
        float ycenter = bp[0] * ah[a] + acy[a];
        float xcenter = bp[1] * aw[a] + acx[a];
        float bhv = std::exp(bp[2]) * ah[a];
        float bwv = std::exp(bp[3]) * aw[a];
        const float* cp = clsp + (size_t)a * nc;
        int best = 0; float bestLogit = cp[0];
        for (int c = 1; c < nc; ++c) if (cp[c] > bestLogit) { bestLogit = cp[c]; best = c; }
        float score = 1.0f / (1.0f + std::exp(-bestLogit));
        if (score > conf)
            dets.push_back({(xcenter - bwv / 2) / sz, (ycenter - bhv / 2) / sz,
                            (xcenter + bwv / 2) / sz, (ycenter + bhv / 2) / sz, score, best});
    }

    std::sort(dets.begin(), dets.end(), [](const Det& a, const Det& b) { return a.score > b.score; });
    std::vector<char> removed(dets.size(), 0);
    std::vector<int> pick;
    for (size_t i = 0; i < dets.size(); ++i) {
        if (removed[i]) continue;
        pick.push_back((int)i);
        for (size_t j = i + 1; j < dets.size(); ++j) {
            if (removed[j]) continue;
            float xx1 = std::max(dets[i].x1, dets[j].x1);
            float yy1 = std::max(dets[i].y1, dets[j].y1);
            float xx2 = std::min(dets[i].x2, dets[j].x2);
            float yy2 = std::min(dets[i].y2, dets[j].y2);
            float inter = std::max(0.0f, xx2 - xx1) * std::max(0.0f, yy2 - yy1);
            float ai = (dets[i].x2 - dets[i].x1) * (dets[i].y2 - dets[i].y1);
            float aj = (dets[j].x2 - dets[j].x1) * (dets[j].y2 - dets[j].y1);
            if (inter / (ai + aj - inter + 1e-9f) > 0.6f) removed[j] = 1;
        }
    }

    std::cout << "efficientdet-d0 " << pick.size() << " detections" << std::endl;
    int w = img.cols, h = img.rows;
    for (int idx : pick) {
        const Det& d = dets[idx];
        std::cout << format("%d %.3f %.3f %.3f %.3f %.3f", d.cid, d.score, d.x1, d.y1, d.x2, d.y2) << std::endl;
        rectangle(img, Point((int)(d.x1 * w), (int)(d.y1 * h)), Point((int)(d.x2 * w), (int)(d.y2 * h)), Scalar(0, 255, 0), 2);
        putText(img, format("%d:%.2f", d.cid, d.score), Point((int)(d.x1 * w), (int)(d.y1 * h) - 5), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0, 255, 0), 1);
    }
    imwrite(output, img);
    std::cout << "wrote " << output << std::endl;
    return 0;
}