from tensorflow.keras.models import Model, Input
from tensorflow.keras.layers import Dense, LSTM, GRU, Dropout
from tensorflow.keras.optimizers import Adam
import tensorflow as tf

def create_lstm_reg(input_shape: tuple[int, int],
                    lstm_units_1: int,
                    lstm_units_2: int,
                    # dense_units,
                    eta: float = 0.001,
                    dropout_rate: float = 0.1):
    inputs = Input(shape=input_shape)
    x = LSTM(units=lstm_units_1, return_sequences=True)(inputs)
    x = Dropout(rate=dropout_rate)(x)
    x = LSTM(units=lstm_units_2)(x)
    #     x = Dropout(rate=dropout_rate)(x)
    #     x = Dense(units=dense_units, activation='relu')(x)
    outputs = Dense(units=1)(x)

    model = Model(inputs=inputs, outputs=outputs)

    model.compile(optimizer=Adam(eta),
                  loss='mean_squared_error',
                  metrics=[
                      tf.keras.metrics.RootMeanSquaredError(),
                      tf.keras.metrics.MeanAbsolutePercentageError(name='mape'),
                      tf.keras.metrics.MeanAbsoluteError()
                  ]
                  )
    return model


def create_gru_reg(input_shape, gru_units_1, gru_units_2, dense_units, eta, dropout_rate=0.1):
    inputs = Input(shape=input_shape)
    x = GRU(units=gru_units_1, return_sequences=True)(inputs)
    x = Dropout(rate=dropout_rate)(x)
    x = GRU(units=gru_units_2)(x)
    #     x = Dropout(rate=dropout_rate)(x)
    #     x = Dense(units=dense_units, activation='relu')(x)
    outputs = Dense(units=1)(x)

    model = Model(inputs=inputs, outputs=outputs)

    model.compile(optimizer=Adam(eta),
                  loss='mean_squared_error',
                  metrics=[
                      tf.keras.metrics.RootMeanSquaredError(),
                      tf.keras.metrics.MeanAbsolutePercentageError(name='mape'),
                      tf.keras.metrics.MeanAbsoluteError()
                  ]
                  )
    return model