library(shiny)
library(ggplot2)

# In this app, the population model is as follows:
# Each outcome Y is a standard normal random variable.
# For each outcome, there are 10 covariates, X_1, ..., X_10.
# Each is obtained from the outcome by the relationship:
# X_j = Y + epsilon_j, 
# where epsilon_j is a standard normal random variable.

# Function to simulate OLS errors for a given number of covariates
simulate_errors <- function(num_covariates, num_universes = 2000, num_samples = 100) {
  errors <- numeric(num_universes)
  
  for (i in 1:num_universes) {
    Y <- rnorm(num_samples, mean = 0, sd = 1)
    X <- matrix(nrow = num_samples, ncol = 10)
    
    for (j in 1:10) {
      epsilon_j <- rnorm(num_samples, mean = 0, sd = 1)
      X[, j] <- Y + epsilon_j
    }
    
    # Use only num_covariates covariates
    model <- lm(Y ~ 0 + X[, 1:num_covariates])
    prediction <- sum(coef(model))
    
    # Can show: Conditional expectation of Y given X is (X_1 + ... + X_10)/11
    # 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:10) {
  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 = 10, 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(-0.75,0.25) +
      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 10 covariates for each observation.",
          "Each observation Y is Normal(0,1), and each covariate X_j = Y + epsilon_j, where",
          "epsilon_j is Normal(0,1).  In other words, each covariate is a noisy measurement of the",
          "corresponding outcome.",
          "",
          "",
          "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)