library(shiny)
library(ggplot2)
library(dplyr)
library(tidyr)
library(DT)
library(ggformula)
ui <- fluidPage(
  wellPanel(
    h2("Time-independent confounder / DiD plots"),
    
    textInput('vec1', 'Effect sizes (comma delimited)', "1,2,3"),
    numericInput('order', 'Order of the timeseries', 2),
    sliderInput("xvalue", "X-axis", min = 1, max = 50, value = 30, step = 0.01),
    sliderInput("yvalue", "Y-axis", min = 0, max = 50, value = 30, step = 0.01),
    
    textInput('legendInput', 'Legendtitle: Label 1, Label 2, Label 3', 
              "Week: Week 1, Week 2, Week 3"),
    
    selectInput("plotType", "Plot Type", 
                choices = c("Confounder", "Difference-in-Difference")),
    
    DTOutput("myTable"),
    uiOutput("dynamicPlotUI")
  )
)

server <- function(input, output, session) {
  
  observe({
    vec1_values <- as.numeric(unlist(strsplit(input$vec1, ",")))
    max_val <- max(vec1_values, na.rm = TRUE)
    denom <- ifelse(input$xvalue < 2, input$xvalue - 1, 1)
    new_max <- (max_val * 10) / denom
    updateSliderInput(session, "yvalue", max = new_max)
  })
  
  parseLegend <- reactive({
    parts <- unlist(strsplit(input$legendInput, ":"))
    if (length(parts) < 2) return(list(title = "Legend", labels = NULL))
    list(title = trimws(parts[1]), labels = trimws(unlist(strsplit(parts[2], ","))))
  })
  
  # ---- Confounder plots ----
  MeanPlot <- reactive({
    EF <- as.numeric(unlist(strsplit(input$vec1, ",")))
    meanEF <- mean(EF)
    UX <- seq(1001, 1000 * input$xvalue, length.out = 1899) / 1000
    UY <- meanEF / (UX - 1)
    ggplot(data.frame(UX, UY), aes(x = UX, y = UY)) +
      geom_line() + ggtitle("Mean approach") +
      xlab("Exposure-confounder association") + ylab("Outcome-confounder association") +
      xlim(1, input$xvalue) + ylim(0, input$yvalue)
  })
  
  TSPlot <- function(order) {
    EF <- as.numeric(unlist(strsplit(input$vec1, ",")))
    d <- order
    EF_pad <- c(rep(0, d), EF)
    ts_bias <- diff(EF_pad, differences = d)
    UX <- seq(1001, 1000 * input$xvalue, length.out = 1899) / 1000
    avg_bias <- sapply(1:length(ts_bias), function(i) sum(ts_bias[1:i]) / i)
    matrix_data <- sapply(1:length(ts_bias), function(i) avg_bias[i] / (UX - 1))
    inverse_nabla <- function(x) { for(i in 1:d) x <- cumsum(x); x }
    df <- as.data.frame(t(apply(matrix_data, 1, inverse_nabla)))
    df$UX <- UX
    colnames(df)[1:(ncol(df)-1)] <- as.character(1:(ncol(df)-1))
    df_long <- pivot_longer(df, -UX, names_to = "Series", values_to = "Value")
    df_long$Series <- factor(df_long$Series, levels = unique(df_long$Series))
    legend_info <- parseLegend()
    n_labels <- length(unique(df_long$Series))
    if (!is.null(legend_info$labels) && length(legend_info$labels) >= n_labels)
      levels(df_long$Series) <- legend_info$labels[1:n_labels]
    ggplot(df_long, aes(x = UX, y = Value, color = Series)) +
      geom_line() +
      xlab("Exposure-confounder association") +
      ylab("Outcome-confounder association") +
      xlim(1, input$xvalue) + ylim(0, input$yvalue) +
      labs(color = legend_info$title)
  }
  
  # ---- Difference-in-Difference plots ----
  DiDPlot <- function(order) {
    values <- as.numeric(unlist(strsplit(input$vec1, ",")))
    df <- data.frame(
      Period = seq_along(values),
      Value2 = values
    )
    diff_n <- function(x, n) { for(i in seq_len(n)) x <- diff(x); x }
    plots_data <- lapply(order, function(n) {
      data.frame(
        Diff_Order = n,
        Period = df$Period[(n+1):nrow(df)],
        Diff_Value = diff_n(df$Value2, n)
      )
    }) %>% bind_rows()
    df_plot <- filter(plots_data, Diff_Order == order)
    ggplot(df_plot, aes(x = Period, y = Diff_Value)) +
      geom_point() +
      geom_line() +
      stat_spline(df=length(df_plot$Diff_Value)/4, color = "blue") +
      theme_minimal() +
      labs(title = paste(order, "Order Difference (Spline)"),
           x = "Period",
           y = paste0("Difference (Order ", order, ")")) +
      theme(axis.text.x = element_text(angle = 45, hjust = 1))
  }
  
  # ---- Dynamic UI ----
  output$dynamicPlotUI <- renderUI({
    if (input$plotType == "Confounder") {
      plot_name <- if (input$order == 0) "Meanapproach" else if (input$order == 1) "Levyapproach" else "Timeseries"
      tagList(
        plotOutput(plot_name),
        downloadButton("downloadPlot", "Download Confounder Plot")
      )
    } else {
      tagList(
        plotOutput("DiDplot"),
        downloadButton("downloadPlot", "Download DiD Plot")
      )
    }
  })
  
  # ---- Render plots ----
  output$Meanapproach <- renderPlot({ MeanPlot() })
  output$Levyapproach <- renderPlot({ TSPlot(order = 1) })
  output$Timeseries <- renderPlot({ TSPlot(order = input$order) })
  output$DiDplot <- renderPlot({ DiDPlot(order = input$order) })
  
  # ---- Download handler for both plot types ----
  output$downloadPlot <- downloadHandler(
    filename = function() {
      if (input$plotType == "Confounder") {
        paste0("Confounder_order_", input$order, ".pdf")
      } else {
        paste0("DiD_order_", input$order, ".pdf")
      }
    },
    content = function(file) {
      p <- if (input$plotType == "Confounder") {
        if (input$order == 0) MeanPlot() else if (input$order == 1) TSPlot(order = 1) else TSPlot(order = input$order)
      } else {
        DiDPlot(order = input$order)
      }
      ggsave(file, plot = p, device = "pdf", width = 6, height = 4)
    }
  )
  
  # ---- Table ----
  output$myTable <- renderDT({
    data.frame(Time = 1:length(unlist(strsplit(input$vec1, ","))),
               Effect = as.numeric(unlist(strsplit(input$vec1, ","))))
  })
  
}

shinyApp(ui, server)
