source("R/globals.R")
source("R/utils_validate.R")
source("R/utils_preprocess.R")
source("R/utils_sampling.R")
source("R/utils_annotation.R")
source("R/utils_qpt.R")

ui <- fluidPage(
    titlePanel("QPT Shiny App"),
    
    sidebarLayout(
        sidebarPanel(
            width = 2,
            style = "font-size: 12px;",
            
            fileInput("datafile", "Upload CSV", accept = ".csv"),
            hr(),
            uiOutput("setup_ui"),
            hr(),
            numericInput("n_tasks", "Number of sampled tasks", value = 20, min = 1),
            numericInput("seed", "Random seed", value = 123, min = 1),
            actionButton("sample_btn", "Sample Tasks", width = "100%"),
            hr(),
            actionButton("opt_btn", "Run Optimization", width = "100%"),
            hr(),
            downloadButton("export_results", "Export Annotations + Results", width = "100%")
        ),
        
        mainPanel(
            width = 10,
            
            tabsetPanel(
                tabPanel(
                    "Upload",
                    h4("Validation"),
                    verbatimTextOutput("validation_text"),
                    h4("Preview"),
                    DTOutput("data_preview")
                ),
                
                tabPanel(
                    "Annotate",
                    br(),
                    
                    fluidRow(
                        column(
                            width = 12,
                            h4(textOutput("task_counter"))
                        )
                    ),
                    
                    div(
                        style = "height: 760px; display: flex; flex-direction: column;",
                        
                        div(
                            style = "flex: 1; min-height: 0;",
                            fluidRow(
                                style = "height: 100%;",
                                
                                column(
                                    width = 2,
                                    style = "height: 100%; overflow-y: auto;",
                                    h4("From"),
                                    uiOutput("edge_from_ui")
                                ),
                                
                                column(
                                    width = 2,
                                    style = "height: 100%; overflow-y: auto;",
                                    h4("To"),
                                    uiOutput("edge_to_ui")
                                ),
                                
                                column(
                                    width = 2,
                                    style = "height: 100%; padding-top: 40px;",
                                    actionButton("add_edge_btn", "Add Edge", width = "100%"),
                                    br(), br(),
                                    actionButton("clear_edges_btn", "Clear Current Task", width = "100%")
                                ),
                                
                                column(
                                    width = 6,
                                    style = "height: 100%;",
                                    h4("Annotation Network"),
                                    div(
                                        style = "height: calc(100% - 40px);",
                                        visNetworkOutput("annotation_network", height = "100%")
                                    )
                                )
                            )
                        ),
                        
                        div(
                            style = "flex: 1; min-height: 0; margin-top: 15px;",
                            
                            fluidRow(
                                column(
                                    width = 6,
                                    actionButton("prev_task_btn", "<-", width = "60px")
                                ),
                                column(
                                    width = 6,
                                    div(
                                        style = "text-align: right;",
                                        actionButton("next_task_btn", "->", width = "60px")
                                    )
                                )
                            ),
                            
                            br(),
                            
                            h4("Context"),
                            div(
                                style = "height: calc(100% - 50px);",
                                DTOutput("current_task_context")
                            )
                        )
                    )
                ),
                
                tabPanel(
                    "Optimize",
                    h4("Recommended Window"),
                    verbatimTextOutput("best_window_text"),
                    h4("Best Result"),
                    DTOutput("best_window_table"),
                    h4("Refined Search Results"),
                    DTOutput("loss_table"),
                    h4("Loss Visualization"),
                    plotOutput("loss_plot", height = "550px")
                )
            )
        )
    )
)

