#include <stan/services/sample/hmc_nuts_diag_e_adapt.hpp>
#include <stan/callbacks/json_writer.hpp>
#include <stan/callbacks/unique_stream_writer.hpp>
#include <stan/callbacks/structured_writer.hpp>
#include <stan/io/empty_var_context.hpp>
#include <src/test/unit/services/util.hpp>
#include <test/test-models/good/optimization/rosenbrock.hpp>
#include <test/unit/services/instrumented_callbacks.hpp>
#include <test/unit/util.hpp>
#include <gtest/gtest.h>
#include <iostream>

auto&& blah = stan::math::init_threadpool_tbb();

static constexpr size_t num_chains = 4;

struct deleter_noop {
  template <typename T>
  constexpr void operator()(T* arg) const {}
};
class ServicesSampleHmcNutsDiagEAdaptParMatch : public testing::Test {
 public:
  ServicesSampleHmcNutsDiagEAdaptParMatch()
      : ss_par(num_chains),
        ss_seq(num_chains),
        ss_metric(num_chains),
        model(std::make_unique<rosenbrock_model_namespace::rosenbrock_model>(
            data_context, 0, &model_log)) {
    for (int i = 0; i < num_chains; ++i) {
      init.push_back(stan::test::unit::instrumented_writer{});
      par_parameters.emplace_back(
          std::unique_ptr<std::stringstream, deleter_noop>(&ss_par[i]), "#");
      seq_parameters.emplace_back(
          std::unique_ptr<std::stringstream, deleter_noop>(&ss_seq[i]), "#");
      metrics.emplace_back(
          stan::callbacks::json_writer<std::stringstream, deleter_noop>(
              std::unique_ptr<std::stringstream, deleter_noop>(&ss_metric[i])));
      diagnostics.push_back(stan::test::unit::instrumented_writer{});
      context.push_back(std::make_shared<stan::io::empty_var_context>());
    }
  }

  void SetUp() {
    for (int i = 0; i < num_chains; ++i) {
      ss_metric[i].str(std::string());
      ss_metric[i].clear();
    }
  }

  stan::io::empty_var_context data_context;
  std::stringstream model_log;
  stan::test::unit::instrumented_logger logger;
  std::vector<stan::test::unit::instrumented_writer> init;
  std::vector<std::stringstream> ss_par;
  std::vector<std::stringstream> ss_seq;
  std::vector<std::stringstream> ss_metric;
  using str_writer
      = stan::callbacks::unique_stream_writer<std::stringstream, deleter_noop>;
  std::vector<str_writer> par_parameters;
  std::vector<str_writer> seq_parameters;
  std::vector<stan::callbacks::json_writer<std::stringstream, deleter_noop>>
      metrics;
  std::vector<stan::test::unit::instrumented_writer> diagnostics;
  std::vector<std::shared_ptr<stan::io::empty_var_context>> context;
  std::unique_ptr<rosenbrock_model_namespace::rosenbrock_model> model;
};

/**
 * This test checks that running multiple chains in one call
 * with the same initial id is the same as running multiple calls
 * with incrementing chain ids.
 */
TEST_F(ServicesSampleHmcNutsDiagEAdaptParMatch, single_multi_match) {
  constexpr unsigned int random_seed = 0;
  constexpr unsigned int chain = 0;
  constexpr double init_radius = 0;
  constexpr int num_warmup = 200;
  constexpr int num_samples = 400;
  constexpr int num_thin = 5;
  constexpr bool save_warmup = true;
  constexpr int refresh = 0;
  constexpr double stepsize = 0.1;
  constexpr double stepsize_jitter = 0;
  constexpr int max_depth = 8;
  constexpr double delta = .1;
  constexpr double gamma = .1;
  constexpr double kappa = .1;
  constexpr double t0 = .1;
  constexpr unsigned int init_buffer = 50;
  constexpr unsigned int term_buffer = 50;
  constexpr unsigned int window = 100;
  stan::test::unit::instrumented_interrupt interrupt;
  EXPECT_EQ(interrupt.call_count(), 0);

  int return_code = stan::services::sample::hmc_nuts_diag_e_adapt(
      *model, num_chains, context, random_seed, chain, init_radius, num_warmup,
      num_samples, num_thin, save_warmup, refresh, stepsize, stepsize_jitter,
      max_depth, delta, gamma, kappa, t0, init_buffer, term_buffer, window,
      interrupt, logger, init, par_parameters, diagnostics, metrics);
  EXPECT_EQ(0, return_code);

  int num_output_lines = (num_warmup + num_samples) / num_thin;
  EXPECT_EQ((num_warmup + num_samples) * num_chains, interrupt.call_count());
  for (int i = 0; i < num_chains; ++i) {
    stan::test::unit::instrumented_writer seq_init;
    stan::test::unit::instrumented_writer seq_diagnostic;
    return_code = stan::services::sample::hmc_nuts_diag_e_adapt(
        *model, *(context[i]), random_seed, i, init_radius, num_warmup,
        num_samples, num_thin, save_warmup, refresh, stepsize, stepsize_jitter,
        max_depth, delta, gamma, kappa, t0, init_buffer, term_buffer, window,
        interrupt, logger, seq_init, seq_parameters[i], seq_diagnostic);
    EXPECT_EQ(0, return_code);
  }

  std::vector<Eigen::MatrixXd> par_res;
  std::vector<std::string> par_metrics;
  for (int i = 0; i < num_chains; ++i) {
    auto par_str = par_parameters[i].get_stream().str();
    auto sub_par_str = par_str.substr(par_str.find("Diagonal") - 1);
    std::istringstream sub_par_stream(sub_par_str);
    Eigen::MatrixXd par_mat
        = stan::test::read_stan_sample_csv(sub_par_stream, 80, 9);
    par_res.push_back(par_mat);
    par_metrics.push_back(ss_metric[i].str());
    ASSERT_TRUE(stan::test::is_valid_JSON(par_metrics[i]));
    EXPECT_EQ(count_matches("stepsize", par_metrics[i]), 1);
    EXPECT_EQ(count_matches("metric_type", par_metrics[i]), 1);
    EXPECT_EQ(count_matches("inv_metric", par_metrics[i]), 1);
    EXPECT_EQ(count_matches("[", par_metrics[i]), 1);  // single list
  }

  std::vector<Eigen::MatrixXd> seq_res;
  std::vector<std::string> seq_metrics;
  for (int i = 0; i < num_chains; ++i) {
    auto seq_str = seq_parameters[i].get_stream().str();
    auto sub_seq_str = seq_str.substr(seq_str.find("Diagonal") - 1);
    std::istringstream sub_seq_stream(sub_seq_str);
    Eigen::MatrixXd seq_mat
        = stan::test::read_stan_sample_csv(sub_seq_stream, 80, 9);
    seq_res.push_back(seq_mat);
    seq_metrics.push_back(ss_metric[i].str());
  }
  for (int i = 0; i < num_chains; ++i) {
    Eigen::MatrixXd diff_res
        = (par_res[i].array() - seq_res[i].array()).matrix();
    EXPECT_MATRIX_EQ(diff_res, Eigen::MatrixXd::Zero(80, 9));
    EXPECT_EQ(par_metrics[i], seq_metrics[i]);
  }
}
