# -----------------------------------------------------------------------------
#   STEP 1: prepare environment 

# Set directory from command line argument
args <- commandArgs(trailingOnly = TRUE)
if (length(args) > 0) {
  current_dir <- args[1]
} else {
  current_dir <- dirname(rstudioapi::getSourceEditorContext()$path)
}
parent_dir <- dirname(current_dir)
data_path <- file.path(parent_dir, "Data")
interim_path <- file.path(data_path, "Interim")
output_path <- file.path(parent_dir, "Output", "figures")

# Install and load packages
package_list <- c("igraph", "readxl", "sfnetworks", "maps", "ggplot2", "rworldmap", "dplyr")
installed_packages <- installed.packages()[, "Package"]
missing_packages <- package_list[!(package_list %in% installed_packages)]

if (length(missing_packages) > 0) {
  install.packages(missing_packages, dependencies = TRUE, repos = "https://cloud.r-project.org")
}

load_status <- vapply(
  package_list,
  require,
  logical(1),
  character.only = TRUE,
  quietly = TRUE
)
if (!all(load_status)) {
  stop(
    "Failed to load required package(s): ",
    paste(package_list[!load_status], collapse = ", "),
    call. = FALSE
  )
}

rm(current_dir, parent_dir, package_list, installed_packages, missing_packages, load_status)

# -----------------------------------------------------------------------------
#   STEP 2: Load and prepare data

setwd(data_path)

network_data <- read_excel("liquidity_lines.xlsx", sheet = "LiquidityLines")
country_data <- read_excel("liquidity_lines.xlsx", sheet = "List")
EA_data <- read_excel("liquidity_lines.xlsx", sheet = "EuroArea")

# bilateral network
setwd(interim_path)
countries <- c("USA", "EA", "CHN")
paths_filenames <- paste0("monthly_full_data_yes_", countries, "_reduced.xlsx")

# EA membership
years <- c(2025, 2023, 2020, 2009, 2000)
for (year in years) {
  column_name <- paste("EA_member_", year, sep = "")
  EA_data <- EA_data %>%
    mutate(!!column_name := ifelse(as.Date(joining_date) <= as.Date(paste(year, "-01-01", sep = "")), 1, 0))
}
rm(year, column_name)

# Initialize a list to store the dataframes for each year
data_list <- list()

# Iterate over the countries
# Take shortest path at any point in year
for (i in 1:length(countries)) {
  year <- 2025
  path_filename <- paths_filenames[i]
  
  # Load paths data for the current year
  paths_data <- read_xlsx(path_filename, range = "A1:D19297")
  paths_data <- paths_data %>% filter(Year == year)
  paths_data <- paths_data %>% select(-c("Year", "Month"))
  
  colname_new <- paste("min_dist_", countries[i], sep = "")
  colname_old <- paste("dist_", countries[i], sep = "")
  
  paths_data <- paths_data %>%
    group_by(country_ID) %>%
    mutate(
      !!colname_new := ifelse(all(is.na(!!sym(colname_old))), NA_real_, min(!!sym(colname_old), na.rm = TRUE))) %>%
    ungroup() %>%
    select(-colname_old)
  
  paths_data <- distinct(paths_data)
  
  # Merge country data with paths for the current year
  merged_data <- merge(country_data, paths_data, by = "country_ID")
  
  # Filter for EA members at cutoff date
  filtered_rows <- filter(EA_data, get(paste0("EA_member_", year)) == 1)
  
  # Drop countries that are in the EA but also have separate swap agreements
  countries_to_drop <- filtered_rows$country
  merged_data <- merged_data[!(merged_data$country %in% countries_to_drop), ]
  
  # Select the EA countries at cutoff date from filtered_rows
  selected_rows <- select(filtered_rows, country, country_code)
  
  # Append data for EA countries
  adjusted_data <- bind_rows(merged_data, selected_rows)
  
  # Fill the remaining columns for the newly added EA countries with connectivity data from the "EA" row
  EA_row <- filter(merged_data, country == "EA")
  exclude_cols <- c("country", "country_code")
  
  for (col in colnames(adjusted_data)) {
    if (col %in% exclude_cols) next
    adjusted_data[!(adjusted_data$country %in% merged_data$country), col] <- EA_row[[col]]
  }
  
  adjusted_data <- adjusted_data %>% select(-c("currency", "bank_code", "country", "centralbank", "currency", "currency_value"))
  
  # Store the adjusted data in the list
  data_list[[i]] <- adjusted_data
  rm(colname_new, colname_old, year, adjusted_data, selected_rows, countries_to_drop, filtered_rows, merged_data, paths_data, path_filename, exclude_cols, EA_row, col)
}

