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

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

// Run one pyramid head and rebuild a colour image from its luma output.
// Mirrors dnn_superres' preprocess_YCrCb / reconstruct_YCrCb: the network only
// ever sees the Y channel, and Cr/Cb are bicubically upscaled and merged back.
static cv::Mat upsample(cv::dnn::Net& net, const cv::Mat& img, const std::string& nodeName, int scale)
{
    cv::Mat ycrcb;
    cv::cvtColor(img, ycrcb, cv::COLOR_BGR2YCrCb);
    ycrcb.convertTo(ycrcb, CV_32F, 1.0 / 255.0);

    cv::Mat ch[3];
    cv::split(ycrcb, ch);
    cv::Mat Y = ch[0];
    if (!Y.isContinuous()) Y = Y.clone();

    // This model takes NHWC [1, H, W, 1] even though its outputs are NCHW.
    int blobShape[] = {1, Y.rows, Y.cols, 1};
    cv::Mat blob(4, blobShape, CV_32F, Y.data);

    net.setInput(blob);
    cv::Mat outBlob = net.forward(nodeName);

    // Output blob is NCHW [1, 1, H*scale, W*scale] -> wrap as a single-channel Mat.
    cv::Mat yHr(outBlob.size[2], outBlob.size[3], CV_32F, outBlob.ptr<float>());

    cv::Mat crHr, cbHr;
    cv::resize(ch[1], crHr, cv::Size(), scale, scale);
    cv::resize(ch[2], cbHr, cv::Size(), scale, scale);

    std::vector<cv::Mat> merged = {yHr, crHr, cbHr};
    cv::Mat hr;
    cv::merge(merged, hr);
    hr.convertTo(hr, CV_8U, 255.0);
    cv::cvtColor(hr, hr, cv::COLOR_YCrCb2BGR);
    return hr;
}

int main(int argc, char** argv)
{
    std::string model     = argVal(argc, argv, "--model",      "lapsrn_x4_2026sep.onnx");
    std::string image     = argVal(argc, argv, "--image",      "example_outputs/input_image.png");
    std::string outputDir = argVal(argc, argv, "--output-dir", "example_outputs");

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

    // The two heads of the Laplacian pyramid, and the scale each one produces.
    const std::vector<std::pair<std::string, int> > outputs =
        {{"NCHW_output_2x", 2}, {"NCHW_output_4x", 4}};

    std::cout << "lapsrn_x4 input " << img.cols << "x" << img.rows << std::endl;
    for (size_t i = 0; i < outputs.size(); ++i)
    {
        const std::string& nodeName = outputs[i].first;
        int scale = outputs[i].second;

        cv::Mat hr = upsample(net, img, nodeName, scale);
        std::string path = cv::format("%s/output_image_%dx.png", outputDir.c_str(), scale);
        cv::imwrite(path, hr);
        std::cout << "  " << nodeName << " x" << scale
                  << " -> " << hr.cols << "x" << hr.rows << "  " << path << std::endl;
    }
    return 0;
}