// Copyright (c) by respective owners including Yahoo!, Microsoft, and
// individual contributors. All rights reserved. Released under a BSD (revised)
// license as described in the file LICENSE.

// Notes:
// Implements importance weighted active learning.  Source: https://arxiv.org/pdf/0812.4952.pdf

#include "vw/core/reductions/active.h"

#include "vw/common/random.h"
#include "vw/common/string_view.h"
#include "vw/common/vw_exception.h"
#include "vw/config/options.h"
#include "vw/core/learner.h"
#include "vw/core/model_utils.h"
#include "vw/core/setup_base.h"
#include "vw/core/shared_data.h"
#include "vw/core/simple_label.h"
#include "vw/core/vw.h"
#include "vw/core/vw_math.h"
#include "vw/core/vw_versions.h"
#include "vw/io/errno_handling.h"
#include "vw/io/logger.h"

#include <cerrno>
#include <cfloat>
#include <cmath>

using namespace VW::LEARNER;
using namespace VW::config;
using namespace VW::reductions;
namespace
{
float get_active_coin_bias(float example_count, float avg_loss, float alt_label_error_rate_diff, float mellowness)
{  // implementation follows
   // https://web.archive.org/web/20120525164352/http://books.nips.cc/papers/files/nips23/NIPS2010_0363.pdf
  const float mellow_log_e_count_over_e_count =
      mellowness * (std::log(example_count + 1.f) + 0.0001f) / (example_count + 0.0001f);
  const float sqrt_mellow_lecoec = std::sqrt(mellow_log_e_count_over_e_count);
  // loss should be in [0,1]
  avg_loss = VW::math::clamp(avg_loss, 0.f, 1.f);

  const float sqrt_avg_loss_plus_sqrt_alt_loss =
      std::min(1.f,  // std::sqrt(avg_loss) + // commented out because two square roots appears to conservative.
          std::sqrt(avg_loss + alt_label_error_rate_diff));  // emperical variance deflater.
  // std::cout << "example_count = " << example_count << " avg_loss = " << avg_loss << " alt_label_error_rate_diff = "
  // << alt_label_error_rate_diff << " mellowness = " << mellowness << " mlecoc = " << mellow_log_e_count_over_e_count
  //	    << " sqrt_mellow_lecoec = " << sqrt_mellow_lecoec << " double sqrt = " << sqrt_avg_loss_plus_sqrt_alt_loss
  //<< std::endl;

  if (alt_label_error_rate_diff <= sqrt_mellow_lecoec * sqrt_avg_loss_plus_sqrt_alt_loss  // deflater in use.
          + mellow_log_e_count_over_e_count)
  {
    return 1;
  }
  // old equation
  //   const float rs = (sqrt_avg_loss_plus_sqrt_alt_loss + std::sqrt(sqrt_avg_loss_plus_sqrt_alt_loss *
  //   sqrt_avg_loss_plus_sqrt_alt_loss + 4 * alt_label_error_rate_diff)) / (2 * alt_label_error_rate_diff); return
  //   mellow_log_e_count_over_e_count * rs * rs;
  const float sqrt_s = (sqrt_mellow_lecoec +
                           std::sqrt(mellow_log_e_count_over_e_count +
                               4 * alt_label_error_rate_diff * mellow_log_e_count_over_e_count)) /
      2 * alt_label_error_rate_diff;
  //  std::cout << "sqrt_s = " << sqrt_s << std::endl;
  return sqrt_s * sqrt_s;
}

float query_decision(const active& a, float updates_to_change_prediction, float example_count)
{
  float bias;
  if (example_count <= 1.f) { bias = 1.f; }
  else
  {
    //    const auto weighted_queries = static_cast<float>(a._shared_data->weighted_labeled_examples);
    const float avg_loss = (static_cast<float>(a._shared_data->sum_loss) / example_count);
    //+ std::sqrt((1.f + 0.5f * std::log(example_count)) / (weighted_queries + 0.0001f));  Commented this out, not
    // following why we need it from the theory.
    //    std::cout << "avg_loss = " << avg_loss << " weighted_queries = " << weighted_queries << " sum_loss = " <<
    //    a._shared_data->sum_loss << " example_count = " << example_count << std::endl;
    bias = get_active_coin_bias(example_count, avg_loss, updates_to_change_prediction / example_count, a.active_c0);
  }
  //  std::cout << "bias = " << bias << std::endl;
  return (a._random_state->get_and_update_random() < bias) ? 1.f / bias : -1.f;
}

template <bool is_learn>
void predict_or_learn_simulation(active& a, learner& base, VW::example& ec)
{
  base.predict(ec);

  if (is_learn)
  {
    const auto k = static_cast<float>(a._shared_data->t);
    constexpr float threshold = 0.f;

    ec.confidence = fabsf(ec.pred.scalar - threshold) / base.sensitivity(ec);
    const float importance = query_decision(a, ec.confidence, k);

    if (importance > 0.f)
    {
      a._shared_data->queries += 1;
      ec.weight *= importance;
      base.learn(ec);
    }
    else
    {
      ec.l.simple.label = FLT_MAX;
      ec.weight = 0.f;
    }
  }
}

template <bool is_learn>
void predict_or_learn_active(active& a, learner& base, VW::example& ec)
{
  if (is_learn) { base.learn(ec); }
  else { base.predict(ec); }

  if (ec.l.simple.label == FLT_MAX)
  {
    const float threshold = (a._shared_data->max_label + a._shared_data->min_label) * 0.5f;
    // We want to understand the change in prediction if the label were to be
    // the opposite of what was predicted. 0 and 1 are used for the expected min
    // and max labels to be coming in from the active interactor.
    ec.l.simple.label = (ec.pred.scalar >= threshold) ? a._min_seen_label : a._max_seen_label;
    ec.confidence = std::abs(ec.pred.scalar - threshold) / base.sensitivity(ec);
    ec.l.simple.label = FLT_MAX;
  }
  else
  {
    // Update seen labels based on the current example's label.
    a._min_seen_label = std::min(ec.l.simple.label, a._min_seen_label);
    a._max_seen_label = std::max(ec.l.simple.label, a._max_seen_label);
  }
}

template <bool is_learn>
void predict_or_learn_active_direct(active& a, learner& base, VW::example& ec)
{
  if (is_learn) { base.learn(ec); }
  else { base.predict(ec); }

  if (ec.l.simple.label == FLT_MAX)
  {
    if (std::string(ec.tag.begin(), ec.tag.begin() + 6) == "query?")
    {
      const float threshold = (a._shared_data->max_label + a._shared_data->min_label) * 0.5f;
      // We want to understand the change in prediction if the label were to be
      // the opposite of what was predicted. 0 and 1 are used for the expected min
      // and max labels to be coming in from the active interactor.
      ec.l.simple.label = (ec.pred.scalar >= threshold) ? a._min_seen_label : a._max_seen_label;
      ec.confidence = std::abs(ec.pred.scalar - threshold) / base.sensitivity(ec);
      ec.l.simple.label = FLT_MAX;
      ec.pred.scalar =
          query_decision(a, ec.confidence, static_cast<float>(a._shared_data->weighted_unlabeled_examples));
    }
  }
  else
  {
    // Update seen labels based on the current example's label.
    a._min_seen_label = std::min(ec.l.simple.label, a._min_seen_label);
    a._max_seen_label = std::max(ec.l.simple.label, a._max_seen_label);
  }
}

void active_print_result(
    VW::io::writer* f, float res, float weight, const VW::v_array<char>& tag, VW::io::logger& logger)
{
  if (f == nullptr) { return; }

  std::stringstream ss;
  ss << std::fixed << res;
  ss << " ";
  if (!tag.empty()) { ss << VW::string_view{tag.begin(), tag.size()}; }
  if (weight >= 0) { ss << " " << std::fixed << weight; }
  ss << '\n';
  const auto ss_str = ss.str();
  ssize_t len = ss_str.size();
  ssize_t t = f->write(ss_str.c_str(), static_cast<unsigned int>(len));
  if (t != len) { logger.err_error("write error: {}", VW::io::strerror_to_string(errno)); }
}

void save_load(active& a, VW::io_buf& io, bool read, bool text)
{
  using namespace VW::version_definitions;
  if (io.num_files() == 0) { return; }
  // This code is valid if version is within
  // [VERSION_FILE_WITH_ACTIVE_SEEN_LABELS, VERSION_FILE_WITH_ACTIVE_SEEN_LABELS_REVERTED)
  // or >= VERSION_FILE_WITH_ACTIVE_SEEN_LABELS_FIXED
  if ((a._model_version >= VERSION_FILE_WITH_ACTIVE_SEEN_LABELS &&
          a._model_version < VERSION_FILE_WITH_ACTIVE_SEEN_LABELS_REVERTED) ||
      a._model_version >= VERSION_FILE_WITH_ACTIVE_SEEN_LABELS_FIXED)
  {
    if (read)
    {
      VW::model_utils::read_model_field(io, a._min_seen_label);
      VW::model_utils::read_model_field(io, a._max_seen_label);
    }
    else
    {
      VW::model_utils::write_model_field(io, a._min_seen_label, "Active: min_seen_label {}", text);
      VW::model_utils::write_model_field(io, a._max_seen_label, "Active: max_seen_label {}", text);
    }
  }
}

void update_stats_active(const VW::workspace& /* all */, VW::shared_data& sd, const active& /* unused */,
    const VW::example& ec, VW::io::logger& /* logger */)
{
  const auto& ld = ec.l.simple;

  sd.update(ec.test_only, ld.label != FLT_MAX, ec.loss, ec.weight, ec.get_num_features());
  if (ld.label != FLT_MAX && !ec.test_only)
  {
    sd.weighted_labels += (static_cast<double>(ld.label)) * static_cast<double>(ec.weight);
  }
}

void output_example_prediction_active(
    VW::workspace& all, const active& data, const VW::example& ec, VW::io::logger& logger)
{
  const auto& ld = ec.l.simple;

  float ai = -1;
  if (ld.label == FLT_MAX)
  {
    ai = query_decision(data, ec.confidence, static_cast<float>(all.sd->weighted_unlabeled_examples));
  }

  all.print_by_ref(all.output_runtime.raw_prediction.get(), ec.partial_prediction, -1, ec.tag, logger);
  for (auto& i : all.output_runtime.final_prediction_sink)
  {
    active_print_result(i.get(), ec.pred.scalar, ai, ec.tag, logger);
  }
}
}  // namespace

