library(shiny)
library(ggplot2)

# In this app, the population model is as follows:
# There are 20 covariates X_1, ..., X_20.  Each is a standard normal random variable.
# Only the first 10 covariates are relevant to the outcome, which follows the linear population model:
# Y = (X_1 + ... + X_10)/11 + epsilon, 
# where epsilon is normally distributed with mean zero and variance 1/11.

# True number of covariates in population model
true_num_covariates = 10
max_num_covariates = 20

# Function to simulate OLS errors for a given number of covariates
simulate_errors <- function(num_covariates, num_universes = 1000, num_samples = 100) {
  errors <- numeric(num_universes)
  
  for (i in 1:num_universes) {
    X <- matrix(rnorm(num_samples*max_num_covariates, mean = 0, sd = 1),
                nrow = num_samples,
                ncol = max_num_covariates)
    Y <- rowSums(X[,1:true_num_covariates])/11 + rnorm(num_samples, mean = 0, sd = sqrt(1/11))

    # Use only num_covariates covariates
    model <- lm(Y ~ 0 + X[, 1:num_covariates])
    prediction <- sum(coef(model))
    
    # Calculate error at the point X = (1,1,...,1) ==> conditional expectation is 10/11
    errors[i] <- prediction - 10/11 
  }
  
  return(errors)
}

# Cached data for varying number of covariates (1 to 10)
cached_results <- list()
for (cov in 1:20) {
  cached_results[[as.character(cov)]] <- simulate_errors(cov)
}

# Define UI
ui <- fluidPage(
  titlePanel("Bias and Variance of OLS"),
  
  sidebarLayout(
    sidebarPanel(
      sliderInput("covariates", "Number of Covariates:", min = 1, max = 20, value = 1)
    ),
    
    mainPanel(
      plotOutput("errorPlot")
    )
  ),
  
  textOutput("explanationText")
)

# Define server logic
server <- function(input, output) {
  
  output$errorPlot <- renderPlot({
    num_covariates <- input$covariates
    errors <- cached_results[[as.character(num_covariates)]]
    
    # Create a data frame for ggplot
    error_df <- data.frame(Universe = 1:length(errors), Error = errors)
    
    # Plot errors
    ggplot(error_df, aes(x = Error)) +
      geom_density(fill = "blue", alpha = 0.5) +  # Density plot
      labs(title = paste("Density of Prediction Errors for", num_covariates, "Covariates"), 
           x = "Prediction Error", 
           y = "Density") +
      xlim(-1,0.5) +
      ylim(0,16) +
      theme_minimal() +
      theme(
        plot.title = element_text(size = 20, face = "bold"),     # Larger title
        axis.title.x = element_text(size = 16),                  # Larger X-axis title
        axis.title.y = element_text(size = 16),                  # Larger Y-axis title
        axis.text.x = element_text(size = 14),                   # Larger X-axis text
        axis.text.y = element_text(size = 14)                    # Larger Y-axis text
      )
  })
  
  # Add explanation text
  output$explanationText <- renderText({
    paste("The density plot shows the distribution of prediction errors for a linear regression model",
          "that is fit using OLS with", input$covariates, "covariates.",
          "In the population model used here, there are 20 covariates, but only 10 are relevant to",
          "the outcome.  In particular, for each observation, X_1, ..., X_20 are Normal(0,1),",
          "independent of each other, and Y = (X_1 + ... + X_10)/11 + epsilon,",
	  "where epsilon ~ Normal(0,1/11).",
          "",
          "Predictions are made at the point X = (1,1,...,1), though similar results would be",
          "obtained for any chosen test point (or RMSE over a test set)."
    )
  })
}

# Run the application 
shinyApp(ui = ui, server = server)