/*!
 *  Copyright (c) 2023-2025 by Contributors
 * \file serve/sampler/cpu_sampler.cc
 * \brief The implementation for CPU sampler functions.
 */
#include <tvm/ffi/function.h>
#include <tvm/runtime/tensor.h>

#include <algorithm>
#include <cmath>

#include "../../support/random.h"
#include "../../support/threading_backend.h"
#include "sampler.h"

namespace mlc {
namespace llm {
namespace serve {

TVM_FFI_STATIC_INIT_BLOCK() { SamplerObj::RegisterReflection(); }

/*!
 * \brief Sample a value from the input probability distribution with top-p.
 * The input is a batch of distributions, and we use `unit_offset` to specify
 * which distribution to sample from.
 * \param prob The input batch of probability distributions.
 * \param unit_offset The offset specifying which distribution to output
 * \param input_prob_offset The offset specifying which distribution to sample from.
 * \param top_p The top-p value of sampling.
 * \param uniform_sample The random number in [0, 1] for sampling.
 * \return The sampled value and probability.
 * \note This function is an enhancement of SampleTopPFromProb in TVM Unity.
 * We will upstream the enhancement after it gets stable.
 */
TokenProbPair SampleTopPFromProb(Tensor prob, int unit_offset, int input_prob_offset, double top_p,
                                 double uniform_sample) {
  // prob: (*, v)
  // The prob array may have arbitrary ndim and shape.
  // The last dimension corresponds to the prob distribution size.
  // We use the `unit_offset` parameter to determine which slice
  // of the prob array we sample from.

  TVM_FFI_ICHECK(prob.IsContiguous());
  TVM_FFI_ICHECK(prob.DataType() == (DLDataType{kDLFloat, 32, 1}));
  TVM_FFI_ICHECK_EQ(prob->device.device_type, DLDeviceType::kDLCPU);

  int64_t ndata = prob->shape[prob->ndim - 1];
  const float* __restrict p_prob =
      static_cast<float*>(__builtin_assume_aligned(prob->data, 4)) + (input_prob_offset * ndata);
  constexpr double one = 1.0f - 1e-5f;

  if (top_p == 0) {
    // Specially handle case where top_p == 0.
    // This case is equivalent to doing argmax.
    int argmax_pos = -1;
    float max_prob = 0.0;
    float sum_prob = 0.0;
    for (int i = 0; i < ndata; ++i) {
      if (p_prob[i] > max_prob) {
        max_prob = p_prob[i];
        argmax_pos = i;
      }
      // Early exit.
      sum_prob += p_prob[i];
      if (1 - sum_prob <= max_prob) {
        break;
      }
    }
    return {argmax_pos, 1.0};
  }

  if (top_p >= one) {
    // Specially handle case where top_p == 1.
    double prob_sum = 0.0f;
    for (int64_t i = 0; i < ndata; ++i) {
      prob_sum += p_prob[i];
      if (prob_sum >= uniform_sample) {
        return {i, p_prob[i]};
      }
    }
    TVM_FFI_ICHECK(false) << "Possibly prob distribution contains NAN.";
  }

  // Key observation: when we are doing top_p sampling
  // usually we only need to preserve some of the elements with
  // high probabilities before we do sort
  thread_local std::vector<std::pair<float, int>> data;

  auto sample_top_p_with_filter = [&](float cuttoff) -> std::pair<float, int64_t> {
    data.clear();
    // filter the data with cuttoff
    float cutoff_sum = 0.0f;
    for (int64_t i = 0; i < ndata; ++i) {
      if (p_prob[i] >= cuttoff) {
        cutoff_sum += p_prob[i];
        data.emplace_back(std::make_pair(p_prob[i], static_cast<int>(i)));
        if (cutoff_sum > 1 - cuttoff) {
          // Short cut. When the remaining parts cannot have total
          // probability larger than cutoff, we can quit.
          break;
        }
      }
    }
    if (data.size() == 0) return std::make_pair(-1, -1);
    auto fcmp = [](const std::pair<float, int>& lhs, const std::pair<float, int>& rhs) {
      return lhs.first > rhs.first;
    };
    std::sort(data.begin(), data.end(), fcmp);

    // short cut, if we know that
    // uniform sample < p[0] / top_p
    // we know that unform_sample < p[0] / top_p_sum
    // because top_p_sum guarantees to be smaller than top_p
    // so we can simply return the argmax sample
    // without computing anything
    if (uniform_sample < data[0].first / top_p) {
      return std::make_pair(data[0].first, data[0].second);
    }

    // compute top_p_sum
    float cum_sum_prob = 0.0f;
    float top_p_sum = 0.0f;
    for (auto it = data.begin(); it != data.end(); ++it) {
      float prob = it->first;
      if (cum_sum_prob < top_p) {
        top_p_sum += prob;
      } else {
        // we get to the right cutoff pt
        break;
      }
      cum_sum_prob += prob;
      it->first = cum_sum_prob;
    }
    // we find that the current total sum by the given cutoff
    // is not sufficient to cover everything
    // this means we might need to retry a smaller cutoff pt.
    if (cum_sum_prob < top_p && cuttoff != 0.0f) return std::make_pair(-1, -1);

    float last_cum_sum_prob = 0.0;
    for (auto it = data.begin(); it != data.end(); ++it) {
      if (uniform_sample < it->first / top_p_sum) {
        return std::make_pair(it->first - last_cum_sum_prob, it->second);
      }
      last_cum_sum_prob = it->first;
    }
    return std::make_pair(data[static_cast<int64_t>(data.size()) - 1].first - last_cum_sum_prob,
                          data[static_cast<int64_t>(data.size()) - 1].second);
  };

  if (top_p < 1) {
    // sample through cutoff by a number
    // by pigeonhole principle we will get at most 1024 elements
    // usually it is much less by applying this filtering(order of 10 - 20)
    data.reserve(256);
    std::pair<float, int64_t> sampled_index = sample_top_p_with_filter(top_p / 1024);
    if (sampled_index.second >= 0) return {sampled_index.second, sampled_index.first};
  }
  // fallback via full prob, rare case
  data.reserve(ndata);
  std::pair<float, int64_t> sampled_index = sample_top_p_with_filter(0.0f);
  TVM_FFI_ICHECK_GE(sampled_index.second, 0);
  return {sampled_index.second, sampled_index.first};
}

/*!
 * \brief Renormalize the probability distribution by the top p value.
 * \param prob The input batch of probability distributions.
 * \param unit_offset The offset specifying which distribution to output
 * \param top_p The top p value for renormalization.
 * \param eps A small epsilon value for comparison stability.
 */
void RenormalizeProbByTopP(Tensor prob, int unit_offset, double top_p, double eps) {
  // prob: (*, v)
  // The prob array may have arbitrary ndim and shape.
  // The last dimension corresponds to the prob distribution size.
  // We use the `unit_offset` parameter to determine which slice
  // of the prob array we will renormalize.
  TVM_FFI_ICHECK(prob.IsContiguous());
  TVM_FFI_ICHECK(prob.DataType() == (DLDataType{kDLFloat, 32, 1}));
  TVM_FFI_ICHECK_EQ(prob->device.device_type, DLDeviceType::kDLCPU);

  if (top_p == 1.0) {
    // No renormalization is needed if top_p is 1.
    return;
  }

  int vocab_size = prob->shape[prob->ndim - 1];
  float* __restrict p_prob =
      static_cast<float*>(__builtin_assume_aligned(prob->data, 4)) + (unit_offset * vocab_size);

  // We manually choice the cutoff values of "top_p / 256" and "top_p / 8192".
  // In most of the cases, only one round is needed.
  std::vector<double> cutoff_values{top_p / 256, top_p / 8192, 0.0f};

  // Create the upper partition vector and the lower partition rolling vectors.
  std::vector<float> upper_partition;
  std::vector<float> lower_partitions[2];
  upper_partition.reserve(vocab_size);
  lower_partitions[0].reserve(vocab_size);
  lower_partitions[1].reserve(vocab_size);
  float upper_partition_sum = 0.0;
  for (int round = 0; round < static_cast<int>(cutoff_values.size()); ++round) {
    const float* lower_partition_begin;
    const float* lower_partition_end;
    if (round == 0) {
      lower_partition_begin = p_prob;
      lower_partition_end = p_prob + vocab_size;
    } else {
      int idx = (round - 1) & 1;
      lower_partition_begin = lower_partitions[idx].data();
      lower_partition_end = lower_partitions[idx].data() + lower_partitions[idx].size();
    }

    // - Partition the last round lower partition into upper and lower
    // based on the new cutoff value.
    std::vector<float>& lower_partition = lower_partitions[round & 1];
    lower_partition.clear();
    for (const float* ptr = lower_partition_begin; ptr != lower_partition_end; ++ptr) {
      if (*ptr >= cutoff_values[round]) {
        upper_partition.push_back(*ptr);
        upper_partition_sum += *ptr;
      } else {
        lower_partition.push_back(*ptr);
      }
    }
    // - If the upper partition sum is at least top p, exit the loop.
    if (upper_partition_sum >= top_p - eps) {
      break;
    }
  }

  // - Sort the upper partition in descending order.
  std::sort(upper_partition.begin(), upper_partition.end(), std::greater<>());
  // - Find the top p boundary prob value.
  float boundary_value = -1.0;
  upper_partition_sum = 0.0;
  for (float upper_value : upper_partition) {
    upper_partition_sum += upper_value;
    if (upper_partition_sum >= top_p - eps) {
      boundary_value = upper_value;
      break;
    }
  }
  // - Mask all values smaller than the boundary to 0.
  float renormalize_sum = 0.0;
  std::vector<int> upper_partition_indices;
  upper_partition_indices.reserve(vocab_size);
  for (int i = 0; i < vocab_size; ++i) {
    if (p_prob[i] >= boundary_value) {
      upper_partition_indices.push_back(i);
      renormalize_sum += p_prob[i];
    } else {
      p_prob[i] = 0.0;
    }
  }
  // - Renormalize.
  for (int idx : upper_partition_indices) {
    p_prob[idx] /= renormalize_sum;
  }
}

namespace detail {

/*! \brief Implementation of getting top probs on CPU. */
template <int num_top_probs>
std::vector<TokenProbPair> ComputeTopProbsImpl(const float* p_prob, int ndata) {
  std::vector<TokenProbPair> top_probs;
  top_probs.reserve(num_top_probs);
  for (int i = 0; i < num_top_probs; ++i) {
    top_probs.emplace_back(-1, -1.0f);
  }

  float sum_prob = 0.0;
  // Selection argsort.
  for (int p = 0; p < ndata; ++p) {
    int i = num_top_probs - 1;
    for (; i >= 0; --i) {
      if (p_prob[p] > top_probs[i].second) {
        if (i != num_top_probs - 1) {
          top_probs[i + 1] = top_probs[i];
        }
      } else {
        break;
      }
    }
    if (i != num_top_probs - 1) {
      top_probs[i + 1] = {p, p_prob[p]};
    }

    // Early exit.
    sum_prob += p_prob[p];
    if (1 - sum_prob <= top_probs[num_top_probs - 1].second) {
      break;
    }
  }
  return top_probs;
}

}  // namespace detail

/*! \brief Get the probs of a few number of tokens with top probabilities. */
inline std::vector<TokenProbPair> ComputeTopProbs(Tensor prob, int unit_offset, int num_top_probs) {
  TVM_FFI_ICHECK_LE(num_top_probs, 5);
  TVM_FFI_ICHECK_EQ(prob->ndim, 2);
  int ndata = prob->shape[1];
  const float* __restrict p_prob =
      static_cast<float*>(__builtin_assume_aligned(prob->data, 4)) + (unit_offset * ndata);
  switch (num_top_probs) {
    case 0:
      return {};
    case 1:
      return detail::ComputeTopProbsImpl<1>(p_prob, ndata);
    case 2:
      return detail::ComputeTopProbsImpl<2>(p_prob, ndata);
    case 3:
      return detail::ComputeTopProbsImpl<3>(p_prob, ndata);
    case 4:
      return detail::ComputeTopProbsImpl<4>(p_prob, ndata);
    case 5:
      return detail::ComputeTopProbsImpl<5>(p_prob, ndata);
  }
  throw;
}

/********************* CPU Sampler *********************/

class CPUSampler : public SamplerObj {
 public:
  explicit CPUSampler(Optional<EventTraceRecorder> trace_recorder)
      : trace_recorder_(std::move(trace_recorder)) {}

