#include <test/test-models/good/mcmc/hmc/common/gauss3D.hpp>
#include <stan/callbacks/stream_logger.hpp>
#include <stan/mcmc/hmc/nuts/unit_e_nuts.hpp>
#include <stan/services/util/create_rng.hpp>
#include <stan/io/empty_var_context.hpp>
#include <fstream>

#include <gtest/gtest.h>

TEST(McmcUnitENuts, build_tree_test) {
  stan::rng_t base_rng = stan::services::util::create_rng(4839294, 0);

  stan::mcmc::unit_e_point z_init(3);
  z_init.q(0) = 1;
  z_init.q(1) = -1;
  z_init.q(2) = 1;
  z_init.p(0) = -1;
  z_init.p(1) = 1;
  z_init.p(2) = -1;

  std::stringstream debug, info, warn, error, fatal;
  stan::callbacks::stream_logger logger(debug, info, warn, error, fatal);

  stan::io::empty_var_context data_var_context;
  gauss3D_model_namespace::gauss3D_model model(data_var_context);

  stan::mcmc::unit_e_nuts<gauss3D_model_namespace::gauss3D_model, stan::rng_t>
      sampler(model, base_rng);

  sampler.z() = z_init;
  sampler.init_hamiltonian(logger);
  sampler.set_nominal_stepsize(0.1);
  sampler.set_stepsize_jitter(0);
  sampler.sample_stepsize();

  stan::mcmc::ps_point z_propose = z_init;

  Eigen::VectorXd p_begin = Eigen::VectorXd::Zero(z_init.p.size());
  Eigen::VectorXd p_sharp_begin = Eigen::VectorXd::Zero(z_init.p.size());
  Eigen::VectorXd p_end = Eigen::VectorXd::Zero(z_init.p.size());
  Eigen::VectorXd p_sharp_end = Eigen::VectorXd::Zero(z_init.p.size());
  Eigen::VectorXd rho = z_init.p;

  double log_sum_weight = -std::numeric_limits<double>::infinity();

  double H0 = -0.1;
  int n_leapfrog = 0;
  double sum_metro_prob = 0;

  bool valid_subtree = sampler.build_tree(
      3, z_propose, p_sharp_begin, p_sharp_end, rho, p_begin, p_end, H0, 1,
      n_leapfrog, log_sum_weight, sum_metro_prob, logger);

  EXPECT_EQ(0.1, sampler.get_nominal_stepsize());

  EXPECT_TRUE(valid_subtree);

  EXPECT_FLOAT_EQ(-11.401228, rho(0));
  EXPECT_FLOAT_EQ(11.401228, rho(1));
  EXPECT_FLOAT_EQ(-11.401228, rho(2));

  EXPECT_FLOAT_EQ(-1.09475, p_begin(0));
  EXPECT_FLOAT_EQ(1.09475, p_begin(1));
  EXPECT_FLOAT_EQ(-1.09475, p_begin(2));

  EXPECT_FLOAT_EQ(-1.09475, p_sharp_begin(0));
  EXPECT_FLOAT_EQ(1.09475, p_sharp_begin(1));
  EXPECT_FLOAT_EQ(-1.09475, p_sharp_begin(2));

  EXPECT_FLOAT_EQ(-1.4131583, p_end(0));
  EXPECT_FLOAT_EQ(1.4131583, p_end(1));
  EXPECT_FLOAT_EQ(-1.4131583, p_end(2));

  EXPECT_FLOAT_EQ(-1.4131583, p_sharp_end(0));
  EXPECT_FLOAT_EQ(1.4131583, p_sharp_end(1));
  EXPECT_FLOAT_EQ(-1.4131583, p_sharp_end(2));

  EXPECT_FLOAT_EQ(-0.022019938, sampler.z().q(0));
  EXPECT_FLOAT_EQ(0.022019938, sampler.z().q(1));
  EXPECT_FLOAT_EQ(-0.022019938, sampler.z().q(2));

  EXPECT_FLOAT_EQ(-1.4131583, sampler.z().p(0));
  EXPECT_FLOAT_EQ(1.4131583, sampler.z().p(1));
  EXPECT_FLOAT_EQ(-1.4131583, sampler.z().p(2));

  EXPECT_EQ(8, n_leapfrog);
  EXPECT_FLOAT_EQ(std::log(0.36134657), log_sum_weight);
  EXPECT_FLOAT_EQ(0.36134657, sum_metro_prob);

  EXPECT_EQ("", debug.str());
  EXPECT_EQ("", info.str());
  EXPECT_EQ("", warn.str());
  EXPECT_EQ("", error.str());
  EXPECT_EQ("", fatal.str());
}