std::shared_ptr<VW::LEARNER::learner> VW::reductions::active_setup(VW::setup_base_i& stack_builder)
{
  options_i& options = *stack_builder.get_options();
  VW::workspace& all = *stack_builder.get_all_pointer();

  bool active_option = false;
  bool simulation = false;
  bool direct = false;
  float active_c0;
  option_group_definition new_options("[Reduction] Active Learning");
  new_options.add(make_option("active", active_option).keep().necessary().help("Enable active learning"))
      .add(make_option("simulation", simulation).help("Active learning simulation mode"))
      .add(make_option("direct", direct)
               .help("Active learning via the tag and predictions interface.  Tag should start with \"query?\" to get "
                     "query decision.  Returned prediction is either -1 for no or the importance weight for yes."))
      .add(make_option("mellowness", active_c0)
               .keep()
               .default_value(1.f)
               .help("Active learning mellowness parameter c_0. Default 1."));

  if (!options.add_parse_and_check_necessary(new_options)) { return nullptr; }

  if (options.was_supplied("lda")) { THROW("lda cannot be combined with active learning") }
  auto data = VW::make_unique<active>(active_c0, all.sd, all.get_random_state(), all.runtime_state.model_file_ver);
  auto base = require_singleline(stack_builder.setup_base_learner());

  using learn_pred_func_t = void (*)(active&, VW::LEARNER::learner&, VW::example&);
  learn_pred_func_t learn_func;
  learn_pred_func_t pred_func;

  VW::learner_update_stats_func<active, VW::example>* update_stats_func = nullptr;
  VW::learner_output_example_prediction_func<active, VW::example>* output_example_prediction_func = nullptr;
  VW::learner_print_update_func<active, VW::example>* print_update_func = nullptr;

  bool learn_returns_prediction = true;
  std::string reduction_name = stack_builder.get_setupfn_name(VW::reductions::active_setup);
  if (simulation)
  {
    learn_func = predict_or_learn_simulation<true>;
    pred_func = predict_or_learn_simulation<false>;
    update_stats_func = VW::details::update_stats_simple_label<active>;
    output_example_prediction_func = VW::details::output_example_prediction_simple_label<active>;
    print_update_func = VW::details::print_update_simple_label<active>;
    reduction_name.append("-simulation");
  }
  else if (direct)
  {
    learn_func = predict_or_learn_active_direct<true>;
    pred_func = predict_or_learn_active_direct<false>;
    update_stats_func = update_stats_active;
    output_example_prediction_func = VW::details::output_example_prediction_simple_label<active>;
    print_update_func = VW::details::print_update_simple_label<active>;
    learn_returns_prediction = base->learn_returns_prediction;
  }
  else
  {
    all.reduction_state.active = true;
    learn_func = predict_or_learn_active<true>;
    pred_func = predict_or_learn_active<false>;
    update_stats_func = update_stats_active;
    output_example_prediction_func = output_example_prediction_active;
    print_update_func = VW::details::print_update_simple_label<active>;
    learn_returns_prediction = base->learn_returns_prediction;
  }

  // Create new learner
  auto l = make_reduction_learner(std::move(data), base, learn_func, pred_func, reduction_name)
               .set_input_label_type(VW::label_type_t::SIMPLE)
               .set_output_label_type(VW::label_type_t::SIMPLE)
               .set_input_prediction_type(VW::prediction_type_t::SCALAR)
               .set_output_prediction_type(VW::prediction_type_t::SCALAR)
               .set_learn_returns_prediction(learn_returns_prediction)
               .set_save_load(save_load)
               .set_update_stats(update_stats_func)
               .set_output_example_prediction(output_example_prediction_func)
               .set_print_update(print_update_func)
               .build();

  return l;
}
