# Set seed for reproducibility
set.seed(42)

# Generate training data: y = 3x + 2 + noise
x <- torch_randn(100, 1)
y <- 3 * x + 2 + torch_randn(100, 1) * 0.3

# Display the first few data points
head(
  data.frame(
    x = as.numeric(x$squeeze()),
    y = as.numeric(y$squeeze())
  ))

            x         y
1 -0.02329975  1.861628
2  1.92341769  7.555232
3  0.11041667  2.613283
4 -2.55959392 -5.931758
5  0.36482519  3.099005
6  0.97125226  4.551073

# Define a simple feedforward neural network
nnet <- nn_module(
  initialize = function() {
    # Define layers
    self$layer1 <- nn_linear(1, 8)  # Input layer to hidden layer (1 -> 8 neurons)
    self$layer2 <- nn_linear(8, 1)  # Hidden layer to output layer (8 -> 1 neuron)
  },
  forward = function(x) {
    # Define forward pass
    x %>% 
      self$layer1() %>%     # First linear transformation
      nnf_relu() %>%     # ReLU activation function
      self$layer2()         # Second linear transformation
  }
)

# Instantiate the model
model <- nnet()

# Display model structure
print(model)

An `nn_module` containing 25 parameters.

── Modules ─────────────────────────────────────────────────────────────────────
• layer1: <nn_linear> #16 parameters
• layer2: <nn_linear> #9 parameters

# Set up optimizer (Adam optimizer with learning rate 0.02)
optimizer <- optim_adam(model$parameters, lr = 0.02)

# Define loss function (Mean Squared Error for regression)
loss_fn <- nnf_mse_loss

# Store loss values for plotting
loss_history <- numeric(300)

# Training loop
for(epoch in 1:300) {
  
  # Set model to training mode
  model$train()
  
  # Reset gradients
  optimizer$zero_grad()
  
  # Forward pass
  y_pred <- model(x)
  
  # Calculate loss
  loss <- loss_fn(y_pred, y)
  
  # Backward pass
  loss$backward()
  
  # Update parameters
  optimizer$step()
  
  # Store loss for plotting
  loss_history[epoch] <- loss$item()
}

# Create a data frame for plotting
training_df <- data.frame(
  epoch = 1:300,
  loss = loss_history
)

# Plot training loss
ggplot(training_df, aes(x = epoch, y = loss)) +
  geom_line(color = "#2c3e50", size = 1) +
  labs(
    title = "Training Loss Over Time",
    subtitle = "Neural Network Learning Progress",
    x = "Epoch",
    y = "Mean Squared Error Loss"
  ) +
  theme_minimal() +
  theme(
    plot.title = element_text(size = 14, face = "bold"),
    plot.subtitle = element_text(size = 12, color = "gray60")
  )

# Set model to evaluation mode
model$eval()

# Generate predictions
with_no_grad({
  y_pred <- model(x)
})

# Convert to R vectors for plotting
x_np <- as.numeric(x$squeeze())
y_np <- as.numeric(y$squeeze())
y_pred_np <- as.numeric(y_pred$squeeze())

# Create data frame for ggplot
plot_df <- data.frame(
  x = x_np,
  y_actual = y_np,
  y_predicted = y_pred_np
)

# Create the plot
ggplot(plot_df, aes(x = x)) +
  geom_point(aes(y = y_actual, color = "Actual"), alpha = 0.7, size = 2) +
  geom_point(aes(y = y_predicted, color = "Predicted"), alpha = 0.7, size = 2) +
  geom_smooth(aes(y = y_predicted), method = "loess", se = FALSE, 
              color = "#e74c3c", linetype = "dashed") +
  labs(
    title = "Neural Network Regression Results",
    subtitle = "Comparing actual vs predicted values",
    x = "Input (x)",
    y = "Output (y)",
    color = "Data Type"
  ) +
  scale_color_manual(values = c("Actual" = "#3498db", "Predicted" = "#e74c3c")) +
  theme_minimal() +
  theme(
    plot.title = element_text(size = 14, face = "bold"),
    plot.subtitle = element_text(size = 12, color = "gray60"),
    legend.position = "top"
  )

# Calculate performance metrics
mse <- mean((y_pred_np - y_np)^2)
rmse <- sqrt(mse)
mae <- mean(abs(y_pred_np - y_np))
r_squared <- cor(y_pred_np, y_np)^2

# Create performance summary
performance_summary <- data.frame(
  Metric = c("Mean Squared Error", "Root Mean Squared Error", 
             "Mean Absolute Error", "R-squared"),
  Value = c(mse, rmse, mae, r_squared)
)

print(performance_summary)

                   Metric      Value
1      Mean Squared Error 0.09061213
2 Root Mean Squared Error 0.30101848
3     Mean Absolute Error 0.23722124
4               R-squared 0.99000990

# Compare with true relationship (y = 3x + 2)
# Generate predictions on a grid for comparison
x_grid <- torch_linspace(-3, 3, 100)$unsqueeze(2)
with_no_grad({
  y_grid_pred <- model(x_grid)
})

x_grid_np <- as.numeric(x_grid$squeeze())
y_grid_pred_np <- as.numeric(y_grid_pred$squeeze())
y_grid_true <- 3 * x_grid_np + 2

# Plot comparison
comparison_df <- data.frame(
  x = x_grid_np,
  y_true = y_grid_true,
  y_predicted = y_grid_pred_np
)

ggplot(comparison_df, aes(x = x)) +
  geom_line(aes(y = y_true, color = "True Function"), size = 2) +
  geom_line(aes(y = y_predicted, color = "Neural Network"), size = 2, linetype = "dashed") +
  geom_point(data = plot_df, aes(y = y_actual), alpha = 0.3, color = "gray50") +  labs(
    title = "Neural Network vs True Function",
    subtitle = "Model learning assessment against the underlying pattern",
    x = "Input (x)",
    y = "Output (y)",
    color = "Function Type"
  ) +
  scale_color_manual(values = c("True Function" = "#2c3e50", "Neural Network" = "#e74c3c")) +
  theme_minimal() +
  theme(
    plot.title = element_text(size = 14, face = "bold"),
    plot.subtitle = element_text(size = 12, color = "gray60"),
    legend.position = "top"
  )