// encode_bench.cpp - what JPEG XL encoding costs the tower, per input format,
// and whether one save worker keeps up without piling frames up or overheating.
// Colour lossy encodes go through the production src/camera/JpegXlEncoder.cpp
// unchanged; lossless and 16-bit encodes use a bench-only path that sets
// uses_original_profile correctly (production's distance 0 is not lossless,
// known-issues #20). Frames come from camera_bench.py `save`.
//
//   encode_bench matrix   <frames> <out>
//       encode time / size for every input x effort x threads, a distance sweep
//       with PSNR, and the de-Bayer/unpack step on its own. One JSON line each.
//   encode_bench pipeline <frames> <out> <kind> <workers> <threads> <secs> <period_ms> [effort] [distance]
//       kind: rgb8 | bilinear | ea | raw8 | raw12
//       period_ms 0: workers process back to back (maximum throughput).
//       period_ms > 0: one frame is queued every period_ms, as a scan keeps one
//       per waypoint; reports latency and backlog. Samples CPU temperature,
//       clock, load and (if readable) package power once a second.
//   encode_bench matrix12 <frames> <out>
//       only the 12-bit parts of matrix.
//   encode_bench idle     <frames> <out> <secs>
//       the same samples with nothing running: the baseline.
#include "fgc/JpegXlEncoder.h"

#include <jxl/decode.h>
#include <jxl/decode_cxx.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 <sys/resource.h>

#include <algorithm>
#include <atomic>
#include <chrono>
#include <cmath>
#include <condition_variable>
#include <cstdio>
#include <cstdlib>
#include <deque>
#include <filesystem>
#include <fstream>
#include <functional>
#include <mutex>
#include <string>
#include <thread>
#include <vector>

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

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 msSince(Clock::time_point a) {
    return std::chrono::duration<double, std::milli>(Clock::now() - a).count();
}

static long long fsize(const std::string& p) {
    std::error_code ec;
    auto s = fs::file_size(p, ec);
    return ec ? -1 : static_cast<long long>(s);
}

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 encoder: any channel count, 8 or 16-bit container, lossless when
// distance is 0 (with uses_original_profile set, which lossless requires).
static bool encodeBench(const std::string& path, const void* img, int channels, int bits,
                        float distance, 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 = channels;
    bi.uses_original_profile = distance == 0.0f ? JXL_TRUE : JXL_FALSE;
    if (JxlEncoderSetBasicInfo(enc.get(), &bi) != JXL_ENC_SUCCESS) return false;
    JxlColorEncoding ce = {};
    JxlColorEncodingSetToSRGB(&ce, channels < 3 ? JXL_TRUE : JXL_FALSE);
    JxlEncoderSetColorEncoding(enc.get(), &ce);
    auto* fs = JxlEncoderFrameSettingsCreate(enc.get(), nullptr);
    JxlEncoderFrameSettingsSetOption(fs, JXL_ENC_FRAME_SETTING_EFFORT, effort);
    if (distance == 0.0f) {
        JxlEncoderSetFrameDistance(fs, 0);
        JxlEncoderSetFrameLossless(fs, JXL_TRUE);
    } else {
        JxlEncoderSetFrameDistance(fs, distance);
    }
    // 12-bit values in 16-bit words: without this, libjxl reads them as the full
    // 0..65535 range and rescales, so "lossless" keeps only the top bits
    // (measured: every pixel changed, by up to 3839).
    JxlBitDepth depth = {JXL_BIT_DEPTH_FROM_CODESTREAM, 0, 0};
    JxlEncoderSetFrameBitDepth(fs, &depth);
    const bool   wide = bits > 8;
    JxlPixelFormat pf = {static_cast<uint32_t>(channels), wide ? JXL_TYPE_UINT16 : JXL_TYPE_UINT8,
                         JXL_NATIVE_ENDIAN, 0};
    const size_t bytes = static_cast<size_t>(W) * H * channels * (wide ? 2 : 1);
    if (JxlEncoderAddImageFrame(fs, &pf, img, bytes) != JXL_ENC_SUCCESS) return false;
    JxlEncoderCloseInput(enc.get());
    std::vector<uint8_t> buf(1 << 22);
    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;
}

