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

# Config
WINDOW_SIZE = 3
SEQUENCE_LIMIT = 1000

# Prepare dataset from 1 to 1000
seq = np.array(range(1, SEQUENCE_LIMIT + 1))
x, y = [], []
for i in range(len(seq) - WINDOW_SIZE):
    x.append(seq[i:i + WINDOW_SIZE])
    y.append(seq[i + WINDOW_SIZE])
x = np.array(x).reshape(-1, WINDOW_SIZE, 1)
y = np.array(y)

# Build and train the RNN model
model = Sequential([
    LSTM(64, input_shape=(WINDOW_SIZE, 1)),
    Dense(1)
])
model.compile(optimizer="adam", loss="mse")
model.fit(x, y, epochs=1000, batch_size=512, verbose=0)  # use 1 epoch for faster Hugging Face init

# Prediction function
def predict_next(n1, n2, n3):
    try:
        # Input validation
        for val in [n1, n2, n3]:
            if not (1 <= val <= SEQUENCE_LIMIT):
                return f"Input numbers must be between 1 and {SEQUENCE_LIMIT}"
        
        input_seq = np.array([n1, n2, n3]).reshape(1, 3, 1)
        prediction = model.predict(input_seq, verbose=0)[0][0]
        return f"Predicted next number: {prediction:.2f}"
    
    except Exception as e:
        return f"Error: {str(e)}"

# Gradio interface
demo = gr.Interface(
    fn=predict_next,
    inputs=[
        gr.Number(label="1st Number (≤ 1000)"),
        gr.Number(label="2nd Number (≤ 1000)"),
        gr.Number(label="3rd Number (≤ 1000)")
    ],
    outputs="text",
    title="RNN Number Sequence Predictor",
    description="Enter any 3 increasing numbers (up to 1000). Model predicts the next number.",
)

demo.launch()