// run_model.cpp
// Hailo inference script – mirrors run_model.py (up to CSV output).
//
// Build (adjust include/library paths to match your HailoRT installation):
//   g++ -std=c++17 -O2 -o run_model \
//       run_model.cpp \
//       -I/usr/include/hailo \
//       -lhailort \
//       -lpthread
//
// Run:
//   ./run_model
//
// Outputs:
//   predictions_latency.csv  – one row per sample:
//       latency_us

#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <string>
#include <vector>

#include "hailo/hailort.hpp"

// ─────────────────────────────────────────────────────────────────────────────
// Configuration
// ─────────────────────────────────────────────────────────────────────────────
static constexpr const char* HEF_PATH          = "model_t8.hef";
static constexpr const char* CSV_OUTPUT_PATH   = "model_t8_inferences_latency.csv";
static constexpr int         INPUT_FEATURES    = 10;    // number of input features
static constexpr int         OUTPUT_FEATURES   = 12;    // number of output features
static constexpr int         WORK_ITERS        = 100;  // CPU warm-up iterations
static constexpr int         NUM_DUMMY_ROWS    = 100;  // number of dummy input rows

// ─────────────────────────────────────────────────────────────────────────────
// Helpers
// ─────────────────────────────────────────────────────────────────────────────

// Generate dummy float32 inputs of shape [NUM_DUMMY_ROWS, INPUT_FEATURES].
static std::vector<std::vector<float>> generate_dummy_inputs() {
    std::vector<std::vector<float>> rows;
    rows.reserve(NUM_DUMMY_ROWS);

    for (int i = 0; i < NUM_DUMMY_ROWS; ++i) {
        std::vector<float> row;
        row.reserve(INPUT_FEATURES);

        for (int j = 0; j < INPUT_FEATURES; ++j) {
            // Deterministic dummy float32 data
            float value = 0.01f * static_cast<float>(i * INPUT_FEATURES + j + 1);
            row.push_back(value);
        }

        rows.push_back(std::move(row));
    }

    return rows;
}

// Lightweight CPU work to mirror the Python dummy-compute loop (keeps CPU
// out of deep idle states between inference calls).
static void dummy_compute() {
    volatile double acc = 0.0;
    double x = 0.001;
    for (int i = 0; i < WORK_ITERS; ++i) {
        x = std::fmod(x * 1.0000001 + 0.000001, 1.0);
        acc += std::sin(x) * std::cos(x) + std::sqrt(x + 1e-9);
    }
    (void)acc; // prevent optimisation
}

