# setwd('~/Dropbox/ImageSeq/')

library(shiny)
library(dplyr)
library(plotly)
library(fields)  # For image.plot in heatMap
library(akima)   # For interpolation in heatMap

# Load the data from sm.csv
sm <- read.csv("sm.csv")

# Define function to convert to numeric
f2n <- function(x) as.numeric(as.character(x))

# Compute MaxImageDimsLeft and MaxImageDimsRight from MaxImageDims
sm$MaxImageDimsLeft <- unlist(lapply(strsplit(sm$MaxImageDims, split = "_"), function(x) sort(f2n(x))[1]))
sm$MaxImageDimsRight <- unlist(lapply(strsplit(sm$MaxImageDims, split = "_"), function(x) sort(f2n(x))[2]))

# Define the heatMap function (unchanged except for updated default color palette)
heatMap <- function(x, y, z,
                    main = "",
                    N, yaxt = NULL,
                    xlab = "",
                    ylab = "",
                    horizontal = FALSE,
                    useLog = "",
                    legend.width = 1,
                    ylim = NULL,
                    xlim = NULL,
                    zlim = NULL,
                    add.legend = TRUE,
                    legend.only = FALSE,
                    vline = NULL,
                    col_vline = "black",
                    hline = NULL,
                    col_hline = "black",
                    cex.lab = 2,
                    cex.main = 2,
                    myCol = NULL,
                    includeMarginals = FALSE,
                    marginalJitterSD_x = 0.01,
                    marginalJitterSD_y = 0.01,
                    openBrowser = FALSE) {
  if (openBrowser) { browser() }
  s_ <- akima::interp(x = x, y = y, z = z,
                      xo = seq(min(x), max(x), length = N),
                      yo = seq(min(y), max(y), length = N),
                      duplicate = "mean")
  if (is.null(xlim)) { xlim = range(s_$x, finite = TRUE) }
  if (is.null(ylim)) { ylim = range(s_$y, finite = TRUE) }
  imageFxn <- if (add.legend) fields::image.plot else graphics::image
  if (!grepl(useLog, pattern = "z")) {
    imageFxn(s_, xlab = xlab, ylab = ylab, log = useLog, cex.lab = cex.lab, main = main,
             cex.main = cex.main, col = myCol, xlim = xlim, ylim = ylim,
             legend.width = legend.width, horizontal = horizontal, yaxt = yaxt,
             zlim = zlim, legend.only = legend.only)
  } else {
    useLog <- gsub(useLog, pattern = "z", replace = "")
    zTicks <- summary(c(s_$z))
    ep_ <- 0.001
    zTicks[zTicks < ep_] <- ep_
    zTicks <- exp(seq(log(min(zTicks)), log(max(zTicks)), length.out = 10))
    zTicks <- round(zTicks, abs(min(log(zTicks, base = 10))))
    s_$z[s_$z < ep_] <- ep_
    imageFxn(s_$x, s_$y, log(s_$z), yaxt = yaxt,
             axis.args = list(at = log(zTicks), labels = zTicks),
             main = main, cex.main = cex.main, xlab = xlab, ylab = ylab,
             log = useLog, cex.lab = cex.lab, xlim = xlim, ylim = ylim,
             horizontal = horizontal, col = myCol, legend.width = legend.width,
             zlim = zlim, legend.only = legend.only)
  }
  if (!is.null(vline)) { abline(v = vline, lwd = 10, col = col_vline) }
  if (!is.null(hline)) { abline(h = hline, lwd = 10, col = col_hline) }
  
  if (includeMarginals) {
    points(x + rnorm(length(y), sd = marginalJitterSD_x * sd(x)),
           rep(ylim[1] * 1.1, length(y)), pch = "|", col = "darkgray")
    points(rep(xlim[1] * 1.1, length(x)),
           y + rnorm(length(y), sd = sd(y) * marginalJitterSD_y), pch = "-", col = "darkgray")
  }
}

# UI Definition
ui <- fluidPage(
  titlePanel("Multiscale Heatmap & Surface Explorer"),
  sidebarLayout(
    sidebarPanel(
      selectInput("application", "Application",
                  choices = unique(sm$application),
                  selected = unique(sm$application)[1]),
      selectInput("model", "Model",
                  choices = unique(sm$optimizeImageRep),
                  selected = "clip"),
      # Removed "Perturb Center" input
      selectInput("metric", "Metric",
                  choices = c("AUTOC_rate_std_ratio_mean", "AUTOC_rate_mean", "AUTOC_rate_std_mean",
                              "AUTOC_rate_std_ratio_mean_pc", "AUTOC_rate_mean_pc", "AUTOC_rate_std_mean_pc",
                              "MeanVImportHalf1", "MeanVImportHalf2", "FracTopkHalf1", "RMSE"),
                  selected = "AUTOC_rate_std_ratio_mean"),
      radioButtons("plotType", "Plot Type",
                   choices = c("Heatmap", "Surface"),
                   selected = "Heatmap")
    ),
    mainPanel(
      uiOutput("plotOutput")
    )
  )
)