TEST(McmcUnitENuts, tree_boundary_test) {
  stan::rng_t base_rng = stan::services::util::create_rng(4839294, 0);

  stan::mcmc::unit_e_point z_init(3);
  z_init.q(0) = 1;
  z_init.q(1) = -1;
  z_init.q(2) = 1;
  z_init.p(0) = -1;
  z_init.p(1) = 1;
  z_init.p(2) = -1;

  std::stringstream debug, info, warn, error, fatal;
  stan::callbacks::stream_logger logger(debug, info, warn, error, fatal);

  stan::io::empty_var_context data_var_context;

  typedef gauss3D_model_namespace::gauss3D_model model_t;
  model_t model(data_var_context);

  // Compute expected tree boundaries
  typedef stan::mcmc::unit_e_metric<model_t, stan::rng_t> metric_t;
  metric_t metric(model);

  stan::mcmc::expl_leapfrog<metric_t> unit_e_integrator;
  double epsilon = 0.1;

  stan::mcmc::unit_e_point z_test = z_init;
  metric.init(z_test, logger);

  unit_e_integrator.evolve(z_test, metric, epsilon, logger);
  Eigen::VectorXd p_forward_1 = z_test.p;
  Eigen::VectorXd p_sharp_forward_1 = metric.dtau_dp(z_test);

  unit_e_integrator.evolve(z_test, metric, epsilon, logger);
  Eigen::VectorXd p_forward_2 = z_test.p;
  Eigen::VectorXd p_sharp_forward_2 = metric.dtau_dp(z_test);

  unit_e_integrator.evolve(z_test, metric, epsilon, logger);
  unit_e_integrator.evolve(z_test, metric, epsilon, logger);
  Eigen::VectorXd p_forward_3 = z_test.p;
  Eigen::VectorXd p_sharp_forward_3 = metric.dtau_dp(z_test);

  unit_e_integrator.evolve(z_test, metric, epsilon, logger);
  unit_e_integrator.evolve(z_test, metric, epsilon, logger);
  unit_e_integrator.evolve(z_test, metric, epsilon, logger);
  unit_e_integrator.evolve(z_test, metric, epsilon, logger);
  Eigen::VectorXd p_forward_4 = z_test.p;
  Eigen::VectorXd p_sharp_forward_4 = metric.dtau_dp(z_test);

  z_test = z_init;
  metric.init(z_test, logger);

  unit_e_integrator.evolve(z_test, metric, -epsilon, logger);
  Eigen::VectorXd p_backward_1 = z_test.p;
  Eigen::VectorXd p_sharp_backward_1 = metric.dtau_dp(z_test);

  unit_e_integrator.evolve(z_test, metric, -epsilon, logger);
  Eigen::VectorXd p_backward_2 = z_test.p;
  Eigen::VectorXd p_sharp_backward_2 = metric.dtau_dp(z_test);

  unit_e_integrator.evolve(z_test, metric, -epsilon, logger);
  unit_e_integrator.evolve(z_test, metric, -epsilon, logger);
  Eigen::VectorXd p_backward_3 = z_test.p;
  Eigen::VectorXd p_sharp_backward_3 = metric.dtau_dp(z_test);

  unit_e_integrator.evolve(z_test, metric, -epsilon, logger);
  unit_e_integrator.evolve(z_test, metric, -epsilon, logger);
  unit_e_integrator.evolve(z_test, metric, -epsilon, logger);
  unit_e_integrator.evolve(z_test, metric, -epsilon, logger);
  Eigen::VectorXd p_backward_4 = z_test.p;
  Eigen::VectorXd p_sharp_backward_4 = metric.dtau_dp(z_test);

  // Check expected tree boundaries to those dynamically generated by NUTS
  stan::mcmc::unit_e_nuts<model_t, stan::rng_t> sampler(model, base_rng);

  sampler.set_nominal_stepsize(epsilon);
  sampler.set_stepsize_jitter(0);
  sampler.sample_stepsize();

  stan::mcmc::ps_point z_propose = z_init;

  Eigen::VectorXd p_begin = Eigen::VectorXd::Zero(z_init.p.size());
  Eigen::VectorXd p_sharp_begin = Eigen::VectorXd::Zero(z_init.p.size());
  Eigen::VectorXd p_end = Eigen::VectorXd::Zero(z_init.p.size());
  Eigen::VectorXd p_sharp_end = Eigen::VectorXd::Zero(z_init.p.size());
  Eigen::VectorXd rho = z_init.p;

  double log_sum_weight = -std::numeric_limits<double>::infinity();

  double H0 = -0.1;
  int n_leapfrog = 0;
  double sum_metro_prob = 0;

  // Depth 0 forward
  sampler.z() = z_init;
  sampler.init_hamiltonian(logger);
  sampler.build_tree(0, z_propose, p_sharp_begin, p_sharp_end, rho, p_begin,
                     p_end, H0, 1, n_leapfrog, log_sum_weight, sum_metro_prob,
                     logger);

  for (int n = 0; n < rho.size(); ++n) {
    EXPECT_FLOAT_EQ(p_forward_1(n), p_begin(n));
    EXPECT_FLOAT_EQ(p_sharp_forward_1(n), p_sharp_begin(n));

    EXPECT_FLOAT_EQ(p_forward_1(n), p_end(n));
    EXPECT_FLOAT_EQ(p_sharp_forward_1(n), p_sharp_end(n));
  }

  // Depth 1 forward
  sampler.z() = z_init;
  sampler.init_hamiltonian(logger);
  sampler.build_tree(1, z_propose, p_sharp_begin, p_sharp_end, rho, p_begin,
                     p_end, H0, 1, n_leapfrog, log_sum_weight, sum_metro_prob,
                     logger);

  for (int n = 0; n < rho.size(); ++n) {
    EXPECT_FLOAT_EQ(p_forward_1(n), p_begin(n));
    EXPECT_FLOAT_EQ(p_sharp_forward_1(n), p_sharp_begin(n));

    EXPECT_FLOAT_EQ(p_forward_2(n), p_end(n));
    EXPECT_FLOAT_EQ(p_sharp_forward_2(n), p_sharp_end(n));
  }

  // Depth 2 forward
  sampler.z() = z_init;
  sampler.init_hamiltonian(logger);
  sampler.build_tree(2, z_propose, p_sharp_begin, p_sharp_end, rho, p_begin,
                     p_end, H0, 1, n_leapfrog, log_sum_weight, sum_metro_prob,
                     logger);

  for (int n = 0; n < rho.size(); ++n) {
    EXPECT_FLOAT_EQ(p_forward_1(n), p_begin(n));
    EXPECT_FLOAT_EQ(p_sharp_forward_1(n), p_sharp_begin(n));

    EXPECT_FLOAT_EQ(p_forward_3(n), p_end(n));
    EXPECT_FLOAT_EQ(p_sharp_forward_3(n), p_sharp_end(n));
  }

  // Depth 3 forward
  sampler.z() = z_init;
  sampler.init_hamiltonian(logger);
  sampler.build_tree(3, z_propose, p_sharp_begin, p_sharp_end, rho, p_begin,
                     p_end, H0, 1, n_leapfrog, log_sum_weight, sum_metro_prob,
                     logger);

  for (int n = 0; n < rho.size(); ++n) {
    EXPECT_FLOAT_EQ(p_forward_1(n), p_begin(n));
    EXPECT_FLOAT_EQ(p_sharp_forward_1(n), p_sharp_begin(n));

    EXPECT_FLOAT_EQ(p_forward_4(n), p_end(n));
    EXPECT_FLOAT_EQ(p_sharp_forward_4(n), p_sharp_end(n));
  }

  // Depth 0 backward
  sampler.z() = z_init;
  sampler.init_hamiltonian(logger);
  sampler.build_tree(0, z_propose, p_sharp_begin, p_sharp_end, rho, p_begin,
                     p_end, H0, -1, n_leapfrog, log_sum_weight, sum_metro_prob,
                     logger);

  for (int n = 0; n < rho.size(); ++n) {
    EXPECT_FLOAT_EQ(p_backward_1(n), p_begin(n));
    EXPECT_FLOAT_EQ(p_sharp_backward_1(n), p_sharp_begin(n));

    EXPECT_FLOAT_EQ(p_backward_1(n), p_end(n));
    EXPECT_FLOAT_EQ(p_sharp_backward_1(n), p_sharp_end(n));
  }

  // Depth 1 backward
  sampler.z() = z_init;
  sampler.init_hamiltonian(logger);
  sampler.build_tree(1, z_propose, p_sharp_begin, p_sharp_end, rho, p_begin,
                     p_end, H0, -1, n_leapfrog, log_sum_weight, sum_metro_prob,
                     logger);

  for (int n = 0; n < rho.size(); ++n) {
    EXPECT_FLOAT_EQ(p_backward_1(n), p_begin(n));
    EXPECT_FLOAT_EQ(p_sharp_backward_1(n), p_sharp_begin(n));

    EXPECT_FLOAT_EQ(p_backward_2(n), p_end(n));
    EXPECT_FLOAT_EQ(p_sharp_backward_2(n), p_sharp_end(n));
  }

  // Depth 2 backward
  sampler.z() = z_init;
  sampler.init_hamiltonian(logger);
  sampler.build_tree(2, z_propose, p_sharp_begin, p_sharp_end, rho, p_begin,
                     p_end, H0, -1, n_leapfrog, log_sum_weight, sum_metro_prob,
                     logger);

  for (int n = 0; n < rho.size(); ++n) {
    EXPECT_FLOAT_EQ(p_backward_1(n), p_begin(n));
    EXPECT_FLOAT_EQ(p_sharp_backward_1(n), p_sharp_begin(n));

    EXPECT_FLOAT_EQ(p_backward_3(n), p_end(n));
    EXPECT_FLOAT_EQ(p_sharp_backward_3(n), p_sharp_end(n));
  }

  // Depth 3 backward
  sampler.z() = z_init;
  sampler.init_hamiltonian(logger);
  sampler.build_tree(3, z_propose, p_sharp_begin, p_sharp_end, rho, p_begin,
                     p_end, H0, -1, n_leapfrog, log_sum_weight, sum_metro_prob,
                     logger);

  for (int n = 0; n < rho.size(); ++n) {
    EXPECT_FLOAT_EQ(p_backward_1(n), p_begin(n));
    EXPECT_FLOAT_EQ(p_sharp_backward_1(n), p_sharp_begin(n));

    EXPECT_FLOAT_EQ(p_backward_4(n), p_end(n));
    EXPECT_FLOAT_EQ(p_sharp_backward_4(n), p_sharp_end(n));
  }
}