  Tensor BatchRenormalizeProbsByTopP(Tensor probs_on_device,                  //
                                     const std::vector<int>& sample_indices,  //
                                     const Array<String>& request_ids,        //
                                     const Array<GenerationConfig>& generation_cfg) final {
    // probs_on_device: (n, v)
    TVM_FFI_ICHECK_EQ(probs_on_device->ndim, 2);
    // - Copy probs to CPU
    RECORD_EVENT(trace_recorder_, request_ids, "start copy probs to CPU");
    Tensor probs_on_host = CopyProbsToCPU(probs_on_device);
    RECORD_EVENT(trace_recorder_, request_ids, "finish copy probs to CPU");
    int num_samples = sample_indices.size();
    int num_probs = probs_on_device->shape[0];
    int vocab_size = probs_on_device->shape[1];
    TVM_FFI_ICHECK_EQ(request_ids.size(), num_samples);
    TVM_FFI_ICHECK_EQ(generation_cfg.size(), num_samples);

    std::vector<int> top_p_indices;
    std::vector<double> top_p_values;
    for (int i = 0; i < num_samples; ++i) {
      if (top_p_indices.empty() || top_p_indices.back() != sample_indices[i]) {
        top_p_indices.push_back(sample_indices[i]);
        top_p_values.push_back(generation_cfg[i]->top_p);
      } else {
        TVM_FFI_ICHECK(fabs(top_p_values.back() - generation_cfg[i]->top_p) < eps_)
            << "Sampler requires the top_p values for each prob distribution are the same.";
      }
    }
    if (top_p_indices.empty()) {
      // Return if no top p needs to apply.
      return probs_on_host;
    }

    tvm::runtime::parallel_for_with_threading_backend(
        [this, &probs_on_host, &request_ids, &top_p_indices, &top_p_values](int i) {
          RECORD_EVENT(this->trace_recorder_, request_ids[i], "start renormalize by top p");
          RenormalizeProbByTopP(probs_on_host, top_p_indices[i], top_p_values[i], eps_);
          RECORD_EVENT(this->trace_recorder_, request_ids[i], "finish renormalize by top p");
        },
        0, static_cast<int64_t>(top_p_indices.size()));

    return probs_on_host;
  }

