"""
Lecture 18 -- BiLSTM layer + a toy Encoder-Decoder skeleton (Keras).
Illustrative only: shows how the pieces connect, not a trainable model.
Full trainable seq2seq-with-attention tutorial:
https://www.tensorflow.org/text/tutorials/nmt_with_attention
"""
from tensorflow.keras import layers, Model

VOCAB_SIZE, EMBED_DIM, HIDDEN = 5000, 64, 128

# ---- 1. A Bidirectional LSTM layer (Section 2) ----
inputs = layers.Input(shape=(None,), dtype="int32")
x = layers.Embedding(VOCAB_SIZE, EMBED_DIM, mask_zero=True)(inputs)
bilstm_out = layers.Bidirectional(layers.LSTM(HIDDEN, return_sequences=True))(x)
# bilstm_out shape: (batch, time, 2*HIDDEN) -- forward/backward concatenated
bilstm_model = Model(inputs, bilstm_out, name="bilstm_demo")
bilstm_model.summary()

# ---- 2. A minimal Encoder-Decoder skeleton (Section 4) ----
# Encoder: (Bi)LSTM compresses the source sequence into final states.
enc_inputs = layers.Input(shape=(None,), name="encoder_tokens")
enc_emb = layers.Embedding(VOCAB_SIZE, EMBED_DIM, mask_zero=True)(enc_inputs)
enc_out, fh, fc, bh, bc = layers.Bidirectional(
    layers.LSTM(HIDDEN, return_state=True))(enc_emb)
state_h = layers.Concatenate()([fh, bh])   # H_forward || H_backward
state_c = layers.Concatenate()([fc, bc])   # C_forward || C_backward

# Decoder: unidirectional LSTM, initialised from the encoder's context,
# starts from a <SOS> token (handled upstream in the input pipeline),
# predicts one token at a time via softmax.
dec_inputs = layers.Input(shape=(None,), name="decoder_tokens")  # teacher-forced
dec_emb = layers.Embedding(VOCAB_SIZE, EMBED_DIM, mask_zero=True)(dec_inputs)
dec_lstm = layers.LSTM(2 * HIDDEN, return_sequences=True)
dec_out = dec_lstm(dec_emb, initial_state=[state_h, state_c])
dec_softmax = layers.Dense(VOCAB_SIZE, activation="softmax")(dec_out)

seq2seq = Model([enc_inputs, dec_inputs], dec_softmax, name="encoder_decoder_demo")
seq2seq.summary()