TEST(McmcUnitENuts, transition_test) {
  stan::rng_t base_rng = stan::services::util::create_rng(4839294, 0);

  stan::mcmc::unit_e_point z_init(3);
  z_init.q(0) = 1;
  z_init.q(1) = -1;
  z_init.q(2) = 1;
  z_init.p(0) = -1;
  z_init.p(1) = 1;
  z_init.p(2) = -1;

  std::stringstream debug, info, warn, error, fatal;
  stan::callbacks::stream_logger logger(debug, info, warn, error, fatal);

  stan::io::empty_var_context data_var_context;
  gauss3D_model_namespace::gauss3D_model model(data_var_context);

  stan::mcmc::unit_e_nuts<gauss3D_model_namespace::gauss3D_model, stan::rng_t>
      sampler(model, base_rng);

  sampler.z() = z_init;
  sampler.init_hamiltonian(logger);
  sampler.set_nominal_stepsize(0.1);
  sampler.set_stepsize_jitter(0);
  sampler.sample_stepsize();

  stan::mcmc::sample init_sample(z_init.q, 0, 0);

  stan::mcmc::sample s = sampler.transition(init_sample, logger);

  EXPECT_EQ(3, sampler.depth_);
  EXPECT_EQ((2 << 3) - 1, sampler.n_leapfrog_);
  EXPECT_FALSE(sampler.divergent_);

  EXPECT_FLOAT_EQ(0.70149082, s.cont_params()(0));
  EXPECT_FLOAT_EQ(-0.69831347, s.cont_params()(1));
  EXPECT_FLOAT_EQ(0.54392564, s.cont_params()(2));
  EXPECT_FLOAT_EQ(-0.63779306, s.log_prob());
  EXPECT_FLOAT_EQ(0.99912512, s.accept_stat());
  EXPECT_EQ("", debug.str());
  EXPECT_EQ("", info.str());
  EXPECT_EQ("", warn.str());
  EXPECT_EQ("", error.str());
  EXPECT_EQ("", fatal.str());
}