# Server Definition
server <- function(input, output) {
  # Reactive data processing
  filteredData <- reactive({
    # Removed filtering by 'perturbCenter'
    df <- sm %>%
      filter(application == input$application,
             optimizeImageRep == input$model) %>%
      mutate(MaxImageDimsRight = ifelse(is.na(MaxImageDimsRight),
                                        MaxImageDimsLeft,
                                        MaxImageDimsRight))
    if (nrow(df) == 0) return(NULL)
    df
  })
  
  # Render the plot output dynamically
  output$plotOutput <- renderUI({
    data <- filteredData()
    if (is.null(data)) {
      return(tags$p("No data available for the selected filters."))
    }
    
    if (input$plotType == "Heatmap") {
      plotOutput("heatmapPlot", height = "600px")
    } else {
      plotlyOutput("surfacePlot", height = "600px")
    }
  })
  
  # Heatmap Output
  output$heatmapPlot <- renderPlot({
    data <- filteredData()
    if (is.null(data)) return(NULL)
    
    # Group data for heatmap
    grouped_data <- data %>%
      group_by(MaxImageDimsLeft, MaxImageDimsRight) %>%
      summarise(
        mean_metric = mean(as.numeric(get(input$metric)), na.rm = TRUE),
        se_metric = sd(as.numeric(get(input$metric)), na.rm = TRUE) / sqrt(n()),
        n = n(),
        .groups = "drop"
      )
    
    # Check for sufficient data points for interpolation
    if (nrow(grouped_data) < 3) {
      plot.new()
      text(0.5, 0.5, "Insufficient data points for interpolation", cex = 1.5)
    } else {
      x <- grouped_data$MaxImageDimsLeft
      y <- grouped_data$MaxImageDimsRight
      z <- grouped_data$mean_metric
      
      # Slightly more appealing color palette
      customPalette <- colorRampPalette(c("blue", "white", "red"))(50)
      
      heatMap(x = x,
              y = y,
              z = z,
              N = 50,
              main = paste(input$application, "-", input$metric),
              # More descriptive axis labels
              xlab = "Maximum Image Dimensions (Left)",
              ylab = "Maximum Image Dimensions (Right)",
              useLog = "xy",
              myCol = customPalette,
              cex.lab = 1.4)
    }
  })
  
  # Surface Plot Output
  output$surfacePlot <- renderPlotly({
    data <- filteredData()
    if (is.null(data)) return(NULL)
    
    # Group data for surface plot
    grouped_data <- data %>%
      group_by(MaxImageDimsLeft, MaxImageDimsRight) %>%
      summarise(
        mean_metric = mean(as.numeric(get(input$metric)), na.rm = TRUE),
        se_metric = sd(as.numeric(get(input$metric)), na.rm = TRUE) / sqrt(n()),
        n = n(),
        .groups = "drop"
      )
    
    # Create grid for surface plot
    all_scales <- sort(unique(c(grouped_data$MaxImageDimsLeft, grouped_data$MaxImageDimsRight)))
    z_matrix <- matrix(NA, nrow = length(all_scales), ncol = length(all_scales))
    tooltip_matrix <- matrix("", nrow = length(all_scales), ncol = length(all_scales))
    
    for (i in 1:nrow(grouped_data)) {
      left_idx <- which(all_scales == grouped_data$MaxImageDimsLeft[i])
      right_idx <- which(all_scales == grouped_data$MaxImageDimsRight[i])
      z_matrix[left_idx, right_idx] <- grouped_data$mean_metric[i]
      tooltip_matrix[left_idx, right_idx] <- sprintf("Mean: %.2f<br>SE: %.2f<br>n: %d",
                                                     grouped_data$mean_metric[i],
                                                     grouped_data$se_metric[i],
                                                     grouped_data$n[i])
    }
    
    # Render interactive 3D surface plot
    plot_ly(
      x = all_scales,
      y = all_scales,
      z = z_matrix,
      type = "surface",
      text = tooltip_matrix,
      hoverinfo = "text"
    ) %>%
      layout(
        title = paste("Surface Plot for", input$metric, "in", input$application),
        scene = list(
          xaxis = list(title = "Maximum Image Dimensions (Right)"),
          yaxis = list(title = "Maximum Image Dimensions (Left)"),
          zaxis = list(title = input$metric)
        )
      )
  })
}

# Run the Shiny App
shinyApp(ui = ui, server = server)