library(shiny)
library(ggplot2)
library(DT)

# Default data with line breaks
default_positions <- "113\n145\n178\n234\n267\n289\n345\n378\n423\n456"
default_wrapping <- "147\n147\n147\n147\n147\n147\n147\n147\n147\n147"

default_annotations <- data.frame(
  type = c("CENP-B Box Type 4", "CENP-B Box Type 4", "CENP-B Box Type 3", "Alpha Satellite", "Alpha Satellite"),
  start = c(236, 577, 915, 101, 956),
  end = c(246, 587, 925, 101, 956),
  color = c("blue", "blue", "green", "black", "black"),
  annotation_type = c("box", "box", "box", "line", "line")
)

ui <- fluidPage(
  titlePanel("Nucleosome Heatmap Generator"),
  
  sidebarLayout(
    sidebarPanel(
      textAreaInput("positions", "Nucleosome Positions (one per line)", value = default_positions, height = "300px"),
      textAreaInput("wrapping", "Wrapping Lengths (one per line)", value = default_wrapping, height = "300px"),
      
      hr(),
      h4("Add Annotation"),
      selectInput("annotation_type", "Type", choices = c("box", "line")),
      textInput("annotation_label", "Label"),
      numericInput("annotation_start", "Start Position", value = 0),
      numericInput("annotation_end", "End Position", value = 0),
      selectInput("annotation_color", "Color", 
                  choices = c("blue", "green", "red", "black", "purple", "orange")),
      actionButton("add_annotation", "Add Annotation"),
      
      hr(),
      numericInput("max_intensity", "Max Intensity", value = 40),
      downloadButton("download_plot", "Download Plot")
    ),
    
    mainPanel(
      plotOutput("heatmap"),
      hr(),
      h4("Current Annotations"),
      DTOutput("annotation_table"),
      actionButton("remove_selected", "Remove Selected Annotations")
    )
  )
)

server <- function(input, output, session) {
  rv <- reactiveValues(
    annotations = default_annotations
  )
  
  process_data <- function(positions_string, wrapping_string) {
    positions <- as.numeric(unlist(strsplit(positions_string, "\n")))
    wrapping <- as.numeric(unlist(strsplit(wrapping_string, "\n")))
    
    heatmap_values <- rep(0, 1000)
    
    for (i in seq_along(positions)) {
      if (!is.na(positions[i]) && !is.na(wrapping[i]) && positions[i] != "" && wrapping[i] != "") {
        wrap_start <- positions[i] - floor(wrapping[i]/2)
        wrap_end <- positions[i] + floor(wrapping[i]/2)
        
        for (pos in wrap_start:wrap_end) {
          if (pos >= 0 && pos <= 999) {
            heatmap_values[pos + 1] <- heatmap_values[pos + 1] + 1
          }
        }
      }
    }
    
    return(heatmap_values)
  }
  
  create_heatmap <- function(data, title) {
    df <- data.frame(Position = 0:999, Intensity = data)
    
    p <- ggplot(df, aes(x = Position, y = 1, fill = Intensity)) +
      geom_tile() +
      scale_fill_gradientn(
        colors = c("white", "yellow", "orange", "red", "darkred"),
        limits = c(0, input$max_intensity)
      ) +
      labs(title = title, x = "Position (bp)", y = "") +
      theme_minimal() +
      theme(axis.text.y = element_blank(), axis.ticks.y = element_blank())
    
    for (i in 1:nrow(rv$annotations)) {
      ann <- rv$annotations[i,]
      if (ann$annotation_type == "box") {
        p <- p + annotate("rect", 
                          xmin = ann$start, xmax = ann$end,
                          ymin = 0.5, ymax = 1.5,
                          fill = ann$color, alpha = 0.5)
        p <- p + annotate("text",
                          x = (ann$start + ann$end)/2,
                          y = 1.6,
                          label = ann$type,
                          size = 3)
      } else if (ann$annotation_type == "line") {
        p <- p + annotate("segment",
                          x = ann$start, xend = ann$start,
                          y = 0.5, yend = 1.5,
                          color = ann$color,
                          size = 1)
        p <- p + annotate("text",
                          x = ann$start,
                          y = 0.4,
                          label = ann$type,
                          size = 3)
      }
    }
    
    return(p)
  }
  
  output$heatmap <- renderPlot({
    data <- process_data(input$positions, input$wrapping)
    create_heatmap(data, "Nucleosome Heatmap")
  })
  
  observeEvent(input$add_annotation, {
    new_annotation <- data.frame(
      type = input$annotation_label,
      start = input$annotation_start,
      end = input$annotation_end,
      color = input$annotation_color,
      annotation_type = input$annotation_type
    )
    rv$annotations <- rbind(rv$annotations, new_annotation)
  })
  
  output$annotation_table <- renderDT({
    rv$annotations
  }, selection = 'multiple')
  
  observeEvent(input$remove_selected, {
    if (!is.null(input$annotation_table_rows_selected)) {
      rv$annotations <- rv$annotations[-input$annotation_table_rows_selected,]
    }
  })
  
  output$download_plot <- downloadHandler(
    filename = function() {
      paste0("nucleosome_heatmap_", Sys.Date(), ".png")
    },
    content = function(file) {
      data <- process_data(input$positions, input$wrapping)
      ggsave(file, plot = create_heatmap(data, "Nucleosome Heatmap"),
             width = 10, height = 2)
    }
  )
}

shinyApp(ui = ui, server = server)