#pragma once

// This function exactly tracks the flop count for each function
// Usage
// compile, e.g.:
//      make run-bench FLOPCOUNT=1 N_P=60 N_ITER=4000
//      if not compiled with FLOPCOUNT all associated function are written as ((void)0) and removed when using -O3 flag
// In code:
//      FLOP_RECORD(flop_counter::<FUNC_NAME>, int add, int mult, int div, int sqrt, int exp)
//      <FUNC_NAME> see FunctionId in util.hpp
// Print out metrics:
//      flop_counter::report();
//      flop_counter::report_csv();
// Additonal remarks:
//      make sure to call flop_counter::reset(); after the warmup phase, because else counting flops double
//      make sure to import util.hpp #include "util.hpp"

// get the cycles count per flop
#include "constants.hpp"   // cost_add, cost_mult, cost_div, cost_sqrt, cost_exp

#include <array>
#include <cstdint>
#include <cstdio>

#ifndef FLOPCOUNT
#define FLOPCOUNT 0
#endif

namespace flop_counter {

// defined functions, add new one if needed
// create enumeration, which must be in sync with kFunctionNames
enum FunctionId : int {
  FN_W,
  FN_ACCELERATION,
  FN_PEDESTRIAN_REPULSION,
  FN_OBSTACLE_REPULSION,
  FN_NEAREST_OBSTACLE_POINT,
  FN_UPDATE,
  FN_OTHER,
  FN_COUNT
};

// track the flop count for each function separately
// add more functions, do not mix up the order
inline constexpr std::array<const char *, FN_COUNT> kFunctionNames = {
    "w",
    "acceleration_term_dt",
    "pedestrian_repulsion_dt",
    "obstacle_repulsion_dt",
    "find_nearest_obstacle_point",
    "update_pass",
    "other",
};

#if FLOPCOUNT

// for each function store add, mult, div, sqrt, exp
inline std::uint64_t counters[FN_COUNT][5]{};

// update
inline void record(FunctionId fn, std::uint64_t add, std::uint64_t mult,
                   std::uint64_t div, std::uint64_t sqrt_n,
                   std::uint64_t exp_n) {
  counters[fn][0] += add;
  counters[fn][1] += mult;
  counters[fn][2] += div;
  counters[fn][3] += sqrt_n;
  counters[fn][4] += exp_n;
}

// set all metrics to 0
inline void reset() {
  for (int f = 0; f < FN_COUNT; ++f)
    for (int o = 0; o < 5; ++o)
      counters[f][o] = 0;
}

inline void report() {
  std::printf("\n=== flop count ===\n");
  std::printf("%-30s %12s %12s %12s %12s %12s %16s\n",
              "function", "add", "mult", "div", "sqrt", "exp", "per function");
  long double total = 0;
  for (int f = 0; f < FN_COUNT; ++f) {
    auto a = counters[f][0], m = counters[f][1], d = counters[f][2],
         s = counters[f][3], e = counters[f][4];
    long double w = a * (long double)cost_add  + m * (long double)cost_mult
                  + d * (long double)cost_div  + s * (long double)cost_sqrt
                  + e * (long double)cost_exp;
    total += w;
    if ((a | m | d | s | e) == 0) continue;
    std::printf("%-30s %12llu %12llu %12llu %12llu %12llu %16.2Lf\n",
                kFunctionNames[f],
                (unsigned long long)a, (unsigned long long)m,
                (unsigned long long)d, (unsigned long long)s,
                (unsigned long long)e, w);
  }
  std::printf("%-30s %12s %12s %12s %12s %12s %16.2Lf\n",
              "TOTAL", "", "", "", "", "", total);
}

#define FLOP_RECORD(fn, a, m, d, s, e) \
  ::flop_counter::record((fn), (a), (m), (d), (s), (e))

#else  // FLOPCOUNT == 0 -> no-op, removed by -O3

inline void record(FunctionId, std::uint64_t, std::uint64_t, std::uint64_t,
                   std::uint64_t, std::uint64_t) {}
inline void reset() {}
inline void report() {}

#define FLOP_RECORD(fn, a, m, d, s, e) ((void)0)

#endif

}  // namespace flop_counter