#ifndef FESA_ANALYSIS_HHT_DEVELOPMENT_PROBE_H_
#define FESA_ANALYSIS_HHT_DEVELOPMENT_PROBE_H_

// This header is compiled out of normal builds, including environment/file I/O.
#if defined(FESA_HHT_DEVELOPMENT_PROBE)
#include <array>
#include <atomic>
#include <chrono>
#include <cstdint>
#include <cstdlib>
#include <fstream>
#include <string>

#include "fesa/solvers/linear/linear_solver.h"
#include "fesa/solvers/linear/tangent_solver.h"

namespace fesa::hht_development_probe {
enum Stage { kProcedure, kInner, kAssembly, kOuterSolve, kRates,
             kFactor, kSolve, kStageCount };
enum Counter { kOuterCorrection, kInnerCorrection, kInternalEvaluation,
               kInertiaEvaluation, kFactorCall, kSolveCall, kRhs,
               kConstraintInternalEvaluation, kInternalFailure, kCounterCount };
struct Statistics {
  std::array<std::atomic<std::uint64_t>, kCounterCount> counts{};
  // Timers/stage attribution are serial procedure-thread data. Only the three
  // element counters below and internal failures may be touched by workers.
  std::array<std::int64_t, kStageCount> inclusive_ns{};
  std::array<std::int64_t, kStageCount> exclusive_ns{};
  std::array<std::uint64_t, kStageCount> stage_calls{};
  std::array<std::array<std::uint64_t, kCounterCount>, kStageCount> by_stage{};
};
inline Statistics* active = nullptr;
class Scope {
 public:
  explicit Scope(Stage stage) : stage_(stage), parent_(current_),
      start_(std::chrono::steady_clock::now()) {
    current_ = this;
    if (active != nullptr) ++active->stage_calls[stage_];
  }
  ~Scope() {
    const auto elapsed = std::chrono::duration_cast<std::chrono::nanoseconds>(
        std::chrono::steady_clock::now() - start_).count();
    if (active != nullptr) {
      active->inclusive_ns[stage_] += elapsed;
      active->exclusive_ns[stage_] += elapsed - children_;
    }
    if (parent_ != nullptr) parent_->children_ += elapsed;
    current_ = parent_;
  }
  static Stage Current() { return current_ == nullptr ? kProcedure : current_->stage_; }
 private:
  Stage stage_;
  Scope* parent_;
  std::chrono::steady_clock::time_point start_;
  std::int64_t children_ = 0;
  inline static thread_local Scope* current_ = nullptr;
};
inline void Count(Counter counter, std::uint64_t count = 1U) {
  if (active == nullptr) return;
  active->counts[counter].fetch_add(count, std::memory_order_relaxed);
  // Element callbacks are parallel; their totals use only the atomics above.
  if (counter != kInternalEvaluation && counter != kInertiaEvaluation &&
      counter != kConstraintInternalEvaluation && counter != kInternalFailure)
    active->by_stage[Scope::Current()][counter] += count;
}
class Session {
 public:
  Session() : previous_(active) { active = &statistics_; }
  ~Session() {
    active = previous_;
#if defined(_WIN32)
    char* environment = nullptr;
    std::size_t size = 0U;
    if (_dupenv_s(&environment, &size, "FESA_HHT_PROBE_OUTPUT") != 0) return;
    const std::string path = environment == nullptr ? "" : environment;
    std::free(environment);
#else
    const char* environment = std::getenv("FESA_HHT_PROBE_OUTPUT");
    const std::string path = environment == nullptr ? "" : environment;
#endif
    if (path.empty()) return;
    std::ofstream output(path, std::ios::out | std::ios::binary);
    constexpr const char* counters[] = {"outer_applied_correction", "inner_applied_correction",
        "assembly_internal_evaluation", "assembly_inertia_evaluation",
        "logical_factor_call", "logical_solve_call", "logical_rhs",
        "constraint_internal_evaluation", "assembly_internal_failure"};
    constexpr const char* stages[] = {"procedure_other", "inner_graph", "assembly",
        "outer_linear", "physical_rates", "factor", "solve"};
    for (int c = 0; c < kCounterCount; ++c)
      output << counters[c] << '=' << statistics_.counts[c].load() << '\n';
    output << "procedure_inclusive_ns=" << statistics_.inclusive_ns[kProcedure] << '\n';
    for (int s = 0; s < kStageCount; ++s) {
      output << stages[s] << "_calls=" << statistics_.stage_calls[s] << '\n';
      output << stages[s] << "_exclusive_ns=" << statistics_.exclusive_ns[s] << '\n';
      for (int c = 0; c < kCounterCount; ++c)
        if (statistics_.by_stage[s][c] != 0U)
          output << stages[s] << '_' << counters[c] << '='
                 << statistics_.by_stage[s][c] << '\n';
    }
  }
 private:
  Statistics statistics_;
  Statistics* previous_;
};
class PositiveSolver final : public LinearSolver {
 public:
  explicit PositiveSolver(LinearSolver& solver) : solver_(solver) {}
  Status Factorize(const SparseMatrix& matrix) override {
    Count(kFactorCall); Scope scope{kFactor}; return solver_.Factorize(matrix);
  }
  Status Solve(const Vector& rhs, Vector& solution) const override {
    Count(kSolveCall); Count(kRhs); Scope scope{kSolve};
    return solver_.Solve(rhs, solution);
  }
 private:
  LinearSolver& solver_;
};
class CorrectionSolver final : public TangentSolver {
 public:
  explicit CorrectionSolver(TangentSolver& solver) : solver_(solver) {}
  Status FactorizeFresh(const SparseMatrix& matrix) override {
    Count(kFactorCall); Scope scope{kFactor}; return solver_.FactorizeFresh(matrix);
  }
  Status SolveCorrection(const Vector& rhs, Vector& solution) const override {
    Count(kSolveCall); Count(kRhs); Scope scope{kSolve};
    return solver_.SolveCorrection(rhs, solution);
  }
 private:
  TangentSolver& solver_;
};
}  // namespace fesa::hht_development_probe
#define FESA_HHT_PROBE_COUNT(counter) \
  ::fesa::hht_development_probe::Count(::fesa::hht_development_probe::counter)
#define FESA_HHT_PROBE_SCOPE(stage) \
  ::fesa::hht_development_probe::Scope probe_scope{ \
      ::fesa::hht_development_probe::stage}
#else
#define FESA_HHT_PROBE_COUNT(counter) ((void)0)
#define FESA_HHT_PROBE_SCOPE(stage) ((void)0)
#endif
#endif  // FESA_ANALYSIS_HHT_DEVELOPMENT_PROBE_H_