#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 Layer { float mn, mx; std::vector<int> ars; int step, fm; };

int main(int argc, char** argv)
{
    std::string model  = argVal(argc, argv, "--model",  "opencv_face_detector_uint8_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 thr = std::stof(argVal(argc, argv, "--conf", "0.4"));

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

    Mat inp;
    resize(img, inp, Size(sz, sz));
    inp.convertTo(inp, CV_32F);
    subtract(inp, Scalar(104, 177, 123), inp);
    if (!inp.isContinuous()) inp = inp.clone();

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

    const float* loc = nullptr;
    const float* conf = nullptr;
    for (size_t i = 0; i < outs.size(); ++i)
    {
        const Mat& o = outs[i];
        size_t tot = o.total();
        const float* p = (const float*)o.data;
        if (tot == 35568) loc = p;
        else if (tot == 17784) conf = p;
    }

    std::vector<Layer> layers = {
        {30, 60, {2}, 8, 38},
        {60, 111, {2, 3}, 16, 19},
        {111, 162, {2, 3}, 32, 10},
        {162, 213, {2, 3}, 64, 5},
        {213, 264, {2}, 100, 5},
        {264, 315, {2}, 300, 5},
    };
    std::vector<Vec4f> priors;
    for (const Layer& L : layers)
    {
        std::vector<float> ratios = {1.0f};
        for (int a : L.ars) { ratios.push_back((float)a); ratios.push_back(1.0f / a); }
        for (int y = 0; y < L.fm; ++y)
            for (int x = 0; x < L.fm; ++x)
            {
                float cx = (x + 0.5f) * L.step;
                float cy = (y + 0.5f) * L.step;
                std::vector<Vec2f> boxes = {{L.mn, L.mn}, {std::sqrt(L.mn * L.mx), std::sqrt(L.mn * L.mx)}};
                for (size_t k = 1; k < ratios.size(); ++k)
                {
                    float a = ratios[k];
                    boxes.push_back({L.mn * std::sqrt(a), L.mn / std::sqrt(a)});
                }
                for (const Vec2f& b : boxes)
                    priors.push_back({cx, cy, b[0], b[1]});
            }
    }

    const float var[4] = {0.1f, 0.1f, 0.2f, 0.2f};
    int n = (int)priors.size();
    std::vector<Rect2f> boxes;
    std::vector<float> scores;
    for (int i = 0; i < n; ++i)
    {
        float c0 = conf[i * 2], c1 = conf[i * 2 + 1];
        float m = std::max(c0, c1);
        float e0 = std::exp(c0 - m), e1 = std::exp(c1 - m);
        float s = e1 / (e0 + e1);
        if (s <= thr) continue;
        float pcx = priors[i][0] / sz, pcy = priors[i][1] / sz;
        float pw = priors[i][2] / sz, ph = priors[i][3] / sz;
        float cx = pcx + loc[i * 4] * var[0] * pw;
        float cy = pcy + loc[i * 4 + 1] * var[1] * ph;
        float bw = pw * std::exp(loc[i * 4 + 2] * var[2]);
        float bh = ph * std::exp(loc[i * 4 + 3] * var[3]);
        boxes.push_back(Rect2f(cx - bw / 2, cy - bh / 2, bw, bh));
        scores.push_back(s);
    }

    std::vector<int> order(scores.size());
    for (size_t i = 0; i < order.size(); ++i) order[i] = (int)i;
    std::sort(order.begin(), order.end(), [&](int a, int b){ return scores[a] > scores[b]; });
    std::vector<char> removed(order.size(), 0);
    std::vector<int> pick;
    for (size_t oi = 0; oi < order.size(); ++oi)
    {
        if (removed[oi]) continue;
        int i = order[oi];
        pick.push_back(i);
        for (size_t oj = oi + 1; oj < order.size(); ++oj)
        {
            if (removed[oj]) continue;
            int j = order[oj];
            const Rect2f& a = boxes[i];
            const Rect2f& b = boxes[j];
            float xx1 = std::max(a.x, b.x), yy1 = std::max(a.y, b.y);
            float xx2 = std::min(a.x + a.width, b.x + b.width);
            float yy2 = std::min(a.y + a.height, b.y + b.height);
            float inter = std::max(0.f, xx2 - xx1) * std::max(0.f, yy2 - yy1);
            float iou = inter / (a.area() + b.area() - inter + 1e-9f);
            if (iou > 0.3f) removed[oj] = 1;
        }
    }

    int W = img.cols, H = img.rows;
    for (int i : pick)
    {
        const Rect2f& b = boxes[i];
        rectangle(img, Point(int(b.x * W), int(b.y * H)),
                  Point(int((b.x + b.width) * W), int((b.y + b.height) * H)), Scalar(0, 255, 0), 2);
    }
    imwrite(output, img);
    std::cout << "opencv_face_detector_uint8 " << pick.size() << " faces" << std::endl;
    return 0;
}