// ─────────────────────────────────────────────────────────────────────────────
// Main
// ─────────────────────────────────────────────────────────────────────────────
int main() {
    using namespace hailort;

    // ── 1. Load HEF ──────────────────────────────────────────────────────────
    auto hef_exp = Hef::create(HEF_PATH);
    if (!hef_exp) {
        std::cerr << "Failed to load HEF: " << hef_exp.status() << "\n";
        return EXIT_FAILURE;
    }
    auto hef = std::move(hef_exp.value());

    // ── 2. Create VDevice ────────────────────────────────────────────────────
    auto vdevice_exp = VDevice::create();
    if (!vdevice_exp) {
        std::cerr << "Failed to create VDevice: " << vdevice_exp.status() << "\n";
        return EXIT_FAILURE;
    }
    auto vdevice = std::move(vdevice_exp.value());

    // ── 3. Configure network group ───────────────────────────────────────────
    auto configure_params_exp = hef.create_configure_params(HAILO_STREAM_INTERFACE_PCIE);
    if (!configure_params_exp) {
        std::cerr << "Failed to create configure params: " << configure_params_exp.status() << "\n";
        return EXIT_FAILURE;
    }

    auto network_groups_exp = vdevice->configure(hef, configure_params_exp.value());
    if (!network_groups_exp) {
        std::cerr << "Failed to configure network: " << network_groups_exp.status() << "\n";
        return EXIT_FAILURE;
    }
    auto& network_groups = network_groups_exp.value();
    auto& network_group  = network_groups[0];

    // ── 4. Build VStream params ───────────────────────────────────────────────
    auto input_params_exp = network_group->make_input_vstream_params(
        false,
        HAILO_FORMAT_TYPE_FLOAT32,
        HAILO_DEFAULT_VSTREAM_TIMEOUT_MS,
        HAILO_DEFAULT_VSTREAM_QUEUE_SIZE
    );
    if (!input_params_exp) {
        std::cerr << "Failed to create input vstream params: " << input_params_exp.status() << "\n";
        return EXIT_FAILURE;
    }

    auto output_params_exp = network_group->make_output_vstream_params(
        false,
        HAILO_FORMAT_TYPE_FLOAT32,
        HAILO_DEFAULT_VSTREAM_TIMEOUT_MS,
        HAILO_DEFAULT_VSTREAM_QUEUE_SIZE
    );
    if (!output_params_exp) {
        std::cerr << "Failed to create output vstream params: " << output_params_exp.status() << "\n";
        return EXIT_FAILURE;
    }

    // ── 5. Create VStreams ────────────────────────────────────────────────────
    auto input_vstreams_exp = VStreamsBuilder::create_input_vstreams(
        *network_group, input_params_exp.value());
    if (!input_vstreams_exp) {
        std::cerr << "Failed to create input vstreams: " << input_vstreams_exp.status() << "\n";
        return EXIT_FAILURE;
    }
    auto& input_vstreams = input_vstreams_exp.value();

    auto output_vstreams_exp = VStreamsBuilder::create_output_vstreams(
        *network_group, output_params_exp.value());
    if (!output_vstreams_exp) {
        std::cerr << "Failed to create output vstreams: " << output_vstreams_exp.status() << "\n";
        return EXIT_FAILURE;
    }
    auto& output_vstreams = output_vstreams_exp.value();

    // Print stream shapes for quick sanity-check
    {
        const auto& in_info  = input_vstreams[0].get_info();
        const auto& out_info = output_vstreams[0].get_info();
        std::cout << "Input  stream : " << in_info.name
                  << "  shape=" << in_info.shape.height
                  << "x" << in_info.shape.width
                  << "x" << in_info.shape.features << "\n";
        std::cout << "Output stream : " << out_info.name
                  << "  shape=" << out_info.shape.height
                  << "x" << out_info.shape.width
                  << "x" << out_info.shape.features << "\n";
    }

    // ── 6. Generate dummy inputs instead of loading CSV ──────────────────────
    auto inputs = generate_dummy_inputs();
    std::cout << "Generated " << inputs.size()
              << " dummy samples with " << INPUT_FEATURES
              << " float32 features each\n";

    // Pre-allocate output buffer (float32, OUTPUT_FEATURES elements, batch=1)
    const std::size_t out_frame_size = OUTPUT_FEATURES * sizeof(float);
    std::vector<float> out_buf(OUTPUT_FEATURES);

    std::vector<double> latencies_us;
    latencies_us.reserve(inputs.size());

    // ── 7. Run inference ──────────────────────────────────────────────────────
    // NOTE: When using VDevice the built-in scheduler manages network-group
    // activation automatically. Just write/read the vstreams directly.
    for (auto& sample : inputs) {
        dummy_compute();

        auto write_status = input_vstreams[0].write(
            MemoryView(sample.data(), INPUT_FEATURES * sizeof(float)));
        if (write_status != HAILO_SUCCESS) {
            std::cerr << "Write failed: " << write_status << "\n";
            return EXIT_FAILURE;
        }

        auto t0 = std::chrono::steady_clock::now();
        auto read_status = output_vstreams[0].read(
            MemoryView(out_buf.data(), out_frame_size));
        auto t1 = std::chrono::steady_clock::now();

        if (read_status != HAILO_SUCCESS) {
            std::cerr << "Read failed: " << read_status << "\n";
            return EXIT_FAILURE;
        }

        double lat_us = std::chrono::duration<double, std::micro>(t1 - t0).count();
        latencies_us.push_back(lat_us);
    }

    // ── 8. Print basic latency statistics ────────────────────────────────────
    {
        const std::size_t n = latencies_us.size();
        double sum = 0.0, min_v = latencies_us[0], max_v = latencies_us[0];
        for (double v : latencies_us) {
            sum += v;
            min_v = std::min(min_v, v);
            max_v = std::max(max_v, v);
        }
        double mean_v = sum / static_cast<double>(n);

        std::cout << std::fixed << std::setprecision(3)
                  << "\nLatency statistics (µs):\n"
                  << "  Samples : " << n      << "\n"
                  << "  Mean    : " << mean_v << "\n"
                  << "  Min     : " << min_v  << "\n"
                  << "  Max     : " << max_v  << "\n";
    }

    // ── 9. Save results to CSV ────────────────────────────────────────────────
    {
        std::ofstream out(CSV_OUTPUT_PATH);
        if (!out.is_open()) {
            std::cerr << "Cannot open output CSV: " << CSV_OUTPUT_PATH << "\n";
            return EXIT_FAILURE;
        }

        out << "latency_us\n";

        out << std::fixed << std::setprecision(6);
        for (double latency : latencies_us) {
            out << latency << "\n";
        }

        std::cout << "Saved results to: " << CSV_OUTPUT_PATH << "\n";
    }

    return EXIT_SUCCESS;
}
