# Function to check if a package is installed, and install it if not
install_if_missing <- function(p) {
  if (!require(p, character.only = TRUE)) {
    install.packages(p, dependencies = TRUE)
    library(p, character.only = TRUE)
  }
}

# List of required packages
packages <- c("shiny", "shinycssloaders", "shinyjs", "prophet", "readr", "ggplot2", "dplyr", "plotly", "future", "bsplus", "shinythemes", "DT", "forecast")
lapply(packages, install_if_missing)

# Load the necessary libraries
library(shiny)
library(shinycssloaders)
library(shinyjs)
library(prophet)
library(readr)
library(ggplot2)
library(dplyr)
library(plotly)
library(future)
library(bsplus)
library(shinythemes)
library(DT)
library(forecast)

plan(multisession)  # Enable parallel processing

# Load default dataset
default_data <- read_csv("filtered_mortality_data_for_r.csv") %>%
  rename(ds = date, y = mortality_rate) %>%
  mutate(ds = as.Date(ds)) %>%
  distinct(ds, .keep_all = TRUE)

# Define UI
ui <- fluidPage(
  theme = shinytheme("cyborg"),
  
  # Add custom CSS for Comfortaa font and white text
  tags$head(
    tags$link(
      href = "https://fonts.googleapis.com/css2?family=Comfortaa:wght@300;400;700&display=swap", 
      rel = "stylesheet"
    ),
    tags$style(HTML("
      body {
        font-family: 'Comfortaa', sans-serif;
      
      }
      .shiny-output-error {
        color: white;
      }
      .shiny-input-container {
        color: white;
      }
      h2, h3, h4, h5, h6, p, label {
        color: white;
        font-family: 'Comfortaa', sans-serif;
      }
      .tab-pane {
        background-color: black;
        color: white;
      }
      .nav-tabs > li > a {
        color: white;
      }
      .nav-tabs > li > a:hover {
        color: white;
      }
      .nav-tabs > .active > a, .nav-tabs > .active > a:focus, .nav-tabs > .active > a:hover {
        color: white;
      }
      .slider-container {
        display: flex;
        justify-content: center;
      }
      .slider-container .shiny-input-container {
        width: 80%;
      }
      #progress {
        color: white;
        margin-top: 10px;
      }
      #update {
        background-color: #337ab7;
        color: white;
        font-weight: bold;
        font-size: 16px;
      }
.well {
  min-height: 20px;
  padding: 19px;
  margin-bottom: 20px;
  background-color: #15151500;
  border: 1px solid #030303;
  border-radius: 4px;
  -webkit-box-shadow: inset 0 1px 1px rgba(0,0,0,0.05);
  box-shadow: inset 0 1px 1px rgba(0,0,0,0.05);
  backdrop-filter: blur(3px);
}
img {
filter:invert(0);
transition:0.5s;
}
img:hover {
filter:invert(1);
transition:0.5s;
}
.shiny-notification {
    position: fixed;
    background-color: #e8e8e8;
    color: #333;
    border: 1px solid #ccc;
    border-radius: 3px;
    opacity: .85;
    padding: 10px 2rem 10px 10px;
    margin: 5px;
    top: 333px;
    left:300px;
}
body {

    background-color: #06060600;
}
.js-plotly-plot .plotly .main-svg {
    position: absolute;
    top: 0px;
    left: 0px;
    pointer-events: none;
    filter: invert(1);
    transition:0.5s;
}
.js-plotly-plot .plotly .main-svg:hover {
    position: absolute;
    top: 0px;
    left: 0px;
    pointer-events: none;
    filter: invert(0);
}
    ")),
    tags$script(src = "https://polyfill.io/v3/polyfill.min.js?features=es6"),
    tags$script(src = "https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"),
    tags$script(HTML("
      MathJax.Hub.Config({
        tex2jax: {inlineMath: [['$','$'], ['\\(','\\)']]}
      });
    "))
  ),
  
  useShinyjs(),  # Initialize shinyjs for better UI interactions
  titlePanel("Professional Prophet Forecasting"),
  
  sidebarLayout(
    sidebarPanel(
      actionButton("update", "Update Forecast", class = "btn btn-primary"),
      div(id = "progress", ""),  # Progress display here
      sliderInput("periods_slider", "Forecast Period (days):", min = 30, max = 730, value = 365) %>%
        shinyInput_label_embed(
          shiny_iconlink() %>%
            bs_embed_tooltip(title = "Number of days to forecast into the future.")
        ),
      numericInput("changepoint_prior_scale", "Changepoint Prior Scale (\\alpha)", 0.05) %>%
        shinyInput_label_embed(
          shiny_iconlink() %>%
            bs_embed_tooltip(title = "Adjust the flexibility of the trend (\\alpha). Higher values allow more changepoints.")
        ),
      numericInput("seasonality_prior_scale", "Seasonality Prior Scale (\\beta)", 10.0) %>%
        shinyInput_label_embed(
          shiny_iconlink() %>%
            bs_embed_tooltip(title = "Adjust the strength of seasonality (\\beta). Higher values allow stronger seasonal effects.")
        ),
      numericInput("holidays_prior_scale", "Holidays Prior Scale (\\gamma)", 10.0) %>%
        shinyInput_label_embed(
          shiny_iconlink() %>%
            bs_embed_tooltip(title = "Adjust the influence of holidays (\\gamma). Higher values allow holidays to have more impact.")
        ),
      numericInput("yearly_seasonality", "Yearly Seasonality", 10) %>%
        shinyInput_label_embed(
          shiny_iconlink() %>%
            bs_embed_tooltip(title = "Number of Fourier terms to capture yearly seasonality.")
        ),
      numericInput("weekly_seasonality", "Weekly Seasonality", 3) %>%
        shinyInput_label_embed(
          shiny_iconlink() %>%
            bs_embed_tooltip(title = "Number of Fourier terms to capture weekly seasonality.")
        ),
      numericInput("daily_seasonality", "Daily Seasonality", 0) %>%
        shinyInput_label_embed(
          shiny_iconlink() %>%
            bs_embed_tooltip(title = "Number of Fourier terms to capture daily seasonality.")
        ),
      checkboxInput("include_history", "Include History in Forecast", TRUE) %>%
        shinyInput_label_embed(
          shiny_iconlink() %>%
            bs_embed_tooltip(title = "Include the historical data in the forecast plot.")
        ),
      radioButtons("method", "Computation Method:",
                   choices = list("Heuristic (fast)" = "heuristic", "Analytic (exact)" = "analytic"),
                   selected = "heuristic") %>%
        shinyInput_label_embed(
          shiny_iconlink() %>%
            bs_embed_tooltip(title = "Choose between heuristic (fast) and analytic (exact) computations.")
        ),
      width = 3
    ),
    
    mainPanel(
      tabsetPanel(
        tabPanel(
          "Forecast",
          withSpinner(plotlyOutput("forecastPlot"), type = 3, color.background = "white"),
          uiOutput("currentParamsMathJax"),
          verbatimTextOutput("modelSummary"),
          verbatimTextOutput("modelParameters"),
          verbatimTextOutput("modelResults")
        ),
        tabPanel(
          "Additional Plots",
          withSpinner(plotlyOutput("residualsPlot"), type = 3, color.background = "white"),
          withSpinner(plotlyOutput("acfPlot"), type = 3, color.background = "white"),
          withSpinner(plotlyOutput("pacfPlot"), type = 3, color.background = "white"),
          withSpinner(plotlyOutput("seasonalDecomposition"), type = 3, color.background = "white"),
          withSpinner(plotlyOutput("uncertaintyIntervals"), type = 3, color.background = "white")
        ),
        tabPanel(
          "Model Performance",
          withSpinner(DTOutput("modelPerformanceMetrics"), type = 3, color.background = "white")
        ),
        tabPanel(
          "Forecast Data",
          withSpinner(DTOutput("forecastTable"), type = 3, color.background = "white")
        ),
        tabPanel(
          "Components",
          withSpinner(plotlyOutput("componentsPlot"), type = 3, color.background = "white")
        ),
        tabPanel(
          "Current Parametrisation",
          verbatimTextOutput("currentParams"),
          tags$div(
            h4("Current Model:"),
            tags$p(HTML("
              \\[
              y(t) = g(t) + s(t) + h(t) + \\epsilon_t
              \\]
              \\[
              g(t) = \\text{Trend} = \\alpha + \\beta \\cdot t
              \\]
              \\[
              s(t) = \\text{Seasonality} = \\sum_k \\gamma_k \\cdot S_k(t)
              \\]
              \\[
              h(t) = \\text{Holidays} = \\sum_l \\delta_l \\cdot H_l(t)
              \\]
            "))
          )
        )
      )
    )
  )
)

# Define server logic
server <- function(input, output, session) {
  dataInput <- reactive({
    return(default_data)
  })
  
  modelFit <- reactive({
    input$update
    isolate({
      # Create a Progress object
      progress <- shiny::Progress$new()
      # Make sure it closes when we exit this reactive, even if there's an error
      on.exit(progress$close())
      
      progress$set(message = "Fitting model", value = 0)
      
      data <- dataInput()
      progress$inc(0.2, detail = "Preparing data")
      
      m <- prophet(
        changepoint.prior.scale = input$changepoint_prior_scale,
        seasonality.prior.scale = input$seasonality_prior_scale,
        holidays.prior.scale = input$holidays_prior_scale,
        yearly.seasonality = input$yearly_seasonality,
        weekly.seasonality = input$weekly_seasonality,
        daily.seasonality = input$daily_seasonality
      )
      progress$inc(0.4, detail = "Fitting model to data")
      
      m <- fit.prophet(m, data)
      progress$inc(0.4, detail = "Model fitting complete")
      
      return(m)
    })
  })
  
  forecastData <- reactive({
    input$update
    isolate({
      # Create a Progress object
      progress <- shiny::Progress$new()
      # Make sure it closes when we exit this reactive, even if there's an error
      on.exit(progress$close())
      
      progress$set(message = "Creating forecast", value = 0)
      
      m <- modelFit()
      progress$inc(0.2, detail = "Preparing future dataframe")
      
      future <- make_future_dataframe(m, periods = input$periods_slider, include_history = input$include_history)
      progress$inc(0.4, detail = "Predicting future values")
      
      forecast <- predict(m, future)
      progress$inc(0.4, detail = "Forecasting complete")
      
      return(forecast)
    })
  })
  
  output$forecastPlot <- renderPlotly({
    forecast <- forecastData()
    m <- modelFit()
    p <- plot(m, forecast) +
      ggtitle("Prophet Forecast") +
      xlab("Time") +
      ylab("Mortality Rate")
    ggplotly(p)
  })
  
  output$componentsPlot <- renderPlotly({
    forecast <- forecastData()
    m <- modelFit()
    components <- prophet_plot_components(m, forecast)
    plots <- lapply(components, ggplotly)
    plotly::subplot(plots, nrows = length(plots), shareX = TRUE)
  })
  
  output$modelPerformanceMetrics <- renderDT({
    input$update
    isolate({
      # Create a Progress object
      progress <- shiny::Progress$new()
      # Make sure it closes when we exit this reactive, even if there's an error
      on.exit(progress$close())
      
      progress$set(message = "Performing cross-validation", value = 0)
      
      m <- modelFit()
      method <- input$method
      df_cv_future <- future({
        if (method == "heuristic") {
          cross_validation(m, initial = 365, period = 180, horizon = 90, units = "days")
        } else {
          cross_validation(m, initial = 365*2, period = 180, horizon = 180, units = "days")
        }
      })
      df_cv <- value(df_cv_future)
      progress$inc(0.5, detail = "Calculating performance metrics")
      
      df_p <- performance_metrics(df_cv)
      progress$inc(0.5, detail = "Cross-validation complete")
      
      datatable(df_p, options = list(pageLength = 10, autoWidth = TRUE, scrollX = TRUE))
    })
  })
  
  output$modelParameters <- renderPrint({
    input$update
    isolate({
      m <- modelFit()
      print(m)
    })
  })
  
  output$modelResults <- renderPrint({
    input$update
    isolate({
      forecast <- forecastData()
      print(head(forecast, 10))
    })
  })
  
  output$modelSummary <- renderPrint({
    input$update
    isolate({
      m <- modelFit()
      summary <- capture.output(summary(m))
      cat(paste(summary, collapse = "\n"))
    })
  })
  
  output$forecastTable <- renderDT({
    forecast <- forecastData()
    datatable(forecast, options = list(pageLength = 10, autoWidth = TRUE, scrollX = TRUE))
  })
  
  output$seasonalDecomposition <- renderPlotly({
    forecast <- forecastData()
    if (nrow(forecast) > 1 && !is.null(forecast$yhat)) {
      ts_data <- ts(forecast$yhat, frequency = 12)
      components <- decompose(ts_data)
      p <- autoplot(components) + ggtitle("Seasonal Decomposition")
      ggplotly(p)
    }
  })
  
  output$uncertaintyIntervals <- renderPlotly({
    forecast <- forecastData()
    if (nrow(forecast) > 0) {
      p <- ggplot(forecast, aes(x = ds)) +
        geom_line(aes(y = yhat), color = 'blue') +
        geom_ribbon(aes(ymin = yhat_lower, ymax = yhat_upper), alpha = 0.2) +
        ggtitle("Forecast with Uncertainty Intervals") +
        xlab("Date") +
        ylab("Forecasted Value")
      ggplotly(p)
    }
  })
  
  output$currentParams <- renderPrint({
    input$update
    isolate({
      params <- list(
        periods_slider = input$periods_slider,
        changepoint_prior_scale = input$changepoint_prior_scale,
        seasonality_prior_scale = input$seasonality_prior_scale,
        holidays_prior_scale = input$holidays_prior_scale,
        yearly_seasonality = input$yearly_seasonality,
        weekly_seasonality = input$weekly_seasonality,
        daily_seasonality = input$daily_seasonality,
        include_history = input$include_history,
        method = input$method
      )
      print(params)
    })
  })
  
  output$currentParamsMathJax <- renderUI({
    input$update
    isolate({
      params <- list(
        periods_slider = input$periods_slider,
        changepoint_prior_scale = input$changepoint_prior_scale,
        seasonality_prior_scale = input$seasonality_prior_scale,
        holidays_prior_scale = input$holidays_prior_scale,
        yearly_seasonality = input$yearly_seasonality,
        weekly_seasonality = input$weekly_seasonality,
        daily_seasonality = input$daily_seasonality,
        include_history = input$include_history,
        method = input$method
      )
      HTML(paste0(
        "<script type='math/tex; mode=display'>
        \\text{Current Parameters:} \\\\
        \\text{Forecast Period (days): ", params$periods_slider, "} \\\\
        \\alpha: ", params$changepoint_prior_scale, " \\\\
        \\beta: ", params$seasonality_prior_scale, " \\\\
        \\gamma: ", params$holidays_prior_scale, " \\\\
        \\text{Yearly Seasonality: ", params$yearly_seasonality, "} \\\\
        \\text{Weekly Seasonality: ", params$weekly_seasonality, "} \\\\
        \\text{Daily Seasonality: ", params$daily_seasonality, "} \\\\
        \\text{Include History: ", ifelse(params$include_history, "Yes", "No"), "} \\\\
        \\text{Method: ", params$method, "}
        </script>"
      ))
    })
  })
  
  updateProgress <- function(session, percent, message) {
    session$sendCustomMessage(type = 'update-progress', message = list(percent = percent, message = message))
  }
  
  observeEvent(input$update, {
    updateProgress(0, "Starting update...")
  })
  
  # Custom handler for progress updates
  session$onFlushed(function() {
    session$sendCustomMessage(type = 'update-progress', message = list(percent = 0, message = "Initializing..."))
  }, once = TRUE)
}

# JavaScript code to handle progress updates
js <- "
Shiny.addCustomMessageHandler('update-progress', function(message) {
  $('#progress').text('Progress: ' + message.percent + '% - ' + message.message);
});
"

# Run the application 
shinyApp(ui = ui, server = server)