#include "engine/framework/modules/recurrent_modules.h"

#include "engine/framework/modules/activation_modules.h"

#include "engine/framework/modules/linear_module.h"

#include "engine/framework/modules/primitive_modules.h"

#include "engine/framework/modules/structural_modules.h"

#include <stdexcept>

#include <vector>

namespace engine::modules {

namespace {

const core::ModulePortSpec kLstmInputs[] = {

{"input", core::PortKind::Activation, false},

{"hidden", core::PortKind::Activation, false},

{"cell", core::PortKind::Activation, false},

{"weight_ih", core::PortKind::Parameter, false},

{"weight_hh", core::PortKind::Parameter, false},

{"bias_ih", core::PortKind::Parameter, false},

{"bias_hh", core::PortKind::Parameter, false},

};

const core::ModulePortSpec kLstmOutputs[] = {

{"hidden", core::PortKind::Activation, false},

{"cell", core::PortKind::Activation, false},

};

const core::ModuleSchema kLstmSchema = {

"LSTMCell",

"nn.recurrent",

kLstmInputs,

7,

kLstmOutputs,

2,

"Applies a single LSTM cell update over [batch, features] inputs.",

};

const core::ModulePortSpec kLstmSequenceInputs[] = {

{"input", core::PortKind::Activation, false},

{"initial_hidden", core::PortKind::Activation, false},

{"initial_cell", core::PortKind::Activation, false},

{"weight_ih", core::PortKind::Parameter, false},

{"weight_hh", core::PortKind::Parameter, false},

{"bias_ih", core::PortKind::Parameter, false},

{"bias_hh", core::PortKind::Parameter, false},

};

const core::ModulePortSpec kLstmSequenceOutputs[] = {

{"sequence", core::PortKind::Activation, false},

{"hidden", core::PortKind::Activation, false},

{"cell", core::PortKind::Activation, false},

};

const core::ModuleSchema kLstmSequenceSchema = {

"LSTMSequence",

"nn.recurrent",

kLstmSequenceInputs,

7,

kLstmSequenceOutputs,

3,

"Applies an LSTM over a full [time, features] sequence.",

};

const core::ModulePortSpec kBidirLstmInputs[] = {

{"input", core::PortKind::Activation, false},

{"initial_forward_hidden", core::PortKind::Activation, false},

{"initial_forward_cell", core::PortKind::Activation, false},

{"initial_reverse_hidden", core::PortKind::Activation, false},

{"initial_reverse_cell", core::PortKind::Activation, false},

{"forward_weight_ih", core::PortKind::Parameter, false},

{"forward_weight_hh", core::PortKind::Parameter, false},

{"forward_bias_ih", core::PortKind::Parameter, false},

{"forward_bias_hh", core::PortKind::Parameter, false},

{"reverse_weight_ih", core::PortKind::Parameter, false},

{"reverse_weight_hh", core::PortKind::Parameter, false},

{"reverse_bias_ih", core::PortKind::Parameter, false},

{"reverse_bias_hh", core::PortKind::Parameter, false},

};

const core::ModulePortSpec kBidirLstmOutputs[] = {

{"sequence", core::PortKind::Activation, false},

{"forward_hidden", core::PortKind::Activation, false},

{"forward_cell", core::PortKind::Activation, false},

{"reverse_hidden", core::PortKind::Activation, false},

{"reverse_cell", core::PortKind::Activation, false},

};

const core::ModuleSchema kBidirLstmSchema = {

"BidirectionalLSTM",

"nn.recurrent",

kBidirLstmInputs,

13,

kBidirLstmOutputs,

5,

"Applies forward and reverse LSTM passes over a full [time, features] sequence and concatenates outputs.",

};

void validate_weight_shapes(const LSTMCellConfig & config, const LSTMCellWeights & weights) {

const int64_t gates = 4 * config.hidden_size;

core::validate_shape(weights.weight_ih, core::TensorShape::from_dims({gates, config.input_size}), "weight_ih");

core::validate_shape(weights.weight_hh, core::TensorShape::from_dims({gates, config.hidden_size}), "weight_hh");

core::validate_shape(weights.bias_ih, core::TensorShape::from_dims({gates}), "bias_ih");

core::validate_shape(weights.bias_hh, core::TensorShape::from_dims({gates}), "bias_hh");

}

core::TensorValue gate_slice(

core::ModuleBuildContext & ctx,

const core::TensorValue & gates,

int64_t hidden_size,

int64_t gate_index) {

return SliceModule({1, gate_index * hidden_size, hidden_size}).build(ctx, gates);

}

void validate_sequence_config(const LSTMSequenceConfig & config) {

if (config.input_size <= 0 || config.hidden_size <= 0) {

throw std::runtime_error("LSTMSequenceConfig dimensions must be positive");

}

}

LSTMCellOutputs build_lstm_cell_from_projected_input(

core::ModuleBuildContext & ctx,

const LSTMCellConfig & config,

const core::TensorValue & projected_input,

const core::TensorValue & hidden,

const core::TensorValue & cell,

const LSTMCellWeights & weights) {

const int64_t gates = 4 * config.hidden_size;

core::validate_shape(projected_input, core::TensorShape::from_dims({hidden.shape.dims[0], gates}), "projected_input");

core::validate_shape(hidden, core::TensorShape::from_dims({hidden.shape.dims[0], config.hidden_size}), "hidden");

core::validate_shape(cell, core::TensorShape::from_dims({hidden.shape.dims[0], config.hidden_size}), "cell");

validate_weight_shapes(config, weights);

const auto projected_hidden = LinearModule({config.hidden_size, gates, true}).build(

ctx,

hidden,

{weights.weight_hh, weights.bias_hh});

const auto gate_values = AddModule().build(ctx, projected_input, projected_hidden);

const auto input_gate = SigmoidModule().build(ctx, gate_slice(ctx, gate_values, config.hidden_size, 0));

const auto forget_gate = SigmoidModule().build(ctx, gate_slice(ctx, gate_values, config.hidden_size, 1));

const auto candidate = TanhModule().build(ctx, gate_slice(ctx, gate_values, config.hidden_size, 2));

const auto output_gate = SigmoidModule().build(ctx, gate_slice(ctx, gate_values, config.hidden_size, 3));

const auto kept_cell = MulModule().build(ctx, forget_gate, cell);

const auto written_cell = MulModule().build(ctx, input_gate, candidate);

const auto new_cell = AddModule().build(ctx, kept_cell, written_cell);

const auto new_hidden = MulModule().build(ctx, output_gate, TanhModule().build(ctx, new_cell));

return {new_hidden, new_cell};

}

LSTMSequenceOutputs build_lstm_sequence(

core::ModuleBuildContext & ctx,

const LSTMSequenceConfig & config,

const core::TensorValue & input,

const core::TensorValue & initial_hidden,

const core::TensorValue & initial_cell,

const LSTMCellWeights & weights) {

if (ctx.ggml == nullptr) {

throw std::runtime_error("ModuleBuildContext.ggml is null");

}

core::validate_shape(input, core::TensorShape::from_dims({input.shape.dims[0], config.input_size}), "input");

core::validate_shape(initial_hidden, core::TensorShape::from_dims({1, config.hidden_size}), "initial_hidden");

core::validate_shape(initial_cell, core::TensorShape::from_dims({1, config.hidden_size}), "initial_cell");

validate_weight_shapes({config.input_size, config.hidden_size}, weights);

const int64_t frames = input.shape.dims[0];

const int64_t gates = 4 * config.hidden_size;

const auto projected_inputs = LinearModule({config.input_size, gates, true}).build(

ctx,

input,

{weights.weight_ih, weights.bias_ih});

std::vector<core::TensorValue> steps(static_cast<size_t>(frames));

auto hidden = initial_hidden;

auto cell = initial_cell;

const ConcatModule concat_rows({0});

for (int64_t step = 0; step < frames; ++step) {

const int64_t t = config.reverse ? (frames - 1 - step) : step;

const auto projected_x_t = SliceModule({0, t, 1}).build(ctx, projected_inputs);

const auto outputs = build_lstm_cell_from_projected_input(

ctx,

{config.input_size, config.hidden_size},

projected_x_t,

hidden,

cell,

weights);

hidden = outputs.hidden;

cell = outputs.cell;

steps[static_cast<size_t>(t)] = hidden;

}

auto sequence = steps[0];

for (int64_t t = 1; t < frames; ++t) {

sequence = concat_rows.build(ctx, sequence, steps[static_cast<size_t>(t)]);

}

return LSTMSequenceOutputs{sequence, hidden, cell};

}

} // namespace

LSTMCellModule::LSTMCellModule(LSTMCellConfig config) : config_(config) {

if (config_.input_size <= 0 || config_.hidden_size <= 0) {

throw std::runtime_error("LSTMCellConfig dimensions must be positive");

}

}

const LSTMCellConfig & LSTMCellModule::config() const noexcept {

return config_;

}

const core::ModuleSchema & LSTMCellModule::schema() const noexcept {

return static_schema();

}

LSTMCellOutputs LSTMCellModule::build(

core::ModuleBuildContext & ctx,

const core::TensorValue & input,

const core::TensorValue & hidden,

const core::TensorValue & cell,

const LSTMCellWeights & weights) const {

if (ctx.ggml == nullptr) {

throw std::runtime_error("ModuleBuildContext.ggml is null");

}

core::validate_rank_between(input, 2, 2, "input");

core::validate_shape(hidden, core::TensorShape::from_dims({input.shape.dims[0], config_.hidden_size}), "hidden");

core::validate_shape(cell, core::TensorShape::from_dims({input.shape.dims[0], config_.hidden_size}), "cell");

core::validate_shape(input, core::TensorShape::from_dims({input.shape.dims[0], config_.input_size}), "input");

validate_weight_shapes(config_, weights);

const int64_t gates = 4 * config_.hidden_size;

const auto projected_input = LinearModule({config_.input_size, gates, true}).build(

ctx,

input,

{weights.weight_ih, weights.bias_ih});

return build_lstm_cell_from_projected_input(ctx, config_, projected_input, hidden, cell, weights);

}

const core::ModuleSchema & LSTMCellModule::static_schema() noexcept {

return kLstmSchema;

}

LSTMSequenceModule::LSTMSequenceModule(LSTMSequenceConfig config) : config_(config) {

validate_sequence_config(config_);

}

const LSTMSequenceConfig & LSTMSequenceModule::config() const noexcept {

return config_;

}

const core::ModuleSchema & LSTMSequenceModule::schema() const noexcept {

return static_schema();

}

LSTMSequenceOutputs LSTMSequenceModule::build(

core::ModuleBuildContext & ctx,

const core::TensorValue & input,

const core::TensorValue & initial_hidden,

const core::TensorValue & initial_cell,

const LSTMSequenceWeights & weights) const {

return build_lstm_sequence(ctx, config_, input, initial_hidden, initial_cell, weights.cell);

}

const core::ModuleSchema & LSTMSequenceModule::static_schema() noexcept {

return kLstmSequenceSchema;

}

BidirectionalLSTMModule::BidirectionalLSTMModule(LSTMSequenceConfig config) : config_(config) {

validate_sequence_config(config_);

}

const LSTMSequenceConfig & BidirectionalLSTMModule::config() const noexcept {

return config_;

}

const core::ModuleSchema & BidirectionalLSTMModule::schema() const noexcept {

return static_schema();

}

BidirectionalLSTMOutputs BidirectionalLSTMModule::build(

core::ModuleBuildContext & ctx,

const core::TensorValue & input,

const core::TensorValue & initial_forward_hidden,

const core::TensorValue & initial_forward_cell,

const core::TensorValue & initial_reverse_hidden,

const core::TensorValue & initial_reverse_cell,

const BidirectionalLSTMWeights & weights) const {

LSTMSequenceConfig forward_config = config_;

forward_config.reverse = false;

LSTMSequenceConfig reverse_config = config_;

reverse_config.reverse = true;

const auto forward = build_lstm_sequence(ctx, forward_config, input, initial_forward_hidden, initial_forward_cell, weights.forward);

const auto reverse = build_lstm_sequence(ctx, reverse_config, input, initial_reverse_hidden, initial_reverse_cell, weights.reverse);

const auto sequence = ConcatModule({1}).build(ctx, forward.sequence, reverse.sequence);

return {sequence, forward.hidden, forward.cell, reverse.hidden, reverse.cell};

}

const core::ModuleSchema & BidirectionalLSTMModule::static_schema() noexcept {

return kBidirLstmSchema;

}

} // namespace engine::modules