// PSNR (dB) of a decoded 8-bit colour file against the original pixels.
static double psnr8(const std::string& path, const uint8_t* ref) {
    auto data = load(path);
    auto dec  = JxlDecoderMake(nullptr);
    JxlDecoderSubscribeEvents(dec.get(), JXL_DEC_FULL_IMAGE);
    JxlDecoderSetInput(dec.get(), data.data(), data.size());
    JxlDecoderCloseInput(dec.get());
    JxlPixelFormat       pf = {3, JXL_TYPE_UINT8, JXL_NATIVE_ENDIAN, 0};
    std::vector<uint8_t> px;
    for (;;) {
        auto s = JxlDecoderProcessInput(dec.get());
        if (s == JXL_DEC_NEED_IMAGE_OUT_BUFFER) {
            size_t n;
            JxlDecoderImageOutBufferSize(dec.get(), &pf, &n);
            px.resize(n);
            JxlDecoderSetImageOutBuffer(dec.get(), &pf, px.data(), n);
        } else if (s == JXL_DEC_FULL_IMAGE || s == JXL_DEC_SUCCESS) {
            break;
        } else {
            return -1;
        }
    }
    double se = 0;
    for (size_t i = 0; i < px.size(); ++i) {
        double d = double(px[i]) - double(ref[i]);
        se += d * d;
    }
    double mse = se / px.size();
    return mse == 0 ? 99.0 : 10 * std::log10(255.0 * 255.0 / mse);
}

struct Frames {
    std::vector<uint8_t>  rgb, bayer, b12;
    std::vector<uint16_t> u12;
};

// One full save job of the given kind: the conversion it needs, then the encode
// and the file write. Returns the conversion time in ms through `prep_ms`.
static bool runJob(const Frames& f, const std::string& kind, const std::string& path, int effort,
                   float distance, int threads, double* prep_ms = nullptr) {
    auto t0 = Clock::now();
    if (kind == "rgb8") {
        if (prep_ms) *prep_ms = 0;
        return fgc::JpegXlEncoder::encodeToFile(path, f.rgb.data(), W, H, 3, distance, effort,
                                                threads);
    }
    if (kind == "bilinear" || kind == "ea") {
        cv::Mat b(H, W, CV_8UC1, const_cast<uint8_t*>(f.bayer.data())), c;
        // The camera's BayerRG8 is OpenCV's BayerBG (measured 2026-09-30).
        cv::cvtColor(b, c, kind == "ea" ? cv::COLOR_BayerBG2RGB_EA : cv::COLOR_BayerBG2RGB);
        if (prep_ms) *prep_ms = msSince(t0);
        return fgc::JpegXlEncoder::encodeToFile(path, c.data, W, H, 3, distance, effort, threads);
    }
    if (kind == "raw8") {
        if (prep_ms) *prep_ms = 0;
        return encodeBench(path, f.bayer.data(), 1, 8, 0.0f, effort, threads);
    }
    if (kind == "raw12") {
        auto u = unpack12p(f.b12);
        if (prep_ms) *prep_ms = msSince(t0);
        return encodeBench(path, u.data(), 1, 12, 0.0f, effort, threads);
    }
    if (kind == "colour12") {  // 12-bit raw -> 16-bit de-Bayer -> lossy colour
        auto    u = unpack12p(f.b12);
        cv::Mat b(H, W, CV_16UC1, u.data()), c;
        cv::cvtColor(b, c, cv::COLOR_BayerBG2RGB);
        if (prep_ms) *prep_ms = msSince(t0);
        return encodeBench(path, c.data, 3, 12, distance, effort, threads);
    }
    return false;
}

// ---- samples: temperature, clock, CPU load, package power ------------------

static double readNum(const std::string& p, double fallback = -1) {
    std::ifstream f(p);
    double        v;
    return (f >> v) ? v : fallback;
}

static double cpuTempC() {
    double best = -1;
    for (auto& z : fs::directory_iterator("/sys/class/thermal")) {
        if (z.path().filename().string().rfind("thermal_zone", 0) != 0) continue;
        std::ifstream t(z.path() / "type");
        std::string   type;
        t >> type;
        double v = readNum((z.path() / "temp").string());
        if (type == "x86_pkg_temp") return v / 1000;
        best = std::max(best, v / 1000);
    }
    return best;
}

static double avgMhz() {
    double sum = 0;
    int    n   = 0;
    for (int i = 0; i < 64; ++i) {
        double v = readNum("/sys/devices/system/cpu/cpu" + std::to_string(i) +
                           "/cpufreq/scaling_cur_freq");
        if (v < 0) break;
        sum += v / 1000;
        ++n;
    }
    return n ? sum / n : -1;
}

