#ifndef STAN_SERVICES_EXPERIMENTAL_ADVI_MEANFIELD_HPP
#define STAN_SERVICES_EXPERIMENTAL_ADVI_MEANFIELD_HPP

#include <stan/callbacks/interrupt.hpp>
#include <stan/callbacks/logger.hpp>
#include <stan/callbacks/writer.hpp>
#include <stan/services/util/experimental_message.hpp>
#include <stan/services/util/initialize.hpp>
#include <stan/services/util/create_rng.hpp>
#include <stan/io/var_context.hpp>
#include <stan/variational/advi.hpp>
#include <boost/random/additive_combine.hpp>
#include <string>
#include <vector>

namespace stan {
namespace services {
namespace experimental {
namespace advi {

/**
 * Runs mean field ADVI.
 *
 * @tparam Model A model implementation
 * @param[in] model Input model to test (with data already instantiated)
 * @param[in] init var context for initialization
 * @param[in] random_seed random seed for the random number generator
 * @param[in] chain chain id to advance the random number generator
 * @param[in] init_radius radius to initialize
 * @param[in] grad_samples number of samples for Monte Carlo estimate
 *   of gradients
 * @param[in] elbo_samples number of samples for Monte Carlo estimate
 *   of ELBO
 * @param[in] max_iterations maximum number of iterations
 * @param[in] tol_rel_obj convergence tolerance on the relative norm
 *   of the objective
 * @param[in] eta stepsize scaling parameter for variational inference
 * @param[in] adapt_engaged adaptation engaged?
 * @param[in] adapt_iterations number of iterations for eta adaptation
 * @param[in] eval_elbo evaluate ELBO every Nth iteration
 * @param[in] output_samples number of posterior samples to draw and
 *   save
 * @param[in,out] interrupt callback to be called every iteration
 * @param[in,out] logger Logger for messages
 * @param[in,out] init_writer Writer callback for unconstrained inits
 * @param[in,out] parameter_writer output for parameter values
 * @param[in,out] diagnostic_writer output for diagnostic values
 * @return error_codes::OK if successful
 */
template <class Model>
int meanfield(Model& model, const stan::io::var_context& init,
              unsigned int random_seed, unsigned int chain, double init_radius,
              int grad_samples, int elbo_samples, int max_iterations,
              double tol_rel_obj, double eta, bool adapt_engaged,
              int adapt_iterations, int eval_elbo, int output_samples,
              callbacks::interrupt& interrupt, callbacks::logger& logger,
              callbacks::writer& init_writer,
              callbacks::writer& parameter_writer,
              callbacks::writer& diagnostic_writer) {
  util::experimental_message(logger);

  boost::ecuyer1988 rng = util::create_rng(random_seed, chain);

  std::vector<int> disc_vector;
  std::vector<double> cont_vector = util::initialize(
      model, init, rng, init_radius, true, logger, init_writer);

  std::vector<std::string> names;
  names.push_back("lp__");
  names.push_back("log_p__");
  names.push_back("log_g__");
  model.constrained_param_names(names, true, true);
  parameter_writer(names);

  Eigen::VectorXd cont_params
      = Eigen::Map<Eigen::VectorXd>(&cont_vector[0], cont_vector.size(), 1);

  stan::variational::advi<Model, stan::variational::normal_meanfield,
                          boost::ecuyer1988>
      cmd_advi(model, cont_params, rng, grad_samples, elbo_samples, eval_elbo,
               output_samples);
  cmd_advi.run(eta, adapt_engaged, adapt_iterations, tol_rel_obj,
               max_iterations, logger, parameter_writer, diagnostic_writer);

  return 0;
}
}  // namespace advi
}  // namespace experimental
}  // namespace services
}  // namespace stan
#endif
