import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, Model

# Vocabulary and helpers
vocab = ['<pad>', '<sos>', '<eos>'] + list("ABCDE")
vocab_size = len(vocab)
char_to_idx = {ch: idx for idx, ch in enumerate(vocab)}
idx_to_char = {idx: ch for ch, idx in char_to_idx.items()}

def encode_seq(seq):
    # Add start (<sos>) and end (<eos>) tokens
    return [char_to_idx['<sos>']] + [char_to_idx[c] for c in seq] + [char_to_idx['<eos>']]

def pad_seq(seq, max_len):
    # Pad with '<pad>' (index 0) to uniform length
    return seq + [char_to_idx['<pad>']] * (max_len - len(seq))

inputs = ['A', 'AB', 'ABC', 'ABCD', 'ABCDE']
input_seqs = [encode_seq(seq) for seq in inputs]
target_seqs = [encode_seq(seq[::-1]) for seq in inputs]  # reversed

max_len = max(len(seq) for seq in input_seqs + target_seqs)
X = np.array([pad_seq(seq, max_len) for seq in input_seqs])
Y = np.array([pad_seq(seq, max_len) for seq in target_seqs])

# Model parameters
embed_dim = 8
hidden_dim = 16

# Encoder
encoder_inputs = layers.Input(shape=(max_len,))
encoder_embedding = layers.Embedding(input_dim=vocab_size, output_dim=embed_dim, mask_zero=True)
x = encoder_embedding(encoder_inputs)
encoder_lstm = layers.LSTM(hidden_dim, return_state=True)
encoder_outputs, state_h, state_c = encoder_lstm(x)
encoder_states = [state_h, state_c]

# Decoder
decoder_inputs = layers.Input(shape=(max_len,))
decoder_embedding = layers.Embedding(input_dim=vocab_size, output_dim=embed_dim, mask_zero=True)
x_dec = decoder_embedding(decoder_inputs)
decoder_lstm = layers.LSTM(hidden_dim, return_sequences=True, return_state=True)
decoder_outputs, _, _ = decoder_lstm(x_dec, initial_state=encoder_states)
decoder_dense = layers.TimeDistributed(layers.Dense(vocab_size, activation="softmax"))
decoder_outputs = decoder_dense(decoder_outputs)

# Assemble model
model = Model([encoder_inputs, decoder_inputs], decoder_outputs)
model.compile(optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"])

# Teacher forcing setup
Y_decoder_input = np.array([pad_seq([char_to_idx['<sos>']] + seq[1:], max_len) for seq in target_seqs])
Y_decoder_target = np.expand_dims(Y, -1)  # shape (batch, seq, 1)

# Train
model.fit([X, Y_decoder_input], Y_decoder_target, epochs=300, verbose=0)

# Inference setup: encoder model
encoder_model = Model(encoder_inputs, encoder_states)

# Inference decoder: predict one token at a time
decoder_state_input_h = layers.Input(shape=(hidden_dim,))
decoder_state_input_c = layers.Input(shape=(hidden_dim,))
decoder_states_inputs = [decoder_state_input_h, decoder_state_input_c]
dec_emb_inf_input = layers.Input(shape=(1,))
dec_emb_inf = decoder_embedding(dec_emb_inf_input)
dec_outputs_inf, state_h_inf, state_c_inf = decoder_lstm(dec_emb_inf, initial_state=decoder_states_inputs)
dec_outputs_inf = decoder_dense(dec_outputs_inf)
decoder_model = Model([dec_emb_inf_input] + decoder_states_inputs, [dec_outputs_inf, state_h_inf, state_c_inf])

def decode_sequence(input_seq):
    # Encode the input as state vectors
    states_value = encoder_model.predict(input_seq, verbose=0)
    # Initial target sequence is '<sos>'
    target_seq = np.array([[char_to_idx['<sos>']]])
    decoded = []
    for _ in range(max_len):
        output_tokens, h, c = decoder_model.predict([target_seq] + states_value, verbose=0)
        sampled_token_index = np.argmax(output_tokens[0, -1, :])
        sampled_char = idx_to_char[sampled_token_index]
        if sampled_char == '<eos>' or len(decoded) > max_len:
            break
        decoded.append(sampled_char)
        target_seq = np.array([[sampled_token_index]])
        states_value = [h, c]
    return ''.join(decoded)

# Run a few test sequences
for test_seq in ['A', 'AB', 'ABC']:
    inp = np.array([pad_seq(encode_seq(test_seq), max_len)])
    print(f"Input: {test_seq} -> Predicted: {decode_sequence(inp)}")