// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "../generators.h"
#include "marian.h"

namespace Generators {

MarianModel::MarianModel(std::unique_ptr<Config> config, OrtEnv& ort_env)
    : Model{std::move(config)} {
  encoder_session_options_ = OrtSessionOptions::Create();
  CreateSessionOptionsFromConfig(config_->model.encoder.session_options.has_value() ? config_->model.encoder.session_options.value() : config_->model.decoder.session_options, *encoder_session_options_, true);

  session_encoder_ = CreateSession(ort_env, config_->model.encoder.filename, encoder_session_options_.get());
  session_decoder_ = CreateSession(ort_env, config_->model.decoder.filename, session_options_.get());

  session_info_.Add(*session_decoder_);
  session_info_.Add(*session_encoder_);
}

MarianLogits::MarianLogits(State& state)
    : state_{state},
      shape_{static_cast<int64_t>(state_.params_->BatchBeamSize()), model_.config_->model.vocab_size},
      type_{model_.session_info_.GetOutputDataType(model_.config_->model.decoder.outputs.logits)} {
  output_raw_ = std::make_unique<Tensor>(model_.p_device_inputs_, type_);
}

std::unique_ptr<State> MarianModel::CreateState(DeviceSpan<int32_t> sequence_lengths, const GeneratorParams& params) const {
  return std::make_unique<MarianState>(*this, sequence_lengths, params);
}

MarianState::MarianState(const MarianModel& model, DeviceSpan<int32_t> sequence_lengths_unk, const GeneratorParams& params)
    : State{params, model},
      model_{model},
      encoder_attention_mask_{model, *this, sequence_lengths_unk, model_.config_->model.encoder.inputs.attention_mask},
      attention_mask_{model, *this, sequence_lengths_unk, model_.config_->model.decoder.inputs.encoder_attention_mask} {
}

DeviceSpan<float> MarianLogits::Get() {
  // The model's output logits are {batch_size*num_beams, vocab_size}

  // TODO(apsonawane): Fix the issue with output_raw not getting updated properly
  // OrtValue* logits_of_last_token = output_raw_->GetOrtTensor();
  OrtValue* logits_of_last_token = state_.outputs_[output_index_];

  // Convert from float16 to float32 if necessary
  if (type_ == Ort::TypeToTensorType<Ort::Float16_t>) {
    Cast(*logits_of_last_token, logits_of_last_token_fp32_, *model_.p_device_inputs_, Ort::TypeToTensorType<float>);
    logits_of_last_token = logits_of_last_token_fp32_.get();
  }

  if (logits_.empty() || logits_of_last_token->GetTensorMutableRawData() != logits_.Span().data())
    logits_ = WrapTensor<float>(*model_.p_device_inputs_, *logits_of_last_token);

  return logits_;
}

void MarianLogits::Update(const DeviceSpan<int32_t>& next_tokens, size_t new_kv_length) {
  if (output_raw_->ort_tensor_ && static_cast<size_t>(output_raw_->GetShape()[1]) == new_kv_length && new_kv_length == 1) {
    return;
  }

  output_raw_->CreateTensor(shape_, state_.params_->use_graph_capture);
  state_.outputs_[output_index_] = output_raw_->GetOrtTensor();
}

void MarianLogits::Add() {
  output_index_ = state_.outputs_.size();

  state_.output_names_.push_back(model_.config_->model.decoder.outputs.logits.c_str());
  state_.outputs_.push_back(output_raw_->GetOrtTensor());
}

MarianInputIDs::MarianInputIDs(State& state)
    : state_{state} {
  name_ = model_.config_->model.decoder.inputs.input_ids.c_str();
  shape_ = {state_.params_->BatchBeamSize()};
  type_ = model_.session_info_.GetInputDataType(name_);

  value_ = std::make_unique<Tensor>(model_.p_device_inputs_, Ort::TypeToTensorType<int32_t>);
  cast_value_ = std::make_unique<Tensor>(model_.p_device_inputs_, Ort::TypeToTensorType<int64_t>);
}

void MarianInputIDs::Add() {
  input_index_ = state_.inputs_.size();

  state_.inputs_.push_back(value_->GetOrtTensor());
  state_.input_names_.push_back(name_);
}

void MarianInputIDs::Update(DeviceSpan<int32_t> new_tokens) {
  // The tensor holds one token per batch*beam entry. The copy below is
  // unchecked in release builds (DeviceSpan::CopyFrom only asserts) and always
  // copies the destination's length, so a shorter source is read past its end
  // and the tail fills with unrelated memory. That is then consumed as token
  // ids by the decoder's embedding gather, which reads out of bounds.
  // Validate explicitly so the caller gets a diagnosable error instead.
  if (new_tokens.size() != static_cast<size_t>(shape_[0])) {
    throw std::runtime_error(
        "MarianInputIDs::Update: got " + std::to_string(new_tokens.size()) +
        " token(s) but the input tensor holds " + std::to_string(shape_[0]) +
        " (batch_size * num_beams). Every batch entry needs exactly one token.");
  }

  value_->CreateTensor(shape_, state_.params_->use_graph_capture);
  state_.inputs_[input_index_] = value_->GetOrtTensor();

  // Update input_ids with next tokens
  auto data_span = value_->GetDeviceSpan<int32_t>();
  data_span.CopyFrom(new_tokens);

  if (type_ == Ort::TypeToTensorType<int64_t>) {
    if (!cast_value_->GetOrtTensor())
      cast_value_->CreateTensor(shape_, state_.params_->use_graph_capture);
    Cast(*value_->GetOrtTensor(), cast_value_->ort_tensor_, *model_.p_device_inputs_, type_);
    state_.inputs_[input_index_] = cast_value_->GetOrtTensor();
  }
}

DeviceSpan<float> MarianState::Run(int current_length, DeviceSpan<int32_t>& next_tokens, DeviceSpan<int32_t> next_indices) {
  const int64_t encoder_hidden_size = model_.config_->model.encoder.hidden_size;
  if (first_run_) {
    // Initialize the encoder and run it once
    encoder_input_ids_.name_ = model_.config_->model.encoder.inputs.input_ids.c_str();
    encoder_input_ids_.Add();

    encoder_attention_mask_.Add();

    // Append eos per sequence at its own unpadded length. PadInputs right-pads,
    // so appending once at the end of the flattened buffer would leave eos
    // behind the padding for every row but the last.
    //
    // The prompt buffer is [batch_size, seq_len]: DefaultInputIDs and
    // DefaultPositionInputs expand it over beams themselves, so only
    // decoder-side tensors are sized batch*beam.
    const size_t batch_size = static_cast<size_t>(params_->search.batch_size);
    const int32_t pad_token_id = model_.config_->model.pad_token_id;
    const int32_t eos_token_id = model_.config_->model.eos_token_id[0];

    auto original_cpu_span = next_tokens.CpuSpan();
    if (batch_size == 0 || original_cpu_span.size() % batch_size != 0) {
      throw std::runtime_error(
          "MarianState::Run: prompt buffer holds " +
          std::to_string(original_cpu_span.size()) +
          " token(s), which is not a positive multiple of batch_size " +
          std::to_string(batch_size) +
          ". The buffer must be a padded [batch_size, sequence_length] block.");
    }
    const size_t src_width = original_cpu_span.size() / batch_size;
    const size_t dst_width = src_width + 1;

    auto extended_tokens = model_.p_device_inputs_->Allocate<int32_t>(batch_size * dst_width);
    auto extended_cpu_span = extended_tokens.CpuSpan();
    for (size_t b = 0; b < batch_size; b++) {
      auto src = original_cpu_span.subspan(b * src_width, src_width);
      auto dst = extended_cpu_span.subspan(b * dst_width, dst_width);

      size_t length = src_width;
      for (size_t i = 0; i < src_width; i++) {
        if (src[i] == pad_token_id) {
          length = i;
          break;
        }
      }

      std::copy(src.begin(), src.begin() + length, dst.begin());
      dst[length] = eos_token_id;
      std::fill(dst.begin() + length + 1, dst.end(), pad_token_id);
    }
    extended_tokens.CopyCpuToDevice();

    encoder_input_ids_.Update(extended_tokens);
    const size_t new_length = static_cast<size_t>(encoder_input_ids_.GetShape()[1]);
    // The mask is built by scanning the token buffer over batch*seq_len, so it
    // must see the extended buffer rather than the original padded one.
    encoder_attention_mask_.Update(extended_tokens, current_length, static_cast<int>(new_length));

    const auto encoder_outputs_type = model_.session_info_.GetOutputDataType(model_.config_->model.encoder.outputs.encoder_outputs.c_str());
    const std::array<int64_t, 3> encoder_outputs_shape{encoder_input_ids_.GetShape()[0], encoder_input_ids_.GetShape()[1], encoder_hidden_size};
    encoder_outputs_ = OrtValue::CreateTensor(model_.p_device_inputs_->GetAllocator(), encoder_outputs_shape, encoder_outputs_type);

    output_names_.push_back(model_.config_->model.encoder.outputs.encoder_outputs.c_str());
    outputs_.push_back(encoder_outputs_.get());

    if (model_.config_->model.encoder.run_options.has_value()) {
      State::SetRunOptions(model_.config_->model.encoder.run_options.value());
    }
    State::Run(*model_.session_encoder_);

    // Clear inputs and outputs for the decoder
    ClearIO();

    // Initialize the decoder inputs and outputs
    decoder_input_ids_.name_ = model_.config_->model.decoder.inputs.input_ids.c_str();
    decoder_input_ids_.Add();

    // Every sequence in the batch starts decoding from bos. Decoder tensors are
    // sized batch*beam because beam expansion has already happened by here.
    auto decoder_start_tokens = model_.p_device_inputs_->Allocate<int32_t>(params_->BatchBeamSize());
    std::fill(decoder_start_tokens.CpuSpan().begin(), decoder_start_tokens.CpuSpan().end(),
              model_.config_->model.bos_token_id);
    decoder_start_tokens.CopyCpuToDevice();

    decoder_input_ids_.Update(decoder_start_tokens);

    const std::array<int64_t, 1> past_key_values_length_shape{1};
    past_key_values_length_ = OrtValue::CreateTensor(model_.allocator_cpu_, past_key_values_length_shape, model_.session_info_.GetInputDataType(model_.config_->model.decoder.inputs.past_key_values_length));
    input_names_.push_back(model_.config_->model.decoder.inputs.past_key_values_length.c_str());
    inputs_.push_back(past_key_values_length_.get());

    *past_key_values_length_->GetTensorMutableData<int64_t>() = -1;

    // attention_mask_.attention_mask_name_ = model_.config_->model.decoder.inputs.encoder_attention_mask;
    attention_mask_.Add();
    attention_mask_.Update(extended_tokens, current_length, static_cast<int>(new_length));

    const auto hidden_states_type = model_.session_info_.GetInputDataType(model_.config_->model.decoder.inputs.encoder_hidden_states.c_str());
    const std::array<int64_t, 3> encoder_hidden_states_shape{decoder_input_ids_.GetMarianInputsShape()[0], encoder_input_ids_.GetShape()[1], encoder_hidden_size};
    encoder_hidden_states_ = OrtValue::CreateTensor(model_.p_device_inputs_->GetAllocator(), encoder_hidden_states_shape, hidden_states_type);

    input_names_.push_back(model_.config_->model.decoder.inputs.encoder_hidden_states.c_str());
    inputs_.push_back(encoder_outputs_.get());

    const auto rnn_states_prev_type = model_.session_info_.GetInputDataType(model_.config_->model.decoder.inputs.rnn_prev_states.c_str());
    const std::array<int64_t, 3> rnn_states_prev_shape{3, decoder_input_ids_.GetMarianInputsShape()[0], encoder_hidden_size};
    rnn_states_prev_ = std::make_unique<Tensor>(model_.p_device_inputs_, rnn_states_prev_type);
    rnn_states_prev_->CreateTensor(rnn_states_prev_shape);

    input_names_.push_back(model_.config_->model.decoder.inputs.rnn_prev_states.c_str());

    auto device_span = rnn_states_prev_->GetDeviceSpan<int32_t>();
    device_span.Zero();
    inputs_.push_back(rnn_states_prev_->GetOrtTensor());

    const auto rnn_states_type = model_.session_info_.GetOutputDataType(model_.config_->model.decoder.outputs.rnn_states.c_str());
    const std::array<int64_t, 3> rnn_states_shape{3, decoder_input_ids_.GetMarianInputsShape()[0], encoder_hidden_size};
    rnn_states_ = std::make_unique<Tensor>(model_.p_device_inputs_, rnn_states_type);
    rnn_states_->CreateTensor(rnn_states_shape);

    output_names_.push_back(model_.config_->model.decoder.outputs.rnn_states.c_str());
    outputs_.push_back(rnn_states_->GetOrtTensor());
    *past_key_values_length_->GetTensorMutableData<int64_t>() += 1;

    logits_.Add();

    logits_.Update(next_tokens.subspan(next_tokens.size() - 1, 1), 1);

    if (model_.config_->model.decoder.run_options.has_value()) {
      State::SetRunOptions(model_.config_->model.decoder.run_options.value());
    }
    State::Run(*model_.session_decoder_);
    first_run_ = false;
    return logits_.Get();
  }

  // Update the decoder inputs with the next tokens
  decoder_input_ids_.Update(next_tokens);
  UpdateRnnStates(next_indices);

  auto data = past_key_values_length_->GetTensorMutableData<int64_t>();
  *data += 1;

  logits_.Update(next_tokens, 1);

  // Run the decoder
  if (model_.config_->model.decoder.run_options.has_value()) {
    State::SetRunOptions(model_.config_->model.decoder.run_options.value());
  }
  State::Run(*model_.session_decoder_);
  return logits_.Get();
}

void MarianState::UpdateRnnStates(DeviceSpan<int32_t> next_indices) {
  if (next_indices.empty()) {
    rnn_states_prev_->GetByteSpan().CopyFrom(rnn_states_->GetByteSpan());
    return;
  }

  const size_t batch_beam_size = params_->BatchBeamSize();
  auto indices = next_indices.CopyDeviceToCpu();
  if (indices.size() != batch_beam_size) {
    throw std::runtime_error("MarianState::UpdateRnnStates: beam index count does not match batch_size * num_beams");
  }

  auto source = rnn_states_->GetByteSpan();
  auto destination = rnn_states_prev_->GetByteSpan();
  constexpr size_t rnn_state_slots = 3;
  if (source.size() != destination.size() || source.size() % (rnn_state_slots * batch_beam_size) != 0) {
    throw std::runtime_error("MarianState::UpdateRnnStates: recurrent state buffers have incompatible sizes");
  }
  const size_t bytes_per_beam = source.size() / (rnn_state_slots * batch_beam_size);

  for (size_t slot = 0; slot < rnn_state_slots; ++slot) {
    for (size_t destination_beam = 0; destination_beam < batch_beam_size; ++destination_beam) {
      const int32_t source_beam = indices[destination_beam];
      if (source_beam < 0 || static_cast<size_t>(source_beam) >= batch_beam_size) {
        throw std::runtime_error("MarianState::UpdateRnnStates: beam index is out of range");
      }

      auto source_state = source.subspan(
          (slot * batch_beam_size + static_cast<size_t>(source_beam)) * bytes_per_beam,
          bytes_per_beam);
      auto destination_state = destination.subspan(
          (slot * batch_beam_size + destination_beam) * bytes_per_beam,
          bytes_per_beam);
      destination_state.CopyFrom(source_state);
    }
  }
}

}  // namespace Generators
