import numpy as np
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense

# Map characters to integers: {'h':0, 'e':1, 'l':2, 'o':3}
char_to_int = {'h':0, 'e':1, 'l':2, 'o':3}
int_to_char = {i: c for c, i in char_to_int.items()}

# Input sequence: "h", "e", "l"
X = np.array([[ [0], [1], [2] ]])  # shape (samples, timesteps, features)
# Next char to predict: "l"
y = np.array([2])
# One-hot encode output
y = np.eye(4)[y]

# Build model: 3 timesteps, 1 feature per step
model = Sequential([
    LSTM(8, input_shape=(3, 1)),
    Dense(4, activation='softmax')
])
model.compile(loss='categorical_crossentropy', optimizer='adam')

# Train briefly (just to show it runs)
model.fit(X, y, epochs=100, verbose=0)

# Predict what comes after "hel"
pred = model.predict(X)
print("Predicted next char:", int_to_char[np.argmax(pred[0])])