struct CpuTimes { unsigned long long busy = 0, total = 0; };
static CpuTimes cpuTimes() {
    std::ifstream f("/proc/stat");
    std::string   cpu;
    unsigned long long v[8] = {};
    f >> cpu >> v[0] >> v[1] >> v[2] >> v[3] >> v[4] >> v[5] >> v[6] >> v[7];
    CpuTimes t;
    for (auto x : v) t.total += x;
    t.busy = t.total - v[3] - v[4];
    return t;
}

// The package energy counter is root-only. FGC_RAPL names a copy of it that a
// root loop refreshes (see README), so the bench itself need not run as root.
static const std::string kRapl = std::getenv("FGC_RAPL")
                                     ? std::getenv("FGC_RAPL")
                                     : "/sys/class/powercap/intel-rapl:0/energy_uj";

struct Sampler {
    std::atomic<bool>       stop{false};
    std::thread             th;
    std::function<void(std::string&)> extra;  // appends ", \"k\": v" fields
    std::vector<double>     temps, mhz, watts, load;

    void start(double secs_total) {
        th = std::thread([this, secs_total] {
            auto     t0   = Clock::now();
            CpuTimes prev = cpuTimes();
            double   e0   = readNum(kRapl);
            auto     tp   = Clock::now();
            while (!stop && msSince(t0) < secs_total * 1000 + 500) {
                std::this_thread::sleep_for(std::chrono::seconds(1));
                CpuTimes now = cpuTimes();
                double   pct = now.total > prev.total
                                   ? 100.0 * (now.busy - prev.busy) / (now.total - prev.total)
                                   : 0;
                prev = now;
                double e1 = readNum(kRapl);
                double w  = (e0 >= 0 && e1 >= e0) ? (e1 - e0) / 1e6 / (msSince(tp) / 1000) : -1;
                e0 = e1;
                tp = Clock::now();
                double tc = cpuTempC(), mh = avgMhz();
                temps.push_back(tc); mhz.push_back(mh); load.push_back(pct);
                if (w >= 0) watts.push_back(w);
                std::string x;
                if (extra) extra(x);
                std::printf("{\"t\": %.0f, \"temp_c\": %.1f, \"mhz\": %.0f, \"cpu_pct\": %.1f, "
                            "\"watts\": %.2f%s}\n",
                            msSince(t0) / 1000, tc, mh, pct, w, x.c_str());
                std::fflush(stdout);
            }
        });
    }
    void finish() { stop = true; if (th.joinable()) th.join(); }
};

static double mean(const std::vector<double>& v) {
    double s = 0;
    for (double x : v) s += x;
    return v.empty() ? -1 : s / v.size();
}
static double maxOf(const std::vector<double>& v) {
    return v.empty() ? -1 : *std::max_element(v.begin(), v.end());
}
static double minOf(const std::vector<double>& v) {
    return v.empty() ? -1 : *std::min_element(v.begin(), v.end());
}
static double pct(std::vector<double> v, double p) {
    if (v.empty()) return -1;
    std::sort(v.begin(), v.end());
    return v[std::min(v.size() - 1, static_cast<size_t>(p * (v.size() - 1) + 0.5))];
}

static double cpuSecondsSelf() {
    rusage r{};
    getrusage(RUSAGE_SELF, &r);
    return r.ru_utime.tv_sec + r.ru_utime.tv_usec / 1e6 + r.ru_stime.tv_sec + r.ru_stime.tv_usec / 1e6;
}

// ---- modes -------------------------------------------------------------------