  std::vector<SampleResult> BatchSampleTokensWithProbBeforeTopP(
      Tensor probs_on_device,                         //
      const std::vector<int>& sample_indices,         //
      const Array<String>& request_ids,               //
      const Array<GenerationConfig>& generation_cfg,  //
      const std::vector<RandomGenerator*>& rngs) final {
    // probs_on_device: (n, v)
    TVM_FFI_ICHECK_EQ(probs_on_device->ndim, 2);
    // - Copy probs to CPU
    RECORD_EVENT(trace_recorder_, request_ids, "start copy probs to CPU");
    Tensor probs_on_host = CopyProbsToCPU(probs_on_device);
    RECORD_EVENT(trace_recorder_, request_ids, "finish copy probs to CPU");

    return BatchSampleTokensImpl(probs_on_host, sample_indices, request_ids, generation_cfg, rngs,
                                 /*top_p_applied=*/false);
  }

  std::vector<SampleResult> BatchSampleTokensWithProbAfterTopP(
      Tensor probs_on_host,                           //
      const std::vector<int>& sample_indices,         //
      const Array<String>& request_ids,               //
      const Array<GenerationConfig>& generation_cfg,  //
      const std::vector<RandomGenerator*>& rngs) final {
    return BatchSampleTokensImpl(probs_on_host, sample_indices, request_ids, generation_cfg, rngs,
                                 /*top_p_applied=*/true);
  }

