// sherpa-onnx/csrc/rknn/online-transducer-modified-beam-search-decoder-rknn.cc
//
// Copyright (c)  2025  Xiaomi Corporation

#include "sherpa-onnx/csrc/rknn/online-transducer-modified-beam-search-decoder-rknn.h"

#include <algorithm>
#include <utility>
#include <vector>

#include "sherpa-onnx/csrc/hypothesis.h"
#include "sherpa-onnx/csrc/macros.h"
#include "sherpa-onnx/csrc/math.h"

namespace sherpa_onnx {

OnlineTransducerDecoderResultRknn
OnlineTransducerModifiedBeamSearchDecoderRknn::GetEmptyResult() const {
  int32_t context_size = model_->ContextSize();
  int32_t blank_id = 0;  // always 0
  OnlineTransducerDecoderResultRknn r;

  std::vector<int64_t> blanks(context_size, -1);
  blanks.back() = blank_id;

  Hypotheses blank_hyp({{blanks, 0}});
  r.hyps = std::move(blank_hyp);
  r.tokens = std::move(blanks);

  return r;
}

void OnlineTransducerModifiedBeamSearchDecoderRknn::StripLeadingBlanks(
    OnlineTransducerDecoderResultRknn *r) const {
  int32_t context_size = model_->ContextSize();
  auto hyp = r->hyps.GetMostProbable(true);

  std::vector<int64_t> tokens(hyp.ys.begin() + context_size, hyp.ys.end());
  r->tokens = std::move(tokens);
  r->timestamps = std::move(hyp.timestamps);

  r->num_trailing_blanks = hyp.num_trailing_blanks;
}

static std::vector<std::vector<float>> GetDecoderOut(
    OnlineZipformerTransducerModelRknn *model, const Hypotheses &hyp_vec) {
  std::vector<std::vector<float>> ans;
  ans.reserve(hyp_vec.Size());

  int32_t context_size = model->ContextSize();
  for (const auto &p : hyp_vec) {
    const auto &hyp = p.second;
    auto start = hyp.ys.begin() + (hyp.ys.size() - context_size);
    auto end = hyp.ys.end();
    auto tokens = std::vector<int64_t>(start, end);
    auto decoder_out = model->RunDecoder(std::move(tokens));

    ans.push_back(std::move(decoder_out));
  }

  return ans;
}

static std::vector<std::vector<float>> GetJoinerOutLogSoftmax(
    OnlineZipformerTransducerModelRknn *model, const float *p_encoder_out,
    const std::vector<std::vector<float>> &decoder_out) {
  std::vector<std::vector<float>> ans;
  ans.reserve(decoder_out.size());

  for (const auto &d : decoder_out) {
    auto joiner_out = model->RunJoiner(p_encoder_out, d.data());

    LogSoftmax(joiner_out.data(), joiner_out.size());

    ans.push_back(std::move(joiner_out));
  }
  return ans;
}

void OnlineTransducerModifiedBeamSearchDecoderRknn::Decode(
    std::vector<float> encoder_out,
    OnlineTransducerDecoderResultRknn *result) const {
  auto &r = result[0];
  auto attr = model_->GetEncoderOutAttr();
  int32_t num_frames = attr.dims[1];
  int32_t encoder_out_dim = attr.dims[2];

  int32_t vocab_size = model_->VocabSize();
  int32_t context_size = model_->ContextSize();

  Hypotheses cur = std::move(result->hyps);
  std::vector<Hypothesis> prev;

  auto decoder_out = std::move(result->previous_decoder_out2);
  if (decoder_out.empty()) {
    decoder_out = GetDecoderOut(model_, cur);
  }

  const float *p_encoder_out = encoder_out.data();

  int32_t frame_offset = result->frame_offset;

  for (int32_t t = 0; t != num_frames; ++t) {
    prev = cur.Vec();
    cur.Clear();

    auto log_probs = GetJoinerOutLogSoftmax(model_, p_encoder_out, decoder_out);
    p_encoder_out += encoder_out_dim;

    for (int32_t i = 0; i != prev.size(); ++i) {
      auto log_prob = prev[i].log_prob;
      for (auto &p : log_probs[i]) {
        p += log_prob;
      }
    }

    auto topk = TopkIndex(log_probs, max_active_paths_);
    for (auto k : topk) {
      int32_t hyp_index = k / vocab_size;
      int32_t new_token = k % vocab_size;

      Hypothesis new_hyp = prev[hyp_index];
      new_hyp.log_prob = log_probs[hyp_index][new_token];

      // blank is hardcoded to 0
      // also, it treats unk as blank
      if (new_token != 0 && new_token != unk_id_) {
        new_hyp.ys.push_back(new_token);
        new_hyp.timestamps.push_back(t + frame_offset);
        new_hyp.num_trailing_blanks = 0;

      } else {
        ++new_hyp.num_trailing_blanks;
      }
      cur.Add(std::move(new_hyp));
    }

    decoder_out = GetDecoderOut(model_, cur);
  }

  result->hyps = std::move(cur);
  result->frame_offset += num_frames;
  result->previous_decoder_out2 = std::move(decoder_out);
}

}  // namespace sherpa_onnx