server <- function(input, output, session) {
    current_task_index <- reactiveVal(1)
    
    annotations_store <- reactiveValues(
        nodes = list(),
        edges = list()
    )
    
    raw_data <- reactive({
        req(input$datafile)
        readr::read_csv(input$datafile$datapath, show_col_types = FALSE)
    })
    
    validated <- reactive({
        req(raw_data())
        validate_uploaded_data(raw_data())
    })
    
    output$validation_text <- renderPrint({
        v <- validated()
        if (v$ok) {
            cat("Validation passed.\n")
            cat(paste(v$messages, collapse = "\n"))
        } else {
            cat("Validation failed.\n")
            cat(paste(v$messages, collapse = "\n"))
        }
    })
    
    output$data_preview <- renderDT({
        req(validated()$ok)
        datatable(
            head(validated()$data, 20),
            options = list(scrollX = TRUE),
            rownames = FALSE
        )
    })
    
    output$setup_ui <- renderUI({
        req(validated()$ok)
        cols <- names(validated()$data)
        
        tagList(
            selectInput(
                "modality_col",
                "Select modality column",
                choices = cols,
                selected = if ("modality" %in% cols) "modality" else cols[1]
            ),
            selectInput(
                "optimization_mode",
                "Optimization mode",
                choices = c(
                    "General window" = "general",
                    "Modality-specific windows" = "multimodal"
                ),
                selected = "multimodal"
            ),
            conditionalPanel(
                condition = "input.optimization_mode == 'multimodal'",
                uiOutput("modality_values_ui")
            ),
            selectizeInput(
                "horizon_cols",
                "Select horizon column(s)",
                choices = cols,
                multiple = TRUE,
                selected = intersect(c("PlayerID", "StuckID", "Learned", "Level"), cols)
            ),
            selectizeInput(
                "code_cols",
                "Select code column(s)",
                choices = cols,
                multiple = TRUE,
                selected = intersect(DEFAULT_CODE_COLS, cols)
            )
        )
    })
    
    output$modality_values_ui <- renderUI({
        req(validated()$ok, input$modality_col)
        vals <- validated()$data[[input$modality_col]]
        vals <- unique(as.character(vals))
        vals <- vals[!is.na(vals)]
        vals <- sort(vals)
        
        selectizeInput(
            "modality_values",
            "Select modality values",
            choices = vals,
            selected = vals,
            multiple = TRUE
        )
    })
    
    events_clean <- reactive({
        req(validated()$ok, input$horizon_cols, input$code_cols)
        validate(
            need(length(input$horizon_cols) >= 1, "Please select at least one horizon column."),
            need(length(input$code_cols) >= 1, "Please select at least one code column.")
        )
        
        preprocess_events(
            df = validated()$data,
            horizon_cols = input$horizon_cols,
            code_cols = input$code_cols
        )
    })
    
    sampled_tasks <- eventReactive(input$sample_btn, {
        req(events_clean())
        sample_tasks(
            df = events_clean(),
            n_tasks = input$n_tasks,
            seed = input$seed
        )
    })
    
    observeEvent(sampled_tasks(), {
        current_task_index(1)
        annotations_store$nodes <- list()
        annotations_store$edges <- list()
    })
    
    current_task <- reactive({
        req(sampled_tasks())
        idx <- current_task_index()
        sampled_tasks()[idx, , drop = FALSE]
    })
    
    current_task_id <- reactive({
        req(current_task())
        current_task()$task_id
    })
    
    ensure_task_annotation <- function(task_id, selected_nodes) {
        if (is.null(annotations_store$nodes[[task_id]])) {
            annotations_store$nodes[[task_id]] <- selected_nodes
        }
        if (is.null(annotations_store$edges[[task_id]])) {
            annotations_store$edges[[task_id]] <- tibble::tibble(
                from = character(0),
                to = character(0)
            )
        }
    }
    
    save_current_task_annotation <- function() {
        req(current_task_id(), input$code_cols)
        task_id <- current_task_id()
        annotations_store$nodes[[task_id]] <- input$code_cols
        
        if (is.null(annotations_store$edges[[task_id]])) {
            annotations_store$edges[[task_id]] <- tibble::tibble(
                from = character(0),
                to = character(0)
            )
        }
    }
    
    observe({
        req(current_task_id(), input$code_cols)
        task_id <- current_task_id()
        ensure_task_annotation(task_id, input$code_cols)
    })
    
    output$edge_from_ui <- renderUI({
        req(input$code_cols)
        radioButtons(
            "edge_from",
            label = NULL,
            choices = input$code_cols,
            selected = input$code_cols[1]
        )
    })
    
    output$edge_to_ui <- renderUI({
        req(input$code_cols)
        radioButtons(
            "edge_to",
            label = NULL,
            choices = input$code_cols,
            selected = input$code_cols[1]
        )
    })
    
    output$task_counter <- renderText({
        req(sampled_tasks())
        paste("Task", current_task_index(), "of", nrow(sampled_tasks()))
    })
    
    output$current_task_context <- renderDT({
        req(events_clean(), current_task())
        
        full_context <- get_task_context(events_clean(), current_task())
        
        dt <- datatable(
            full_context,
            options = list(
                paging = FALSE,
                searching = FALSE,
                info = FALSE,
                scrollX = TRUE,
                scrollY = "450px",
                ordering = FALSE,
                dom = "t"
            ),
            rownames = FALSE,
            callback = JS(
                "table.on('draw.dt', function() {",
                "  var body = $(table.table().container()).find('.dataTables_scrollBody');",
                "  body.scrollTop(body[0].scrollHeight);",
                "});"
            )
        )
        
        dt <- formatStyle(
            dt,
            "is_focal",
            target = "row",
            color = styleEqual(c(TRUE, FALSE), c("red", "black")),
            fontWeight = styleEqual(c(TRUE, FALSE), c("bold", "normal"))
        )
        
        dt
    })
    output$export_results <- downloadHandler(
        filename = function() {
            paste0("qpt_export_", Sys.Date(), ".xlsx")
        },
        content = function(file) {
            req(sampled_tasks(), events_clean())
            
            save_current_task_annotation()
            
            tasks_df <- sampled_tasks()
            
            annotation_edges <- purrr::map_dfr(tasks_df$task_id, function(task_id) {
                edges_i <- annotations_store$edges[[task_id]]
                
                if (is.null(edges_i) || nrow(edges_i) == 0) {
                    return(tibble::tibble(task_id = task_id, from = NA, to = NA))
                }
                
                edges_i %>%
                    dplyr::mutate(task_id = task_id) %>%
                    dplyr::select(task_id, from, to)
            })
            
            contexts_df <- purrr::map_dfr(seq_len(nrow(tasks_df)), function(i) {
                task_row <- tasks_df[i, , drop = FALSE]
                ctx <- get_task_context(events_clean(), task_row)
                
                ctx %>%
                    dplyr::mutate(task_id = task_row$task_id[1]) %>%
                    dplyr::select(task_id, dplyr::everything())
            })
            
            export_list <- list(
                tasks = tasks_df,
                annotation_edges = annotation_edges,
                contexts = contexts_df
            )
            
            if (!is.null(optimization_results())) {
                export_list$best_windows <- optimization_results()$best_windows
            }
            
            writexl::write_xlsx(export_list, path = file)
        }
    )
    observeEvent(input$add_edge_btn, {
        req(current_task_id(), input$edge_from, input$edge_to, input$code_cols)
        
        task_id <- current_task_id()
        ensure_task_annotation(task_id, input$code_cols)
        annotations_store$nodes[[task_id]] <- input$code_cols
        
        new_edge <- tibble::tibble(
            from = input$edge_from,
            to = input$edge_to
        )
        
        annotations_store$edges[[task_id]] <- dplyr::bind_rows(
            annotations_store$edges[[task_id]],
            new_edge
        ) %>%
            dplyr::distinct(from, to)
    })
    
    observeEvent(input$clear_edges_btn, {
        req(current_task_id(), input$code_cols)
        
        task_id <- current_task_id()
        annotations_store$nodes[[task_id]] <- input$code_cols
        annotations_store$edges[[task_id]] <- tibble::tibble(
            from = character(0),
            to = character(0)
        )
    })
    
    output$annotation_network <- renderVisNetwork({
        req(current_task_id(), input$code_cols)
        
        task_id <- current_task_id()
        ensure_task_annotation(task_id, input$code_cols)
        
        nodes_df <- make_vis_nodes(input$code_cols)
        edges_df <- make_vis_edges(annotations_store$edges[[task_id]])
        
        visNetwork(nodes_df, edges_df, width = "100%", height = "100%") %>%
            visEdges(arrows = "to") %>%
            visPhysics(stabilization = FALSE) %>%
            visOptions(highlightNearest = TRUE)
    })
    
    observeEvent(input$next_task_btn, {
        req(sampled_tasks())
        save_current_task_annotation()
        
        idx <- current_task_index()
        if (idx < nrow(sampled_tasks())) {
            current_task_index(idx + 1)
        }
    })
    
    observeEvent(input$prev_task_btn, {
        req(sampled_tasks())
        save_current_task_annotation()
        
        idx <- current_task_index()
        if (idx > 1) {
            current_task_index(idx - 1)
        }
    })
    
    optimization_results <- eventReactive(input$opt_btn, {
        req(events_clean(), sampled_tasks(), input$optimization_mode)
        
        save_current_task_annotation()
        
        annotations <- collect_annotations_from_store(
            tasks = sampled_tasks(),
            store = annotations_store,
            fallback_nodes = input$code_cols
        )
        
        if (input$optimization_mode == "general") {
            optimize_general_window_coarse_to_fine(
                events = events_clean(),
                tasks = sampled_tasks(),
                annotations = annotations,
                coarse_points = 10,
                refine_radius = 3,
                do_hill_climb = TRUE,
                hill_max_iter = 75
            )
        } else {
            req(input$modality_col, input$modality_values)
            validate(
                need(length(input$modality_values) >= 1, "Please select at least one modality value.")
            )
            
            optimize_window_combinations_coarse_to_fine(
                events = events_clean(),
                tasks = sampled_tasks(),
                annotations = annotations,
                modality_col = input$modality_col,
                modality_values = input$modality_values,
                coarse_points = 10,
                refine_radius = 3,
                do_hill_climb = TRUE,
                hill_max_iter = 75
            )
        }
    })
    
    output$best_window_text <- renderPrint({
        req(optimization_results(), input$optimization_mode)
        
        best_tbl <- optimization_results()$best_windows
        
        if (nrow(best_tbl) == 0) {
            cat("No window recommendation available.\n")
            return()
        }
        
        if (input$optimization_mode == "general") {
            cat("Recommended general window:\n\n")
            cat("window: ", best_tbl$window[1], "\n", sep = "")
        } else {
            cat("Recommended window combination:\n\n")
            
            window_cols <- intersect(input$modality_values, names(best_tbl))
            
            for (col in window_cols) {
                cat(col, ": ", best_tbl[[col]][1], "\n", sep = "")
            }
        }
        
        cat("\nAverage misconnection loss: ", round(best_tbl$mean_loss[1], 4), "\n", sep = "")
        cat("Maximum searched window: ", optimization_results()$k_max, "\n", sep = "")
        cat("Search method: ", optimization_results()$search_method, "\n", sep = "")
    })
    
    output$best_window_table <- renderDT({
        req(optimization_results())
        
        datatable(
            optimization_results()$best_windows,
            options = list(
                paging = FALSE,
                searching = FALSE,
                info = FALSE,
                scrollX = TRUE
            ),
            rownames = FALSE
        )
    })
    
    output$loss_table <- renderDT({
        req(optimization_results())
        
        coarse_tbl <- optimization_results()$coarse_results %>%
            dplyr::mutate(search_stage = "coarse")
        
        refine_tbl <- optimization_results()$refine_results %>%
            dplyr::mutate(search_stage = "refine")
        
        hill_tbl <- optimization_results()$hill_results
        
        if (!is.null(hill_tbl)) {
            hill_tbl <- hill_tbl %>%
                dplyr::mutate(search_stage = "hill")
            
            out <- dplyr::bind_rows(coarse_tbl, refine_tbl, hill_tbl)
        } else {
            out <- dplyr::bind_rows(coarse_tbl, refine_tbl)
        }
        
        datatable(
            out,
            options = list(pageLength = 25, scrollX = TRUE),
            rownames = FALSE
        )
    })
    
    output$loss_plot <- renderPlot({
        req(optimization_results(), input$optimization_mode)
        
        if (input$optimization_mode == "general") {
            res <- optimization_results()$refine_results
            
            ggplot(
                res,
                aes(
                    x = window,
                    y = mean_loss
                )
            ) +
                geom_line() +
                geom_point() +
                labs(
                    x = "Window",
                    y = "Mean loss",
                    title = "Refined mean loss by general window"
                ) +
                theme_minimal()
            
        } else {
            res <- optimization_results()$refine_results
            mod_vals <- input$modality_values
            
            if (length(mod_vals) == 2) {
                x_col <- mod_vals[1]
                y_col <- mod_vals[2]
                
                validate(
                    need(x_col %in% names(res), paste("Missing column in results:", x_col)),
                    need(y_col %in% names(res), paste("Missing column in results:", y_col))
                )
                
                ggplot(
                    res,
                    aes(
                        x = .data[[x_col]],
                        y = .data[[y_col]],
                        fill = mean_loss
                    )
                ) +
                    geom_tile() +
                    labs(
                        x = paste0(x_col, " window"),
                        y = paste0(y_col, " window"),
                        fill = "Mean loss",
                        title = "Refined window-combination heatmap"
                    ) +
                    theme_minimal()
                
            } else if (length(mod_vals) == 1) {
                x_col <- mod_vals[1]
                
                validate(
                    need(x_col %in% names(res), paste("Missing column in results:", x_col))
                )
                
                ggplot(
                    res,
                    aes(
                        x = .data[[x_col]],
                        y = mean_loss
                    )
                ) +
                    geom_line() +
                    geom_point() +
                    labs(
                        x = paste0(x_col, " window"),
                        y = "Mean loss",
                        title = "Refined mean loss by window size"
                    ) +
                    theme_minimal()
                
            } else {
                ggplot(
                    res %>% dplyr::mutate(combo_id = dplyr::row_number()),
                    aes(
                        x = combo_id,
                        y = mean_loss
                    )
                ) +
                    geom_point() +
                    labs(
                        x = "Refined combination index",
                        y = "Mean loss",
                        title = "Refined mean loss across combinations"
                    ) +
                    theme_minimal()
            }
        }
    })
}

shinyApp(ui, server)
