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

// The exported graph is frozen at 384x384 and always takes two images: the
// picture and its horizontal mirror.
static const int INPUT_SIZE = 384;
static const cv::Scalar MEAN(104.00698793, 116.66876762, 122.67891434);

// The 20 LIP classes, in channel order.
static const char* CLASSES[] = {
    "Background", "Hat", "Hair", "Glove", "Sunglasses", "UpperClothes",
    "Dress", "Coat", "Socks", "Pants", "Jumpsuits", "Scarf", "Skirt",
    "Face", "LeftArm", "RightArm", "LeftLeg", "RightLeg", "LeftShoe",
    "RightShoe"
};
static const int NUM_CLASSES = 20;

// LIP palette, RGB.
static const unsigned char COLORS[NUM_CLASSES][3] = {
    {0, 0, 0}, {128, 0, 0}, {255, 0, 0}, {0, 85, 0}, {170, 0, 51},
    {255, 85, 0}, {0, 0, 85}, {0, 119, 221}, {85, 85, 0}, {0, 85, 85},
    {85, 51, 0}, {52, 86, 128}, {0, 128, 0}, {0, 0, 255}, {51, 170, 221},
    {0, 255, 255}, {85, 255, 170}, {170, 255, 85}, {255, 255, 0}, {255, 170, 0}
};

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;
}

// Mirroring an image swaps the left and right body parts, so the channels of
// the mirrored prediction have to be swapped back before the two are averaged.
static int mirrorChannel(int c)
{
    switch (c)
    {
        case 14: return 15;  // LeftArm  <-> RightArm
        case 15: return 14;
        case 16: return 17;  // LeftLeg  <-> RightLeg
        case 17: return 16;
        case 18: return 19;  // LeftShoe <-> RightShoe
        case 19: return 18;
        default: return c;
    }
}

// Build the [2, 3, 384, 384] blob: the image and its mirror, mean-subtracted.
static cv::Mat preprocess(const cv::Mat& image)
{
    cv::Mat resized, mirror;
    cv::resize(image, resized, cv::Size(INPUT_SIZE, INPUT_SIZE));
    cv::flip(resized, mirror, 1);

    std::vector<cv::Mat> images;
    images.push_back(resized);
    images.push_back(mirror);

    return cv::dnn::blobFromImages(images, 1.0, cv::Size(), MEAN);
}

// Average the two predictions and turn them into a label map of `size`.
// `out` is [2, 20, 384, 384].
static cv::Mat postprocess(const cv::Mat& out, const cv::Size& size)
{
    cv::Mat labels(size, CV_8U, cv::Scalar(0));
    cv::Mat best(size, CV_32F, cv::Scalar(-FLT_MAX));

    for (int c = 0; c < NUM_CLASSES; ++c)
    {
        cv::Mat direct(INPUT_SIZE, INPUT_SIZE, CV_32F,
                       const_cast<float*>(out.ptr<float>(0, c)));
        cv::Mat mirrored(INPUT_SIZE, INPUT_SIZE, CV_32F,
                         const_cast<float*>(out.ptr<float>(1, mirrorChannel(c))));

        cv::Mat unmirrored;
        cv::flip(mirrored, unmirrored, 1);

        cv::Mat score = 0.5f * (direct + unmirrored);
        cv::resize(score, score, size, 0, 0, cv::INTER_LINEAR);

        cv::Mat better = score > best;
        score.copyTo(best, better);
        labels.setTo(c, better);
    }
    return labels;
}

// Paint a label map with the LIP palette, as a BGR image.
static cv::Mat colorize(const cv::Mat& labels)
{
    cv::Mat lut(1, NUM_CLASSES, CV_8UC3);
    for (int i = 0; i < NUM_CLASSES; ++i)
        lut.at<cv::Vec3b>(0, i) = cv::Vec3b(COLORS[i][2], COLORS[i][1], COLORS[i][0]);

    cv::Mat segmentation(labels.size(), CV_8UC3);
    for (int y = 0; y < labels.rows; ++y)
    {
        const unsigned char* src = labels.ptr<unsigned char>(y);
        cv::Vec3b* dst = segmentation.ptr<cv::Vec3b>(y);
        for (int x = 0; x < labels.cols; ++x)
            dst[x] = lut.at<cv::Vec3b>(0, src[x]);
    }
    return segmentation;
}

int main(int argc, char** argv)
{
    std::string model  = argVal(argc, argv, "--model",  "human_parsing_jppnet_2026sep.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");

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

    cv::dnn::Net net = cv::dnn::readNetFromONNX(model);
    net.setInput(preprocess(img));
    cv::Mat out = net.forward();

    cv::Mat labels = postprocess(out, img.size());
    cv::Mat segmentation = colorize(labels);

    std::cout << "human_parsing_jppnet input " << img.cols << "x" << img.rows << std::endl;
    for (int c = 0; c < NUM_CLASSES; ++c)
    {
        int pixels = cv::countNonZero(labels == c);
        if (pixels > 0)
            std::cout << "  " << CLASSES[c] << " " << pixels << " px" << std::endl;
    }

    cv::imwrite(output, segmentation);
    std::cout << "wrote " << output << std::endl;
    return 0;
}