// bench.cpp - tower CPU cost per full-resolution frame, using the production
// JpegXlEncoder.cpp unchanged. Frames come from camera_bench.py `save`.
//
//   bench all  <frames_dir> <out_dir>   every measurement, one JSON line each
//   bench load <frames_dir> <out_dir> <secs>   encode colour d0.8 e4 in a loop
//                                               (to stream alongside)
#include "fgc/JpegXlEncoder.h"
#include "fgc/Logger.h"

#include <jxl/encode.h>
#include <jxl/encode_cxx.h>
#include <jxl/thread_parallel_runner.h>
#include <jxl/thread_parallel_runner_cxx.h>

#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>

#include <algorithm>
#include <chrono>
#include <cstdio>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <functional>
#include <iostream>
#include <string>
#include <vector>

namespace fs = std::filesystem;
using Clock  = std::chrono::steady_clock;

// Frame size, from the .json camera_bench.py writes beside each frame.
static int W = 0, H = 0;

static int jsonInt(const std::string& text, const std::string& key) {
    auto at = text.find("\"" + key + "\"");
    return at == std::string::npos ? 0 : std::atoi(text.c_str() + text.find(':', at) + 1);
}

static std::vector<uint8_t> load(const std::string& p) {
    std::ifstream f(p, std::ios::binary);
    return {std::istreambuf_iterator<char>(f), {}};
}

static double ms(Clock::time_point a, Clock::time_point b) {
    return std::chrono::duration<double, std::milli>(b - a).count();
}

// Median of `reps` runs of fn (fn returns nothing; timed around the call).
static double timeit(int reps, const std::function<void()>& fn) {
    std::vector<double> t;
    for (int i = 0; i < reps; ++i) {
        auto a = Clock::now();
        fn();
        t.push_back(ms(a, Clock::now()));
    }
    std::sort(t.begin(), t.end());
    return t[t.size() / 2];
}

static void line(const std::string& what, double t, long long bytes = -1) {
    std::printf("{\"test\": \"%s\", \"ms\": %.1f", what.c_str(), t);
    if (bytes >= 0) std::printf(", \"bytes\": %lld", bytes);
    std::printf("}\n");
    std::fflush(stdout);
}

// GenICam 12p: two pixels in three bytes, little-endian.
static std::vector<uint16_t> unpack12p(const std::vector<uint8_t>& in) {
    std::vector<uint16_t> out(static_cast<size_t>(W) * H);
    for (size_t i = 0, o = 0; o + 1 < out.size(); i += 3, o += 2) {
        out[o]     = static_cast<uint16_t>(in[i] | ((in[i + 1] & 0x0F) << 8));
        out[o + 1] = static_cast<uint16_t>((in[i + 1] >> 4) | (in[i + 2] << 4));
    }
    return out;
}

// Bench-only: lossless single-channel 16-bit JXL carrying 12 significant bits.
// The production encoder is 8-bit only.
static bool encodeGray(const std::string& path, const void* img, int bits, int effort, int threads) {
    auto enc    = JxlEncoderMake(nullptr);
    auto runner = JxlThreadParallelRunnerMake(nullptr, threads);
    JxlEncoderSetParallelRunner(enc.get(), JxlThreadParallelRunner, runner.get());
    JxlBasicInfo bi;
    JxlEncoderInitBasicInfo(&bi);
    bi.xsize = W; bi.ysize = H; bi.bits_per_sample = bits; bi.num_color_channels = 1;
    bi.uses_original_profile = JXL_TRUE;
    if (JxlEncoderSetBasicInfo(enc.get(), &bi) != JXL_ENC_SUCCESS) return false;
    JxlColorEncoding ce = {};
    JxlColorEncodingSetToSRGB(&ce, JXL_TRUE);
    JxlEncoderSetColorEncoding(enc.get(), &ce);
    auto* fs = JxlEncoderFrameSettingsCreate(enc.get(), nullptr);
    JxlEncoderFrameSettingsSetOption(fs, JXL_ENC_FRAME_SETTING_EFFORT, effort);
    JxlEncoderSetFrameDistance(fs, 0);
    JxlEncoderSetFrameLossless(fs, JXL_TRUE);
    // 12-bit values in 16-bit words must be read at the header's bit depth, or
    // libjxl rescales them from 0..65535 and the result is not lossless.
    JxlBitDepth depth = {JXL_BIT_DEPTH_FROM_CODESTREAM, 0, 0};
    JxlEncoderSetFrameBitDepth(fs, &depth);
    JxlPixelFormat pf = {1, bits > 8 ? JXL_TYPE_UINT16 : JXL_TYPE_UINT8, JXL_NATIVE_ENDIAN, 0};
    if (JxlEncoderAddImageFrame(fs, &pf, img, static_cast<size_t>(W) * H * (bits > 8 ? 2 : 1)) != JXL_ENC_SUCCESS)
        return false;
    JxlEncoderCloseInput(enc.get());
    std::vector<uint8_t> buf(1 << 20);
    uint8_t* next = buf.data();
    size_t   avail = buf.size();
    JxlEncoderStatus r;
    while ((r = JxlEncoderProcessOutput(enc.get(), &next, &avail)) == JXL_ENC_NEED_MORE_OUTPUT) {
        size_t off = next - buf.data();
        buf.resize(buf.size() * 2);
        next = buf.data() + off;
        avail = buf.size() - off;
    }
    if (r != JXL_ENC_SUCCESS) return false;
    std::ofstream(path, std::ios::binary).write(reinterpret_cast<char*>(buf.data()),
                                                next - buf.data());
    return true;
}

