import streamlit as st
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim

# Streamlit app title
st.title('Simple Linear Regression with PyTorch')

# Sidebar sliders for noise and number of data points
noise_level = st.sidebar.slider('Noise Level', 0.0, 1.0, 0.1, step=0.01)
num_points = st.sidebar.slider('Number of Data Points', 10, 100, 50, step=5)
num_epochs = st.sidebar.slider('Number of Epochs', 1, 50, 10, step=1)
learning_rate = st.sidebar.slider('Learning Rate', 0.001, 0.1, 0.01, step=0.001)

# Generate data
np.random.seed(0)
x = np.linspace(0, 10, num_points)
y = 2 * x + 1 + noise_level * np.random.randn(num_points)

# Convert data to PyTorch tensors
x_tensor = torch.tensor(x, dtype=torch.float32).view(-1, 1)
y_tensor = torch.tensor(y, dtype=torch.float32).view(-1, 1)

# Define the linear regression model
class LinearRegressionModel(nn.Module):
    def __init__(self):
        super(LinearRegressionModel, self).__init__()
        self.linear = nn.Linear(1, 1)
    
    def forward(self, x):
        return self.linear(x)

model = LinearRegressionModel()

# Define the loss function and the optimizer
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate)

# Train the model
losses = []
for epoch in range(num_epochs):
    model.train()
    optimizer.zero_grad()
    outputs = model(x_tensor)
    loss = criterion(outputs, y_tensor)
    loss.backward()
    optimizer.step()
    losses.append(loss.item())

# Get the final model parameters
slope = model.linear.weight.item()
intercept = model.linear.bias.item()

# Make predictions
model.eval()
y_pred_tensor = model(x_tensor)
y_pred = y_pred_tensor.detach().numpy()

# Create scatter plot with regression line
fig, ax = plt.subplots()
ax.scatter(x, y, alpha=0.6, label='Data points')
ax.plot(x, y_pred, color='red', label='Regression line')
ax.set_title('Scatter Plot with Noise and Number of Data Points')
ax.set_xlabel('X-axis')
ax.set_ylabel('Y-axis')
ax.legend()

# Display slope and intercept in Streamlit
st.write(f"**Slope:** {slope}")
st.write(f"**Intercept:** {intercept}")

# Display scatter plot in Streamlit
st.pyplot(fig)

# Plot training loss
fig_loss, ax_loss = plt.subplots()
ax_loss.plot(range(num_epochs), losses)
ax_loss.set_title('Training Loss')
ax_loss.set_xlabel('Epoch')
ax_loss.set_ylabel('Loss')

# Display training loss plot in Streamlit
st.pyplot(fig_loss)