#include <chrono>
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <iostream>
#include <string>
#include <vector>

extern "C" {
void *rlgym_bpu_create(const char *model_path, int bpu_core, int priority,
                       char *err, int err_len);
int rlgym_bpu_infer(void *handle, const float *input, float *output, char *err,
                    int err_len);
int rlgym_bpu_input_floats(void *handle);
int rlgym_bpu_output_floats(void *handle);
void rlgym_bpu_destroy(void *handle);
}

namespace {

constexpr float kReference[12] = {
    0.537439f,  1.032869f,  0.042010f, -0.463015f,
    0.515008f,  1.349627f,  0.014652f, -1.518314f,
    1.012123f, -0.346840f,  0.086353f, -0.755527f,
};

bool read_f32_file(const std::string &path, std::vector<float> *data) {
  std::ifstream ifs(path, std::ios::binary);
  if (!ifs) {
    return false;
  }
  ifs.seekg(0, std::ios::end);
  const auto size = ifs.tellg();
  ifs.seekg(0, std::ios::beg);
  if (size <= 0 || size % static_cast<std::streamoff>(sizeof(float)) != 0) {
    return false;
  }
  data->resize(static_cast<size_t>(size) / sizeof(float));
  ifs.read(reinterpret_cast<char *>(data->data()), size);
  return ifs.good();
}

}  // namespace

int main(int argc, char **argv) {
  const char *model =
      "/root/go1_pro_deploy/deploy_45dim_rl_gym/bpu_quantization/mapper_output_6500_gemm/"
      "policy_robotlab_6500_int16_gemm.bin";
  const char *input =
      "/root/go1_pro_deploy/deploy_45dim_rl_gym/bpu_quantization/calibration_data/00000.bin";
  int repeat = 1000;
  if (argc > 1) model = argv[1];
  if (argc > 2) input = argv[2];
  if (argc > 3) repeat = std::atoi(argv[3]);

  char err[1024] = {};
  void *handle = rlgym_bpu_create(model, 1, 0, err, sizeof(err));
  if (handle == nullptr) {
    std::cerr << err << "\n";
    return 3;
  }

  const int input_floats = rlgym_bpu_input_floats(handle);
  const int output_floats = rlgym_bpu_output_floats(handle);
  if (input_floats <= 0 || output_floats <= 0) {
    std::cerr << "invalid tensor sizes from BPU runtime\n";
    rlgym_bpu_destroy(handle);
    return 2;
  }

  std::vector<float> obs;
  if (!read_f32_file(input, &obs) || static_cast<int>(obs.size()) != input_floats) {
    std::cerr << "failed to read " << input_floats << " float32 input: " << input << "\n";
    rlgym_bpu_destroy(handle);
    return 2;
  }

  std::vector<float> out(static_cast<size_t>(output_floats), 0.0f);
  if (rlgym_bpu_infer(handle, obs.data(), out.data(), err, sizeof(err)) != 0) {
    std::cerr << err << "\n";
    rlgym_bpu_destroy(handle);
    return 4;
  }
  float max_diff = 0.0f;
  if (input_floats == 450 && output_floats == 12) {
    for (int i = 0; i < 12; ++i) {
      max_diff = std::max(max_diff, std::fabs(out[static_cast<size_t>(i)] - kReference[i]));
    }
  }

  const auto t0 = std::chrono::steady_clock::now();
  for (int i = 0; i < repeat; ++i) {
    if (rlgym_bpu_infer(handle, obs.data(), out.data(), err, sizeof(err)) != 0) {
      std::cerr << err << "\n";
      rlgym_bpu_destroy(handle);
      return 5;
    }
  }
  const auto t1 = std::chrono::steady_clock::now();
  const double elapsed_ms =
      std::chrono::duration<double, std::milli>(t1 - t0).count();

  std::cout << "input_floats=" << input_floats << " output_floats=" << output_floats << "\n";
  if (input_floats == 450 && output_floats == 12) {
    std::cout << "reference_max_abs_diff " << max_diff << "\n";
  }
  std::cout << "repeat=" << repeat << " cpp_avg_ms=" << (elapsed_ms / repeat) << "\n";
  std::cout << "action0=" << out[0] << " action_max_abs="
            << *std::max_element(out.begin(), out.end(), [](float a, float b) {
                 return std::fabs(a) < std::fabs(b);
               })
            << "\n";

  rlgym_bpu_destroy(handle);
  return 0;
}