# Bind data into single frame
merged_data <- data.frame(data_list)
merged_data <- merged_data %>% select(-ends_with(".1"), -ends_with(".2"), -ends_with(".3"))

# Calculate max distance (for labeling later)
i <- 1
for (country in countries) {
  country_3rd_years <- c()
  for (year in years) {
    dist_country <- data_list[[i]][[paste0("min_dist_", country)]]
    max_value <- max(dist_country, na.rm = TRUE)
    if (max_value == 3) {
      country_3rd_years <- c(country_3rd_years, year)
    }
  }
  assign(paste0(country, "_3rd_years"), country_3rd_years)
  i <- i + 1
}

# Create short version (only direct connections)
merged_data_short <- merged_data %>%
  mutate(across(c(min_dist_USA, min_dist_CHN, min_dist_EA), ~ifelse(. %in% c(0, 1), ., NA)))

# -----------------------------------------------------------------------------
#   STEP 3: Generate maps

setwd(output_path)

# Function to create and save map
create_network_map <- function(data, country, year, direct_only = FALSE) {
  
  # EA and Kosovo (XKX) are not matched by the map lookup; EA members are
  # filled individually earlier so the world map remains complete.
  Map_obj <- joinCountryData2Map(data, joinCode = "ISO3", nameJoinColumn = "country_code")
  Map_filtered <- Map_obj[Map_obj$ISO3 != "ATA", ]
  
  col_name <- paste0("min_dist_", country)
  
  if (direct_only) {
    filename_base <- paste0("direct_network_bilateral", year, "_", country)
    colorPalette <- c("black", "blue")
    legendText <- c("source", "1st degree")
  } else {
    filename_base <- paste0("network_bilateral", year, "_", country)
    colorPalette <- c("black", "blue", "deepskyblue", "lightblue")
    legendText <- c("source", "1st degree", "2nd degree", "3rd degree", "unconnected")
  }
  
  # Function to draw the map
  draw_map <- function() {
    mapParams <- mapCountryData(
      Map_filtered,
      nameColumnToPlot = col_name,
      catMethod = "categorical",
      missingCountryCol = gray(0.6),
      colourPalette = colorPalette,
      addLegend = FALSE,
      mapTitle = ""
    )
    mapParams$legendText <- legendText
    do.call(addMapLegendBoxes, c(mapParams, x = "left", xpd = TRUE, cex = 0.7))
  }
  
  # Save as PDF
  pdf(paste0(filename_base, ".pdf"), width = 10, height = 7, useDingbats = FALSE)
  draw_map()
  dev.off()

  # Save as EPS
  if (capabilities("cairo")) {
    cairo_ps(
      filename = paste0(filename_base, ".eps"),
      width = 10,
      height = 7,
      onefile = FALSE,
      family = "Helvetica",
      fallback_resolution = 1200
    )
  } else {
    postscript(
      file = paste0(filename_base, ".eps"),
      width = 10,
      height = 7,
      onefile = FALSE,
      horizontal = FALSE,
      paper = "special",
      family = "Helvetica",
      colormodel = "srgb"
    )
  }
  draw_map()
  dev.off()
  
  # Save as PNG
  png(paste0(filename_base, ".png"), width = 1200, height = 800, res = 150)
  draw_map()
  dev.off()

  # Save as EPS
  setEPS()
  postscript(paste0(filename_base, ".eps"), width = 8, height = 5.33)
  draw_map()
  dev.off()

  cat("Saved:", filename_base, ".pdf, .eps and .png\n")
}

# Generate all maps
year <- 2025

for (country in countries) {
  # Full network map
  create_network_map(merged_data, country, year, direct_only = FALSE)
  # Direct connections only
  create_network_map(merged_data_short, country, year, direct_only = TRUE)
}