static void matrix(const Frames& f, const std::string& out, bool only12) {
    auto timeJob = [&](const std::string& kind, int e, float d, int t, int reps,
                       const std::string& path) {
        std::vector<double> tot, prep;
        for (int i = 0; i < reps; ++i) {
            double p  = 0;
            auto   t0 = Clock::now();
            runJob(f, kind, path, e, d, t, &p);
            tot.push_back(msSince(t0));
            prep.push_back(p);
        }
        std::sort(tot.begin(), tot.end());
        std::sort(prep.begin(), prep.end());
        return std::make_pair(tot[tot.size() / 2], prep[prep.size() / 2]);
    };
    auto emit = [&](const std::string& kind, int e, float d, int t, double ms, double prep,
                    long long bytes, double ps = -1) {
        std::printf("{\"kind\": \"%s\", \"effort\": %d, \"distance\": %.1f, \"threads\": %d, "
                    "\"ms\": %.1f, \"prep_ms\": %.1f, \"bytes\": %lld",
                    kind.c_str(), e, d, t, ms, prep, bytes);
        if (ps >= 0) std::printf(", \"psnr_db\": %.2f", ps);
        std::printf("}\n");
        std::fflush(stdout);
    };

    // Effort x threads at today's distance, for every colour input.
    if (!only12)
    for (const std::string kind : {"rgb8", "bilinear", "ea"})
        for (int t : {1, 2, 3, 4})
            for (int e : {1, 2, 3, 4, 5}) {
                const std::string p = out + "/m_" + kind + ".jxl";
                auto r = timeJob(kind, e, 0.8f, t, e == 5 && t < 3 ? 1 : 3, p);
                emit(kind, e, 0.8f, t, r.first, r.second, fsize(p));
            }
    // Distance sweep at effort 4, 4 threads, with the loss it costs.
    if (!only12)
    for (const std::string kind : {"rgb8", "ea"}) {
        cv::Mat ref;
        if (kind == "ea") {
            cv::Mat b(H, W, CV_8UC1, const_cast<uint8_t*>(f.bayer.data()));
            cv::cvtColor(b, ref, cv::COLOR_BayerBG2RGB_EA);
        }
        const uint8_t* refp = kind == "rgb8" ? f.rgb.data() : ref.data;
        for (float d : {0.5f, 0.8f, 1.0f, 1.5f, 2.0f, 3.0f}) {
            const std::string p = out + "/d_" + kind + ".jxl";
            auto r = timeJob(kind, 4, d, 4, 3, p);
            emit(kind, 4, d, 4, r.first, r.second, fsize(p), psnr8(p, refp));
        }
    }
    // Lossless raw mosaic, 8 and 12 bit.
    for (const std::string kind : {"raw8", "raw12"}) {
        if (only12 && kind == "raw8") continue;
        for (int t : {1, 2, 3, 4})
            for (int e : {1, 2, 3}) {
                const std::string p = out + "/m_" + kind + ".jxl";
                auto r = timeJob(kind, e, 0.0f, t, 3, p);
                emit(kind, e, 0.0f, t, r.first, r.second, fsize(p));
            }
    }
    // 12-bit colour, for comparison with the 8-bit colour paths.
    for (int e : {1, 4}) {
        const std::string p = out + "/m_colour12.jxl";
        auto r = timeJob("colour12", e, 0.8f, 4, 3, p);
        emit("colour12", e, 0.8f, 4, r.first, r.second, fsize(p));
    }
}

