#############################################################################
#                                                                           #
#  Step-by-Step Script to Calculate wind-modified Traffic Density buffers   #                                                             
#                                                                           #
#  Computes the traffic buffer index for each categorya and                 #
#  then calculates the normalization based on the whole                     #
#  categories.                                                              #
#                                                                           #
#############################################################################

# Required libraries
library(sf)
library(sp)
library(mapview)
library(dplyr)
library(tidyr)
library(purrr)

# Load traffic data
wdir <- "/esarchive/scratch/ajulian/scripts/Open_Data/Traffic_buffers/Data/"
wdir2 <-"/esarchive/scratch/ajulian/scripts/Open_Data/Traffic_MPCS/Results/"
wwind <- "/esarchive/scratch/ajulian/scripts/Meteo/Results/"
# Define buffer sizes and combination target
buffer_sizes <- c(5, 10, 15, 25, 50, 100, 300, 500, 1000, 2000, 5000, 10000)
combo_buffers <- c(100, 300, 500)
# Traffic indicators to loop over
#traffic_vars <- c("im", "imd", "tm", "tmd")

# Load data
traffic_ind <- readRDS(file.path(wdir2, "traffic_ind_CAT_2024p.RDS"))
refined_grid_sf <- readRDS(file.path(wdir, "stations_sf.RDS")) %>% st_as_sf()
OK_wind <- readRDS(file.path(wwind, "kriged_wind_AQMS_2024.RDS")) %>% st_as_sf()

# Transform CRS
template_crs <- st_crs(refined_grid_sf)
crs_utm <- "+proj=utm +zone=31 +ellps=intl +units=m +no_defs"
traffic_ind_sf <- st_as_sf(traffic_ind, coords = c("lon", "lat"), crs = template_crs) %>% st_transform(crs_utm)
refined_grid_sf <- st_transform(refined_grid_sf, crs_utm)
traffic_ind_info_sf <- st_transform(select(traffic_ind, imd_total, geom), crs_utm)

# --- Helper functions ---
adjust_buffer_distance <- function(base_buffer_dist, wind_speed) {
  base_buffer_dist * (1 + wind_speed / 10)
}

compute_elongation_factor <- function(wind_speed) {
  ifelse(wind_speed <= 1, 1, 1 + (0.75 * (wind_speed / 10)))
}

compute_translation_factor <- function(wind_speed, base_buffer_dist) {
  ifelse(wind_speed <= 1, 0, base_buffer_dist * (wind_speed / 10))
}

create_ellipsoid <- function(buffer, wind_direction, base_buffer_dist, wind_speed) {
  elongation_factor <- compute_elongation_factor(wind_speed)
  translation_factor <- compute_translation_factor(wind_speed, base_buffer_dist)
  
  scale_matrix <- matrix(c(cos(wind_direction), -sin(wind_direction), sin(wind_direction), cos(wind_direction)), nrow = 2)
  scale_matrix <- scale_matrix %*% diag(c(1 / elongation_factor, elongation_factor)) %*% t(scale_matrix)
  
  coords <- st_coordinates(buffer)[, 1:2]
  center <- st_coordinates(st_centroid(buffer))[, 1:2]
  coords <- sweep(coords, 2, center, "-") %*% scale_matrix
  coords <- sweep(coords, 2, center, "+")
  
  translation_vector <- c(sin(wind_direction), cos(wind_direction)) * translation_factor
  coords <- sweep(coords, 2, translation_vector, "+")
  coords <- rbind(coords, coords[1, ])
  
  st_polygon(list(coords))
}

compute_traffic_density_dispersion <- function(grid_points, wind_speeds, wind_directions, base_buffer_dist, traffic_ind_sf) {
  adjusted_buffer_dists <- adjust_buffer_distance(base_buffer_dist, wind_speeds)
  
  ellipsoids <- lapply(seq_along(grid_points), function(i) {
    tryCatch({
      buffer <- st_buffer(grid_points[i], dist = adjusted_buffer_dists[i])
      ellipsoid <- create_ellipsoid(buffer, wind_directions[i], adjusted_buffer_dists[i], wind_speeds[i])
      st_sfc(ellipsoid, crs = st_crs(grid_points))
    }, error = function(e) {
      message(paste("Error at point", i, ":", e$message))
      NULL
    })
  }) %>% discard(is.null)
  
  ellipsoids_sfc <- st_sfc(do.call(c, ellipsoids), crs = st_crs(grid_points))
  ellipsoids_sfc <- st_transform(ellipsoids_sfc, st_crs(traffic_ind_sf))
  
  traffic_density_indices <- map_dbl(seq_along(ellipsoids_sfc), function(i) {
    roads_context <- st_intersection(traffic_ind_sf, ellipsoids_sfc[i])
    lengths <- as.numeric(st_length(roads_context))
    weights <- roads_context$imd_total * lengths
    
    total_length <- sum(lengths, na.rm = TRUE)
    total_weighted <- sum(weights, na.rm = TRUE)
    
    if (total_length > 0) {
      total_weighted / total_length
    } else {
      0    # ← return 0 when there is no intersection (instead of NA_real_)
    }
  })
  
  list(indices = traffic_density_indices, buffers = ellipsoids_sfc)
}

normalize <- function(values) {
  rng <- range(values, na.rm = TRUE)
  if (diff(rng) == 0) return(rep(0, length(values)))
  (values - rng[1]) / diff(rng) * 100
}

# --- Compute all buffer sizes ---
results <- map(buffer_sizes, function(buf_size) {
  message("Computing for buffer size: ", buf_size, "m")
  compute_traffic_density_dispersion(
    st_geometry(refined_grid_sf),
    OK_wind$Speed,
    OK_wind$Direction,
    buf_size,
    traffic_ind_info_sf
  )
})

# Add normalized results to sf
for (i in seq_along(buffer_sizes)) {
  col_name <- paste0("buffer_", buffer_sizes[i], "m_dispersion")
  norm_name <- paste0("buffer_", buffer_sizes[i], "m_normalized")
  refined_grid_sf[[col_name]] <- results[[i]]$indices
  refined_grid_sf[[norm_name]] <- normalize(refined_grid_sf[[col_name]])
}

# --- Composite index for 100, 300, 500 m ---
combo_names <- paste0("buffer_", combo_buffers, "m_normalized")
missing <- setdiff(combo_names, names(refined_grid_sf))
if (length(missing) > 0) stop("Missing required normalized buffer columns: ", paste(missing, collapse = ", "))

# Compute composite index using area-weighted average
area_weights <- 1 / (pi * combo_buffers^2)
normalized_weights <- area_weights / sum(area_weights)

df_buf <- refined_grid_sf %>%
  st_drop_geometry() %>%
  select(all_of(combo_names)) %>%
  mutate(across(everything(), as.numeric))

composite_index <- as.numeric(as.matrix(df_buf) %*% normalized_weights)

# Add composite index and remove originals
refined_grid_sf <- refined_grid_sf %>%
  mutate(traffic_100_300_500 = composite_index) 
  #%>%
  #select(-all_of(combo_names))

# Save output
saveRDS(refined_grid_sf, file.path(wdir2, "Wind_traffic_buffers_AQMS_combined.RDS"))