  std::pair<std::vector<std::vector<SampleResult>>, std::vector<int>>
  BatchVerifyDraftTokensWithProbAfterTopP(
      Tensor probs_on_host, const Array<String>& request_ids,
      const std::vector<int>& cum_verify_lengths, const Array<GenerationConfig>& generation_cfg,
      const std::vector<RandomGenerator*>& rngs,
      const std::vector<std::vector<SampleResult>>& draft_output_tokens,
      const std::vector<int64_t>& token_tree_parent_ptr, Tensor draft_probs_on_device) final {
    // probs_on_host: (n, v)
    RECORD_EVENT(trace_recorder_, request_ids, "start draft verification");
    TVM_FFI_ICHECK_EQ(probs_on_host->ndim, 2);

    int num_sequence = static_cast<int>(cum_verify_lengths.size()) - 1;
    TVM_FFI_ICHECK_EQ(rngs.size(), num_sequence);
    TVM_FFI_ICHECK_EQ(draft_output_tokens.size(), num_sequence);

    Tensor draft_probs_on_host = draft_probs_on_device.CopyTo(DLDevice{kDLCPU, 0});
    std::vector<std::vector<SampleResult>> sample_results;
    sample_results.resize(num_sequence);

    float* __restrict global_p_probs =
        static_cast<float*>(__builtin_assume_aligned(probs_on_host->data, 4));
    int vocab_size = probs_on_host->shape[1];

    std::vector<int> last_accepted_tree_node(num_sequence, 0);
    tvm::runtime::parallel_for_with_threading_backend(
        [&](int i) {
          int verify_start = cum_verify_lengths[i];
          int verify_end = cum_verify_lengths[i + 1];

          TVM_FFI_ICHECK_EQ(token_tree_parent_ptr[verify_start], -1);
          for (int j = verify_start + 1; j < verify_end; ++j) {
            TVM_FFI_ICHECK_EQ(token_tree_parent_ptr[j], j - verify_start - 1)
                << "CPU sampler only supports chain-style draft tokens.";
          }

          int cur_token_idx = 0;
          // Sub 1 to ignore the last prediction.
          for (; cur_token_idx < verify_end - verify_start - 1; ++cur_token_idx) {
            float* p_probs = global_p_probs + (verify_start + cur_token_idx) * vocab_size;
            int cur_token = draft_output_tokens[i][cur_token_idx].GetTokenId();
            float q_value = draft_output_tokens[i][cur_token_idx].sampled_token_id.second;
            float p_value = p_probs[cur_token];

            if (p_value >= q_value) {
              sample_results[i].push_back(
                  SampleResult{{cur_token, p_value},
                               ComputeTopProbs(probs_on_host, verify_start + cur_token_idx,
                                               generation_cfg[i]->top_logprobs)});
              continue;
            }
            float r = rngs[i]->GetRandomNumber();
            if (r < p_value / (q_value + eps_)) {
              sample_results[i].push_back(
                  SampleResult{{cur_token, p_value},
                               ComputeTopProbs(probs_on_host, verify_start + cur_token_idx,
                                               generation_cfg[i]->top_logprobs)});
              continue;
            }

            // normalize a new probability distribution
            double sum_v = 0.0;
            const float* __restrict p_qdist =
                static_cast<float*>(__builtin_assume_aligned(draft_probs_on_host->data, 4)) +
                (verify_start + cur_token_idx + 1) * vocab_size;

            for (int j = 0; j < vocab_size; ++j) {
              p_probs[j] = std::max(p_probs[j] - p_qdist[j], 0.0f);
              sum_v += p_probs[j];
            }
            for (int j = 0; j < vocab_size; ++j) {
              p_probs[j] /= sum_v;
            }

            // sample a new token from the new distribution
            SampleResult sample_result;
            sample_result.sampled_token_id = SampleTopPFromProb(
                probs_on_host, verify_start + cur_token_idx, verify_start + cur_token_idx,
                /*top_p=*/1.0f, rngs[i]->GetRandomNumber());
            sample_result.top_prob_tokens = ComputeTopProbs(
                probs_on_host, verify_start + cur_token_idx, generation_cfg[i]->top_logprobs);
            sample_results[i].push_back(sample_result);
            break;
          }
          last_accepted_tree_node[i] = cur_token_idx;
          // if cur_token_idx == verify_end - verify_start - 1
          // all draft tokens are accepted
          // we sample a new token
          if (cur_token_idx == verify_end - verify_start - 1) {
            SampleResult sample_result;
            // sample a new token from the original distribution
            sample_result.sampled_token_id = SampleTopPFromProb(
                probs_on_host, verify_start + cur_token_idx, verify_start + cur_token_idx,
                /*top_p=*/1.0f, rngs[i]->GetRandomNumber());
            sample_result.top_prob_tokens = ComputeTopProbs(
                probs_on_host, verify_start + cur_token_idx, generation_cfg[i]->top_logprobs);
            sample_results[i].push_back(sample_result);
          }
        },
        0, num_sequence);
    RECORD_EVENT(trace_recorder_, request_ids, "finish draft verification");
    return {sample_results, last_accepted_tree_node};
  }

