import torch
import torch.nn as nn
import numpy as np
import gradio as gr

SEQUENCE_LENGTH = 10
INPUT_SIZE = 1
OUTPUT_SIZE = 1
HIDDEN_UNITS = 128
device = torch.device('cpu')

class Seq2Seq(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, seq_len):
        super(Seq2Seq, self).__init__()
        self.seq_len = seq_len
        self.hidden_size = hidden_size
        self.encoder_lstm = nn.LSTM(input_size, hidden_size, batch_first=True)
        self.decoder_lstm = nn.LSTM(hidden_size, hidden_size, batch_first=True)
        self.decoder_linear = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        _, (hidden, cell) = self.encoder_lstm(x)
        context_vector = hidden.permute(1, 0, 2)
        decoder_input = context_vector.repeat(1, self.seq_len, 1)
        decoder_output, _ = self.decoder_lstm(decoder_input, (hidden, cell))
        prediction = self.decoder_linear(decoder_output)
        return prediction

model_path = 'seq2seq_model_weights.pth'
model = Seq2Seq(INPUT_SIZE, HIDDEN_UNITS, OUTPUT_SIZE, SEQUENCE_LENGTH).to(device)
model.load_state_dict(torch.load(model_path, map_location=device))
model.eval()

def predict_sequence(input_text):
    try:
        numbers = [float(n.strip()) for n in input_text.split(',')]
        
        if len(numbers) != SEQUENCE_LENGTH:
            return f"Error: Please enter exactly {SEQUENCE_LENGTH} numbers, separated by commas."
        
        input_array = np.array(numbers).reshape(1, SEQUENCE_LENGTH, 1)
        input_tensor = torch.from_numpy(input_array).float().to(device)
        
        with torch.no_grad():
            prediction_tensor = model(input_tensor)
            
        output_array = prediction_tensor.cpu().numpy().flatten()
        
        output_text = ", ".join([f"{n:.1f}" for n in output_array])
        
        return output_text
        
    except Exception as e:
        return f"An error occurred: {str(e)}"

demo = gr.Interface(
    fn=predict_sequence,
    inputs=gr.Textbox(
        label="Input Sequence", 
        placeholder=f"Enter {SEQUENCE_LENGTH} numbers, e.g., 1, 2, 3, 4, 5, 6, 7, 8, 9, 10"
    ),
    outputs=gr.Textbox(label="Predicted Sequence"),
    title="Q11: Seq2Seq Model (n -> n+1)",
    description="Enter a sequence of 10 numbers to predict the next sequence.",
    allow_flagging="never"
)

demo.launch()