static long long fsize(const std::string& p) { return static_cast<long long>(fs::file_size(p)); }

int main(int argc, char** argv) {
    if (argc < 4) { std::fprintf(stderr, "usage: bench all|load <frames> <out> [secs]\n"); return 2; }
    const std::string mode = argv[1], dir = argv[2], out = argv[3];
    fs::create_directories(out);
    {
        std::ifstream j(dir + "/BayerRG8_0.json");
        std::string   t{std::istreambuf_iterator<char>(j), {}};
        W = jsonInt(t, "width");
        H = jsonInt(t, "height");
    }
    auto rgb   = load(dir + "/RGB8_0.bin");
    auto bayer = load(dir + "/BayerRG8_0.bin");
    auto b12   = load(dir + "/BayerRG12p_0.bin");
    if (rgb.size() != size_t(W) * H * 3 || bayer.size() != size_t(W) * H || b12.empty()) {
        std::fprintf(stderr, "frames missing or wrong size\n");
        return 1;
    }

    if (mode == "raw8true") {
        for (int e : {1, 2, 3}) {
            std::string p = out + "/raw8true_e" + std::to_string(e) + ".jxl";
            double t = timeit(3, [&] { encodeGray(p, bayer.data(), 8, e, 3); });
            line("raw8_truelossless_e" + std::to_string(e) + "_t3", t, fsize(p));
        }
        return 0;
    }

    if (mode == "load") {
        const double secs = argc > 4 ? std::atof(argv[4]) : 60;
        auto t0 = Clock::now();
        int  n  = 0;
        while (ms(t0, Clock::now()) < secs * 1000) {
            fgc::JpegXlEncoder::encodeToFile(out + "/load.jxl", rgb.data(), W, H, 3, 0.8f, 4);
            ++n;
        }
        line("load_colour_d0.8_e4_count" + std::to_string(n), ms(t0, Clock::now()) / n);
        return 0;
    }

    // 1. The copy FrameReceived makes of every delivered frame.
    std::vector<uint8_t> dst;
    line("copy_rgb8_48.6MB", timeit(5, [&] { dst.assign(rgb.begin(), rgb.end()); }));
    line("copy_bayer8_16.2MB", timeit(5, [&] { dst.assign(bayer.begin(), bayer.end()); }));

    // 2. De-Bayer. The camera's BayerRG8 is OpenCV's BayerBG (measured: matches
    //    the camera's own RGB8 channel means).
    cv::Mat b(H, W, CV_8UC1, bayer.data()), c;
    line("debayer_bilinear", timeit(5, [&] { cv::cvtColor(b, c, cv::COLOR_BayerBG2RGB); }));
    line("debayer_ea", timeit(3, [&] { cv::cvtColor(b, c, cv::COLOR_BayerBG2RGB_EA); }));
    line("debayer_vng", timeit(3, [&] { cv::cvtColor(b, c, cv::COLOR_BayerBG2RGB_VNG); }));
    cv::Mat cam_rgb(H, W, CV_8UC3, rgb.data()), rot;
    line("rotate_180_rgb8", timeit(5, [&] { cv::rotate(cam_rgb, rot, cv::ROTATE_180); }));
    cv::cvtColor(b, c, cv::COLOR_BayerBG2RGB);

    // 3. Colour JXL, lossy as production: camera RGB8 and tower-de-Bayered.
    for (int threads : {3, 4})
        for (int e : {1, 2, 3, 4, 5}) {
            std::string p = out + "/rgb_d08_e" + std::to_string(e) + ".jxl";
            double t = timeit(3, [&] {
                fgc::JpegXlEncoder::encodeToFile(p, rgb.data(), W, H, 3, 0.8f, e, threads);
            });
            line("colour_camRGB8_d0.8_e" + std::to_string(e) + "_t" + std::to_string(threads), t,
                 fsize(p));
        }
    for (int e : {1, 4}) {
        std::string p = out + "/deb_d08_e" + std::to_string(e) + ".jxl";
        double t = timeit(3, [&] {
            fgc::JpegXlEncoder::encodeToFile(p, c.data, W, H, 3, 0.8f, e, 3);
        });
        line("colour_towerDebayer_d0.8_e" + std::to_string(e) + "_t3", t, fsize(p));
    }

    // 4. Raw mosaic, lossless, single channel.
    for (int e : {1, 2, 3, 4}) {
        std::string p = out + "/raw8_e" + std::to_string(e) + ".jxl";
        double t = timeit(3, [&] {
            fgc::JpegXlEncoder::encodeToFile(p, bayer.data(), W, H, 1, 0.0f, e, 3);
        });
        line("raw8_lossless_e" + std::to_string(e) + "_t3", t, fsize(p));
    }
    std::vector<uint16_t> u;
    line("unpack_12p", timeit(3, [&] { u = unpack12p(b12); }));
    for (int e : {1, 2, 3}) {
        std::string p = out + "/raw12_e" + std::to_string(e) + ".jxl";
        double t = timeit(3, [&] { encodeGray(p, u.data(), 12, e, 3); });
        line("raw12_lossless_e" + std::to_string(e) + "_t3", t, fsize(p));
    }
    return 0;
}