
library(shiny)
library(ggplot2)
library(caret)
library(DT)
library(plotly)

# Load model
outputDir <- getwd()  # Get current working directory
model_path <- file.path(outputDir, 'model.rds')

cat('Loading model...
')
cat('Model path:', model_path, '
')

# Check if model file exists
if (!file.exists(model_path)) {
  stop('Model file does not exist. Please ensure model.rds is in the current directory')
}

model_data <- readRDS(model_path)
feature_info <- model_data$feature_info

cat('Model loaded successfully! Number of features:', length(feature_info), '
')

# Safe value extraction function
safe_get_value <- function(info, field) {
  tryCatch({
    val <- as.numeric(info[[field]])
    if (is.na(val) || is.null(val)) {
      if (field == 'min') return(0)
      if (field == 'max') return(100)
      if (field == 'mean') return(50)
      return(25)
    }
    return(val)
  }, error = function(e) {
    if (field == 'min') return(0)
    if (field == 'max') return(100)
    if (field == 'mean') return(50)
    return(25)
  })
}

ui <- fluidPage(
  titlePanel('CPB Postoperative Infection risk Prediction System'),
  
  tags$head(
    tags$style(HTML('
      body {
        background-color: #f5f5f5;
      }
      .well {
        padding: 20px;
        border-radius: 10px;
        box-shadow: 0 2px 5px rgba(0,0,0,0.1);
        background-color: white;
      }
      h2 { margin-top: 0; }
      .sidebar-panel {
        background-color: white;
        padding: 20px;
        border-radius: 10px;
        height: calc(90vh - 10px);  /* 设置固定高度，与右侧保持一致 */
        display: flex;
        flex-direction: column;
      }
      .input-scroll-area {
        flex: 1;  /* 占据剩余所有空间 */
        overflow-y: auto;
        padding-right: 5px;
      }
      .risk-legend {
        margin-top: 20px;
        padding: 15px;
        background-color: #f8f9fa;
        border-radius: 8px;
        border-left: 4px solid #3498db;
      }
      .risk-items-container {
        display: flex;
        flex-direction: row;
        justify-content: space-around;
        align-items: center;
        flex-wrap: wrap;
        gap: 15px;
        margin-top: 10px;
      }
      .risk-item-horizontal {
        display: flex;
        align-items: center;
        margin: 5px 0;
        font-size: 15px;
        flex: 1;
        min-width: 180px;
      }
      .risk-color {
        width: 20px;
        height: 20px;
        border-radius: 4px;
        margin-right: 8px;
        flex-shrink: 0;
      }
      .risk-low { background-color: #28a745; }
      .risk-medium { background-color: #ffaa00; }
      .risk-high { background-color: #dc3545; }
      .risk-text {
        white-space: nowrap;
      }
    '))
  ),
  
  sidebarLayout(
    sidebarPanel(
      width = 4,
      class = 'sidebar-panel',  # 应用新的样式类
      
      h3('Patient Information Input', style = 'color: #2c3e50; margin-top: 0; margin-bottom: 15px;'),
      
      # Scrollable input area
      div(class = 'input-scroll-area',
          uiOutput('patient_inputs')
      )
    ),
    
    mainPanel(
      width = 8,
      fluidRow(
        column(6,
               div(class = 'well',
                   h4('📊 Risk Probability', style = 'color: #3498db; margin-top: 0;'),
                   h1(textOutput('risk_prob'),
                      style = 'color: #e74c3c; text-align: center; margin: 20px 0; font-size: 48px;')
               )
        ),
        column(6,
               div(class = 'well',
                   h4('⚠️ Risk Level', style = 'color: #3498db; margin-top: 0;'),
                   h1(textOutput('risk_level'),
                      style = 'color: #27ae60; text-align: center; margin: 20px 0; font-size: 48px;')
               )
        )
      ),
      
      div(class = 'well',
          h4('📈 Risk Dashboard', style = 'color: #3498db; margin-top: 0;'),
          plotlyOutput('prediction_plot', height = '350px'),
          
          # Risk classification legend - Horizontal layout
          div(class = 'risk-legend',
              h5('Risk Classification Criteria:', style = 'color: #2c3e50; margin-top: 0; margin-bottom: 10px; font-weight: bold;'),
              div(class = 'risk-items-container',
                  div(class = 'risk-item-horizontal',
                      div(class = 'risk-color risk-low'),
                      span(class = 'risk-text', 'Low Risk: ≤33.33%')
                  ),
                  div(class = 'risk-item-horizontal',
                      div(class = 'risk-color risk-medium'),
                      span(class = 'risk-text', 'Medium Risk: 33.33% - 66.66%')
                  ),
                  div(class = 'risk-item-horizontal',
                      div(class = 'risk-color risk-high'),
                      span(class = 'risk-text', 'High Risk: >66.66%')
                  )
              )
          )
      )
      

    )
  )
)

server <- function(input, output, session) {
  
  # Generate input controls
  output$patient_inputs <- renderUI({
    req(feature_info)
    
    input_list <- list()
    
    for (var in names(feature_info)) {
      info <- feature_info[[var]]
      
      if (info$type == 'numeric') {
        # Safely get numeric values
        min_val <- safe_get_value(info, 'min')
        max_val <- safe_get_value(info, 'max')
        mean_val <- safe_get_value(info, 'mean')
        
        # Ensure valid values
        if (min_val >= max_val || is.na(min_val) || is.na(max_val)) {
          max_val <- min_val + 10
        }
        if (is.na(mean_val)) {
          mean_val <- (min_val + max_val) / 2
        }
        mean_val <- max(min_val, min(max_val, mean_val))
        
        input_list[[var]] <- sliderInput(
          inputId = var,
          label = var,
          min = min_val,
          max = max_val,
          value = mean_val,
          step = max(0.01, (max_val - min_val) / 100)
        )
      } else if (info$type == 'categorical') {
        if (length(info$levels) > 0) {
          input_list[[var]] <- selectInput(
            inputId = var,
            label = var,
            choices = info$levels,
            selected = info$levels[1]
          )
        }
      }
    }
    
    if (length(input_list) == 0) {
      return(h4('No available input variables', style = 'color: red;'))
    }
    
    do.call(tagList, input_list)
  })
  
  # Reactive prediction results - triggered by any input change
  predictions <- reactive({
    req(feature_info)
    
    # Collect input values
    input_values <- list()
    for (var in names(feature_info)) {
      if (!is.null(input[[var]])) {
        input_values[[var]] <- input[[var]]
      }
    }
    
    if (length(input_values) == 0) {
      return(data.frame())
    }
    
    input_df <- as.data.frame(input_values, stringsAsFactors = FALSE)
    
    # Convert data types
    for (var in names(feature_info)) {
      if (feature_info[[var]]$type == 'categorical') {
        input_df[[var]] <- as.factor(input_df[[var]])
      } else {
        input_df[[var]] <- as.numeric(input_df[[var]])
      }
    }
    
    # Perform predictions
    results <- data.frame()
    for (model_name in names(model_data$models)) {
      model <- model_data$models[[model_name]]
      prob <- tryCatch({
        pred_result <- predict(model, newdata = input_df, type = 'prob')
        # Handle different prediction result formats
        if (is.data.frame(pred_result)) {
          if ('Yes' %in% names(pred_result)) {
            pred_result$Yes
          } else if ('1' %in% names(pred_result)) {
            pred_result[['1']]
          } else {
            pred_result[[1]]
          }
        } else {
          0.5
        }
      }, error = function(e) {
        cat('Model', model_name, 'prediction error:', e$message, '
')
        0.5
      })
      
      # Ensure prob is numeric
      prob_value <- ifelse(is.numeric(prob) && length(prob) > 0, 
                           as.numeric(prob[1]), 0.5)
      
      # Determine risk category with three levels
      risk_category <- if(prob_value > 0.7) {
        'High Risk'
      } else if(prob_value > 0.5) {
        'Medium Risk'
      } else {
        'Low Risk'
      }
      
      results <- rbind(results, data.frame(
        Model = model_name,
        Probability = round(prob_value, 4),
        Risk = risk_category,
        stringsAsFactors = FALSE
      ))
    }
    
    results
  })
  
  # Calculate average risk probability (for dashboard)
  avg_risk_prob <- reactive({
    pred <- predictions()
    if (nrow(pred) == 0) return(0)
    mean(pred$Probability, na.rm = TRUE) * 100
  })
  
  # Display prediction results
  output$risk_prob <- renderText({
    req(avg_risk_prob())
    paste0(round(avg_risk_prob(), 1), '%')
  })
  
  output$risk_level <- renderText({
    risk_val <- avg_risk_prob()
    if (risk_val == 0) return('Not predicted')
    
    # Three-level risk classification
    if (risk_val > 70) {
      'High Risk'
    } else if (risk_val > 50) {
      'Medium Risk'
    } else {
      'Low Risk'
    }
  })
  
  # Plotly tech-style dashboard
  output$prediction_plot <- plotly::renderPlotly({
    risk_prob <- avg_risk_prob()
    
    # Determine color based on risk probability (three levels)
    color <- if(risk_prob > 66.66) {
      '#ff0080'  # High risk - Neon pink
    } else if(risk_prob > 33.33) {
      '#ffaa00'  # Medium risk - Neon orange
    } else {
      '#00ff88'  # Low risk - Neon green
    }
    
    # Create 3D tech-style gauge with solid vibrant colors
    fig <- plotly::plot_ly(
      type = 'indicator',
      mode = 'gauge+number',
      value = risk_prob,
      number = list(
        suffix = '%',
        font = list(
          size = 48,
          family = 'Arial, sans-serif',
          color = color,
          weight = 'bold'
        )
      ),
      title = list(
        text = 'Average Risk Probability',
        font = list(
          size = 16,
          family = 'Arial, sans-serif',
          color = '#333333'
        )
      ),
      gauge = list(
        axis = list(
          range = list(0, 100),
          tickwidth = 2,
          tickcolor = '#666666',
          tickfont = list(
            family = 'Arial, sans-serif',
            color = '#666666',
            size = 12
          ),
          tickvals = c(0, 20, 40, 60, 80, 100),
          ticktext = c('0%', '20%', '40%', '60%', '80%', '100%'),
          nticks = 6
        ),
        bar = list(
          color = color,
          thickness = 0.4,
          line = list(color = '#ffffff', width = 2)
        ),
        bgcolor = '#e0e0e0',
        borderwidth = 0,
        bordercolor = 'transparent',
        shape = 'angular',
        steps = list(
          list(range = c(0, 33.33),
               color = '#ffff00',
               line = list(color = '#ffff00', width = 0)),
          list(range = c(33.33, 66.66),
               color = '#4169e1',
               line = list(color = '#4169e1', width = 0)),
          list(range = c(66.66, 100),
               color = '#9370db',
               line = list(color = '#9370db', width = 0))
        ),
        threshold = list(
          line = list(color = color, width = 6),
          thickness = 0.8,
          value = risk_prob
        )
      )
    )
    
    # Add subtitle and risk level indicator (three levels)
    risk_text <- if(risk_prob > 66.66) {
      '<span style="color:#ff0080; font-weight:bold; text-shadow: 2px 2px 4px rgba(0,0,0,0.3);">⚡ HIGH RISK</span>'
    } else if(risk_prob > 33.33) {
      '<span style="color:#ffaa00; font-weight:bold; text-shadow: 2px 2px 4px rgba(0,0,0,0.3);">⚠️ MEDIUM RISK</span>'
    } else {
      '<span style="color:#00ff88; font-weight:bold; text-shadow: 2px 2px 4px rgba(0,0,0,0.3);">✓ LOW RISK</span>'
    }
    
    fig <- fig %>%
      plotly::layout(
        paper_bgcolor = 'transparent',
        plot_bgcolor = 'transparent',
        font = list(family = 'Arial, sans-serif', color = '#333333'),
        margin = list(t = 80, b = 40, l = 40, r = 40),
        annotations = list(
          list(
            x = 0.5,
            y = 0.25,
            xref = 'paper',
            yref = 'paper',
            text = risk_text,
            showarrow = FALSE,
            font = list(
              family = 'Arial, sans-serif',
              size = 18,
              color = color,
              weight = 'bold'
            )
          )
        )
      )
    
    fig
  })
}

shinyApp(ui, server)

