# train_model.py
import numpy as np
from keras.models import Sequential
from keras.layers import LSTM, Dense
from keras.utils import to_categorical
from keras.callbacks import ModelCheckpoint

# Load and preprocess data
with open("poem_data.txt", "r", encoding='utf-8') as file:
    text = file.read().lower()

chars = sorted(list(set(text)))
char_to_idx = {c: i for i, c in enumerate(chars)}
idx_to_char = {i: c for i, c in enumerate(chars)}

seq_length = 40
step = 1
X = []
y = []

for i in range(0, len(text) - seq_length, step):
    seq = text[i:i + seq_length]
    label = text[i + seq_length]
    X.append([char_to_idx[c] for c in seq])
    y.append(char_to_idx[label])

X = np.array(X)
y = to_categorical(y, num_classes=len(chars))

# Build the model
model = Sequential()
model.add(LSTM(128, input_shape=(seq_length, 1)))
model.add(Dense(len(chars), activation='softmax'))
model.compile(loss='categorical_crossentropy', optimizer='adam')

# Reshape and normalize
X = np.reshape(X, (X.shape[0], seq_length, 1)) / float(len(chars))

# Train the model
model.fit(X, y, batch_size=128, epochs=20)

# Save model
model.save("model.keras")