static int pipeline(const Frames& f, const std::string& out, const std::string& kind, int workers,
                    int threads, double secs, int period_ms, int effort, float distance) {
    std::mutex                    mu;
    std::condition_variable       cv;
    std::deque<Clock::time_point> queue;  // enqueue times
    std::atomic<long long>        done{0};
    std::atomic<bool>             producing{true}, quit{false};
    std::vector<double>           lat;  // enqueue -> written, ms
    size_t                        max_q = 0;
    const bool                    paced = period_ms > 0;

    Sampler s;
    s.extra = [&](std::string& x) {
        size_t q;
        { std::lock_guard<std::mutex> l(mu); q = queue.size(); }
        x = ", \"done\": " + std::to_string(done.load()) + ", \"queue\": " + std::to_string(q);
    };
    const double cpu0 = cpuSecondsSelf();
    auto         t0   = Clock::now();
    s.start(secs);

    std::vector<std::thread> ws;
    for (int w = 0; w < workers; ++w)
        ws.emplace_back([&, w] {
            const std::string path = out + "/p_" + kind + "_" + std::to_string(w) + ".jxl";
            for (;;) {
                Clock::time_point enq;
                if (paced) {
                    std::unique_lock<std::mutex> l(mu);
                    cv.wait(l, [&] { return !queue.empty() || quit; });
                    if (quit) return;
                    enq = queue.front();
                    queue.pop_front();
                } else {
                    if (msSince(t0) >= secs * 1000) return;
                    enq = Clock::now();
                }
                runJob(f, kind, path, effort, distance, threads);
                ++done;
                std::lock_guard<std::mutex> l(mu);
                lat.push_back(msSince(enq));
            }
        });

    if (paced) {
        auto next = Clock::now();
        while (msSince(t0) < secs * 1000) {
            {
                std::lock_guard<std::mutex> l(mu);
                queue.push_back(Clock::now());
                max_q = std::max(max_q, queue.size());
            }
            cv.notify_one();
            next += std::chrono::milliseconds(period_ms);
            std::this_thread::sleep_until(next);
        }
        std::lock_guard<std::mutex> l(mu);
        quit = true;
    }
    cv.notify_all();
    size_t backlog;
    { std::lock_guard<std::mutex> l(mu); backlog = queue.size(); }
    for (auto& w : ws) w.join();
    const double wall = msSince(t0) / 1000;
    s.finish();
    const double cpu = cpuSecondsSelf() - cpu0;

    std::lock_guard<std::mutex> l(mu);
    std::printf("{\"summary\": \"pipeline\", \"kind\": \"%s\", \"workers\": %d, \"threads\": %d, "
                "\"effort\": %d, \"distance\": %.1f, \"period_ms\": %d, \"secs\": %.1f, "
                "\"frames\": %lld, \"frames_per_s\": %.3f, \"latency_ms_p50\": %.0f, "
                "\"latency_ms_p95\": %.0f, \"latency_ms_max\": %.0f, \"max_queue\": %zu, "
                "\"backlog_at_end\": %zu, \"cpu_s_per_frame\": %.3f, \"temp_c_mean\": %.1f, "
                "\"temp_c_max\": %.1f, \"mhz_mean\": %.0f, \"mhz_min\": %.0f, \"cpu_pct_mean\": %.1f, "
                "\"watts_mean\": %.2f}\n",
                kind.c_str(), workers, threads, effort, distance, period_ms, wall,
                done.load(), done.load() / wall, pct(lat, 0.5), pct(lat, 0.95), maxOf(lat),
                max_q, backlog, done ? cpu / done.load() : -1, mean(s.temps), maxOf(s.temps),
                mean(s.mhz), minOf(s.mhz), mean(s.load), mean(s.watts));
    return 0;
}

int main(int argc, char** argv) {
    if (argc < 4) {
        std::fprintf(stderr, "usage: encode_bench matrix|pipeline|idle <frames> <out> ...\n");
        return 2;
    }
    const std::string mode = argv[1], dir = argv[2], out = argv[3];
    fs::create_directories(out);

    if (mode == "idle") {
        Sampler s;
        const double secs = argc > 4 ? std::atof(argv[4]) : 60;
        s.start(secs);
        std::this_thread::sleep_for(std::chrono::milliseconds(static_cast<int>(secs * 1000)));
        s.finish();
        std::printf("{\"summary\": \"idle\", \"secs\": %.0f, \"temp_c_mean\": %.1f, "
                    "\"temp_c_max\": %.1f, \"mhz_mean\": %.0f, \"cpu_pct_mean\": %.1f, "
                    "\"watts_mean\": %.2f}\n",
                    secs, mean(s.temps), maxOf(s.temps), mean(s.mhz), mean(s.load), mean(s.watts));
        return 0;
    }

    {
        std::ifstream j(dir + "/BayerRG8_0.json");
        std::string   t{std::istreambuf_iterator<char>(j), {}};
        W = jsonInt(t, "width");
        H = jsonInt(t, "height");
    }
    Frames f;
    f.rgb   = load(dir + "/RGB8_0.bin");
    f.bayer = load(dir + "/BayerRG8_0.bin");
    f.b12   = load(dir + "/BayerRG12p_0.bin");
    if (!W || f.rgb.size() != size_t(W) * H * 3 || f.bayer.size() != size_t(W) * H ||
        f.b12.size() != size_t(W) * H * 3 / 2) {
        std::fprintf(stderr, "frames missing or wrong size\n");
        return 1;
    }

    if (mode == "matrix" || mode == "matrix12") {
        matrix(f, out, mode == "matrix12");
        return 0;
    }
    if (mode == "pipeline" && argc >= 9) {
        return pipeline(f, out, argv[4], std::atoi(argv[5]), std::atoi(argv[6]),
                        std::atof(argv[7]), std::atoi(argv[8]),
                        argc > 9 ? std::atoi(argv[9]) : 4,
                        argc > 10 ? static_cast<float>(std::atof(argv[10])) : 0.8f);
    }
    std::fprintf(stderr, "bad arguments\n");
    return 2;
}