 private:
  std::vector<SampleResult> BatchSampleTokensImpl(Tensor probs_on_host,                           //
                                                  const std::vector<int>& sample_indices,         //
                                                  const Array<String>& request_ids,               //
                                                  const Array<GenerationConfig>& generation_cfg,  //
                                                  const std::vector<RandomGenerator*>& rngs,      //
                                                  bool top_p_applied) {
    // probs_on_host: (n, v)
    RECORD_EVENT(trace_recorder_, request_ids, "start sampling");
    TVM_FFI_ICHECK_EQ(probs_on_host->ndim, 2);
    TVM_FFI_ICHECK_EQ(probs_on_host->device.device_type, DLDeviceType::kDLCPU);

    // - Sample tokens from probabilities.
    int n = request_ids.size();
    TVM_FFI_ICHECK_EQ(generation_cfg.size(), n);
    TVM_FFI_ICHECK_EQ(rngs.size(), n);

    std::vector<SampleResult> sample_results;
    sample_results.resize(n);

    tvm::runtime::parallel_for_with_threading_backend(
        [this, &sample_results, &probs_on_host, &generation_cfg, &rngs, &request_ids, top_p_applied,
         sample_indices](int i) {
          RECORD_EVENT(this->trace_recorder_, request_ids[i], "start sample token");
          // Sample top p from probability.
          double top_p =
              top_p_applied
                  ? 1.0f
                  : (generation_cfg[i]->temperature < eps_ ? 0.0 : generation_cfg[i]->top_p);
          sample_results[i].sampled_token_id = SampleTopPFromProb(
              probs_on_host, i, sample_indices[i], top_p, rngs[i]->GetRandomNumber());
          sample_results[i].top_prob_tokens =
              ComputeTopProbs(probs_on_host, i, generation_cfg[i]->top_logprobs);
          RECORD_EVENT(this->trace_recorder_, request_ids[i], "finish sample token");
        },
        0, n);
    RECORD_EVENT(trace_recorder_, request_ids, "finish sampling");
    return sample_results;
  }

