import streamlit as st
import torch
import torch.nn as nn
import torch.optim as optim
from keras.models import Sequential
from keras.layers import Dense
from keras.optimizers import Adam

# Set up Streamlit layout
st.title("PyTorch vs Keras Comparison")

# Define PyTorch model
class PyTorchModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(PyTorchModel, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# Define Keras model
def KerasModel(input_size, hidden_size, output_size):
    model = Sequential()
    model.add(Dense(hidden_size, activation='relu', input_shape=(input_size,)))
    model.add(Dense(output_size))
    return model

# Define example NLP tasks
nlp_tasks = {
    'Task 1: Sentiment Analysis': {
        'PyTorch': {
            'model': PyTorchModel(100, 64, 2),
            'optimizer': optim.Adam,
        },
        'Keras': {
            'model': KerasModel(100, 64, 2),
            'optimizer': Adam,
        }
    },
    'Task 2: Text Classification': {
        'PyTorch': {
            'model': PyTorchModel(200, 128, 5),
            'optimizer': optim.SGD,
        },
        'Keras': {
            'model': KerasModel(200, 128, 5),
            'optimizer': Adam,
        }
    }
}

# Select NLP task
task = st.sidebar.selectbox("Select NLP Task", list(nlp_tasks.keys()))

# Select framework
framework = st.sidebar.selectbox("Select Framework", ['PyTorch', 'Keras'])

# Get model and optimizer for selected task and framework
model = nlp_tasks[task][framework]['model']
optimizer = nlp_tasks[task][framework]['optimizer']

# Display model summary
st.subheader(f"{framework} Model Summary")
st.text(model)

# Display optimizer details
st.subheader(f"{framework} Optimizer Details")
st.text(optimizer)

# Perform example computations
if st.button("Perform Computation"):
    # Perform forward pass
    input_data = torch.randn(1, model.fc1.in_features)
    output = model(input_data)
    st.write(f"Output: {output}")

    # Perform backward pass
    loss = output.mean()
    optimizer = optimizer(model.parameters(), lr=0.01)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    st.write("Backward pass completed.")