  /*! \brief Copy prob distributions from device to CPU. */
  Tensor CopyProbsToCPU(Tensor probs_on_device) {
    // probs_on_device: (n, v)
    if (probs_on_device->device.device_type == kDLCPU) {
      return probs_on_device;
    }

    TVM_FFI_ICHECK(probs_on_device->device.device_type != kDLCPU);
    if (probs_host_.defined()) {
      TVM_FFI_ICHECK_EQ(probs_host_->shape[1], probs_on_device->shape[1]);
    }

    int64_t init_size = probs_host_.defined() ? probs_host_->shape[0] : 32;
    int64_t num_tokens = probs_on_device->shape[0];
    int64_t vocab_size = probs_on_device->shape[1];
    while (init_size < num_tokens) {
      init_size *= 2;
    }
    if (!probs_host_.defined() || init_size != probs_host_->shape[0]) {
      probs_host_ =
          Tensor::Empty({init_size, vocab_size}, probs_on_device->dtype, DLDevice{kDLCPU, 0});
    }
    TVM_FFI_ICHECK_LE(num_tokens, probs_host_->shape[0]);
    Tensor view = probs_host_.CreateView({num_tokens, vocab_size}, probs_on_device->dtype);
    view.CopyFrom(probs_on_device);
    return view;
  }

  /*! \brief The event trace recorder for requests. */
  Optional<EventTraceRecorder> trace_recorder_;
  /*! \brief Customized function which computes prob distribution from logits */
  Function flogits_to_probs_inplace_;
  /*! \brief Probability distribution array on CPU. */
  Tensor probs_host_{nullptr};
  const float eps_ = 1e-5;
};

Sampler Sampler::CreateCPUSampler(Optional<EventTraceRecorder> trace_recorder) {
  return Sampler(tvm::ffi::make_object<CPUSampler>(std::move(trace_recorder)));
}

}  // namespace serve
}  // namespace llm
}  // namespace mlc