# Create world map of altitude
# Author: Martín Lotto Batista - Barcelona Supercomputing Center

# Load required libraries
pacman::p_load(tidyverse, sf, raster, smoothr)

# World map shapefile
world_map <- st_read("data/world_shp/world_map_robinson.shp", quiet=TRUE)
robinson_proj <- "+proj=robin +lon_0=0 +x_0=0 +y_0=0 +ellps=WGS84 +datum=WGS84 +units=m +no_defs"

# Weather stations --------------------------------------------------------
# Read and transform station data
stations <- st_read("data/station_data/stations1.shp", quiet=TRUE) %>% 
  dplyr::select(WCID) %>% 
  st_transform(crs=st_crs(world_map))

st.data <- stations %>% 
  st_coordinates() %>% 
  as.data.frame() %>% 
  mutate(st=1:nrow(.))

# Get bounding box of the world map
bounding_rectangle <- st_as_sfc(st_bbox(world_map), crs=st_crs(stations))

# Plot the map with weather stations
plot_weather_stations <- ggplot() +
  geom_sf(data=world_map, fill="grey75", col="grey90") +
  geom_point(data=st.data, aes(x=X, y=Y, color="#D09E00"),
          alpha=0.5, size=0.05) +
  geom_sf(data=world_map, fill=NA, col="grey90") +
  scale_color_manual("", values=c("#D09E00"), labels=c("Weather\nstations")) +
  theme_void() +
  theme(axis.title=element_blank(),
        legend.position="right", 
        panel.grid.major=element_line(color="grey30", linetype="dashed", linewidth=0.1),
        panel.grid.minor=element_blank(),
        plot.margin=margin(0, 1, 0, 0, "cm")) +
  guides(color=guide_legend(override.aes=list(size=5, alpha=1)))

# Temperature map ---------------------------------------------------------
# Read temperature data and change to Robinson projection
# Note: Warnings are normal because there is not a temporal structure
temperature <- raster("data/temp_era5land_dec_2023.grib")
temperature_robinson <- projectRaster(temperature, crs=robinson_proj)

# Convert raster data to data frame
temperature_data <- raster::as.data.frame(temperature_robinson, xy=TRUE)
temperature_data <- na.omit(temperature_data)
names(temperature_data) <- c("x", "y", "t2m_global")
temperature_data[["t2m_global"]] <- temperature_data[["t2m_global"]]-273.15

# Define cuts for color breaks
cuts <- pretty(c(min(temperature_data[["t2m_global"]]), 
                 max(temperature_data[["t2m_global"]])), 
               n=8)

# Plot temperature map
temperature_plot <- temperature_data %>%
  mutate(cat_var=cut(t2m_global, breaks=cuts)) %>% 
  ggplot() +
  geom_tile(aes(x=x, y=y, fill=t2m_global)) +
  geom_sf(data=world_map, fill=NA, col="grey90") +
  scale_fill_gradient2(low="#02008E", high="#560100",
                       breaks=cuts,
                       labels=cuts, guide="legend") +
  labs(fill="Degree\nCelsius") +
  theme_void() +
  theme(axis.title=element_blank(),
        legend.position="right", 
        legend.direction="vertical",
        legend.title=element_text(face="bold"),
        legend.key.height=unit(0.15, "in"),
        panel.grid.major=element_line(color="grey30", linetype="dashed", linewidth=0.1),
        panel.grid.minor=element_blank(),
        plot.margin=margin(0, 1, 0, 0, "cm"))

# Elevation ---------------------------------------------------------------
# Read elevation data and change to Robinson projection
elevation <- raster("data/elev_0.1_deg.tif")
elevation_robinson <- projectRaster(elevation, crs=robinson_proj)

# Convert raster data to data frame
elevation_data <- raster::as.data.frame(elevation_robinson, xy=TRUE, na.rm=TRUE)
elevation_data <- elevation_data %>% 
  filter(elev_0.1_deg >= 0) %>% 
  mutate(elev=log(elev_0.1_deg))

# Define cuts for color breaks
cuts <- pretty(c(min(elevation_data[["elev_0.1_deg"]]), 
                 max(elevation_data[["elev_0.1_deg"]])), 
               n=8)

# Plot elevation map
elevation_plot <- ggplot() +
  geom_tile(data=elevation_data, 
            aes(x=x, y=y, fill=elev_0.1_deg)) +
  geom_sf(data=world_map, fill=NA, col="grey90") +
  scale_fill_gradient2(mid="#DA4C69", high="#FCFDBF",
                       breaks=cuts,
                       labels=cuts,
                       guide="legend") +
  labs(fill="Metres") +
  theme_void() +
  theme(axis.title=element_blank(),
        legend.position="right", 
        legend.direction="vertical",
        legend.title=element_text(face="bold"),
        legend.key.height=unit(0.15, "in"),
        panel.grid.major=element_line(color="grey30", linetype="dashed", linewidth=0.1),
        panel.grid.minor=element_blank(),
        plot.margin=margin(0, 1, 0, 0, "cm"))

# Environmental weights ---------------------------------------------------
# Process environmental weights
ws <- raster("data/weights.tif")
raster_robinson <- projectRaster(ws, crs=robinson_proj)

# Convert raster data to data frame
df <- raster::as.data.frame(raster_robinson, xy=TRUE, na.rm=TRUE)
df <- df %>% 
  mutate(weights=ifelse(weights < 0, 0, weights))

# Define cuts for color breaks
cuts <- pretty(c(min(na.omit(df[["weights"]])), 
                 max(na.omit(df[["weights"]]))), 
               n=7)

env_plot <- ggplot() +
  geom_sf(data=world_map, fill="grey75", col="grey90") +
  geom_tile(data=df, aes(x=x, y=y, fill=weights)) +
  scale_fill_gradient(low="lightgreen", high="darkgreen", 
                      na.value="white", 
                      breaks=cuts, labels=cuts, guide="legend") +
  labs(fill="Proportion", col="Proportion") +
  theme_void() +
  theme(axis.title=element_blank(),
        legend.position="right", 
        legend.direction="vertical",
        legend.title=element_text(face="bold"),
        legend.key.height=unit(0.15, "in"),
        panel.grid.major=element_line(color="grey30", linetype="dashed", linewidth=0.1),
        panel.grid.minor=element_blank(),
        plot.margin=margin(0, 1, 0, 0, "cm"))

# Flights -----------------------------------------------------------------
# Airport data
airport_info <- read.csv("data/flights_data/AirportInfo.csv")
prediction_data <- read.csv("data/flights_data/prediction.csv")

# Merge airport data with transportation prediction data
merged_data <- prediction_data %>%
  left_join(airport_info %>% 
              mutate(coords=paste0(Lon, "/", Lat)) %>% 
              dplyr::select(NodeName, coords),
            by=c("Origin"="NodeName")) %>% 
  mutate(Origin=coords) %>% 
  dplyr::select(-coords) %>% 
  separate(Origin, into=c("lon.o", "lat.o"), sep="/") %>% 
  left_join(airport_info %>% 
              mutate(coords=paste0(Lon, "/", Lat)) %>% 
              dplyr::select(NodeName, coords),
            by=c("Destination"="NodeName")) %>% 
  mutate(Destination=coords) %>% 
  dplyr::select(-coords) %>% 
  separate(Destination, into=c("lon.d", "lat.d"), sep="/") %>% 
  filter(PredMu > 100000)

# Extract coordinates from the world map
coordinates_df <- st_coordinates(world_map)[, c("X", "Y", "L3")] %>% 
  as.data.frame()

value_threshold <- 100000

# Get bounding box of the world map
bounding_box <- st_as_sfc(st_bbox(world_map), crs=st_crs(world_map)) %>% 
  st_coordinates() %>% 
  as.data.frame()

# Transform and filter prediction data for origin and destination airports
origin_airports <- st_as_sf(merged_data[merged_data[["PredMu"]] > value_threshold, ], 
                            coords=c("lon.o", "lat.o"), crs=4326) %>% 
  st_transform(crs=st_crs(world_map))

destination_airports <- st_as_sf(merged_data[merged_data[["PredMu"]] > value_threshold, ], 
                                 coords=c("lon.d", "lat.d"), crs=4326) %>% 
  st_transform(crs=st_crs(world_map))

# Extract coordinates for plotting
coordinates_origin <- st_coordinates(origin_airports) %>% as.data.frame()
coordinates_destination <- st_coordinates(destination_airports) %>% as.data.frame()

# Combine origin and destination coordinates
plot_data <- bind_cols(coordinates_origin, coordinates_destination) %>%
  magrittr::set_colnames(c("xo", "yo", "xd", "yd"))

# Plot the flights
flight_plot <- ggplot() +
  geom_sf(data=world_map, fill="grey75", col="grey90") +
  geom_curve(data=plot_data,
             aes(x=xo, y=yo, xend=xd, yend=yd), colour="#c994c7",
             linewidth=0.1) +
  geom_point(data=plot_data,
             aes(x=xo, y=yo), colour="#dd1c77",
             size=0.4) +
  geom_point(data=plot_data,
             aes(x=xd, y=yd, colour="#dd1c77"),
             size=0.4) +
  scale_colour_manual("",
                      values=c("#dd1c77"),
                      labels=c("Airport\nconnections")) +
  theme_void() +
  theme(axis.title=element_blank(),
        legend.position="right", 
        legend.direction="vertical",
        legend.title=element_text(face="bold"),
        panel.grid.major=element_line(color="grey30", 
                                      linetype="dashed", 
                                      linewidth=0.1),
        panel.grid.minor=element_blank(),
        plot.margin=margin(0, 1, 0, 0, "cm")) +
  guides(colour=guide_legend(override.aes=list(size=5, alpha=1)))

# API ---------------------------------------------------------------------
api <- read.csv("data/API_country_2017.csv")

# Match missing names
world_map <- world_map %>% 
  left_join(api %>% 
              mutate(Country=str_replace_all(Country, c("^Congo"="Republic of the Congo",
                                                        "Côte d’Ivoire"="Ivory Coast",
                                                        "United Republic of Tanzania"="Tanzania",
                                                        "Timor-Leste"="East Timor",
                                                        "Iran \\(Islamic Republic of\\)"="Iran",
                                                        "Gambia"="The Gambia"))) %>% 
              rename(name_en=Country))

api_plot <- ggplot() +
  geom_sf(data=world_map, aes(fill=API), col="grey90") +
  ggthemes::scale_fill_continuous_tableau(palette="Blue", 
                                          na.value="grey75",
                                          guide="legend") +
  theme_void() +
  theme(axis.title=element_blank(),
        legend.position="right", 
        legend.direction="vertical",
        legend.title=element_text(face="bold"),
        legend.key.height=unit(0.15, "in"),
        panel.grid.major=element_line(color="grey30", linetype="dashed", linewidth=0.1),
        panel.grid.minor=element_blank(),
        plot.margin=margin(0, 1, 0, 0, "cm"))

# Figure 3 ----------------------------------------------------------------
fig3 <- cowplot::plot_grid(plot_weather_stations, temperature_plot,
                           elevation_plot, env_plot, flight_plot, api_plot,
                           ncol=2,
                           labels=c("a.", "b.", "c.", "d.", "e.", "f."))

ggsave(plot=fig3, filename="outputs/figure3.pdf",
       width=14, height=12, units="in", dpi=300)

# Mosaic version
lfig3 <- cowplot::plot_grid(plot_weather_stations, temperature_plot,
                            elevation_plot, env_plot, flight_plot, api_plot,
                            nrow=2,
                            labels=c("a.", "b.", "c.", "d.", "e.", "f."))

ggsave(plot=lfig3, filename="outputs/landscape_figure3.jpeg",
       width=3000, height=3500, units="px")

# Figure 4 ----------------------------------------------------------------
# Newly suitable shapefiles
p_falc <- st_read("data/newly_suitable/p_falciparum_new.shp", quiet=TRUE)
p_vivax <- st_read("data/newly_suitable/p_vivax_new.shp", quiet=TRUE)

# Buffer areas in each map
p_falc <- st_buffer(p_falc, 10000)
p_vivax <- st_buffer(p_vivax, 10000)

plasmodium <- p_falc %>% 
  mutate(parasite="P. falciparum") %>% 
  bind_rows(p_vivax %>% 
              mutate(parasite="P. vivax"))

fig4 <- ggplot() +
  geom_sf(data=world_map, fill="grey75", col="grey90") +
  geom_sf(data=plasmodium, aes(fill=parasite), col=NA, alpha=0.8) +
  scale_fill_manual(values=c("P. falciparum"="#9B0B9E", "P. vivax"="orange")) +
  labs(fill="", col="") +
  theme_void() +
  theme(axis.title=element_blank(),
        legend.position="right", 
        legend.direction="vertical",
        legend.text=element_text(face="italic"),
        legend.title=element_text(face="bold"),
        legend.key.height=unit(0.15, "in"),
        panel.grid.major=element_line(color="grey30", linetype="dashed", linewidth=0.1),
        panel.grid.minor=element_blank(),
        plot.margin=margin(0, 1, 0, 0, "cm"))

ggsave(plot=fig4, filename="outputs/figure4.pdf",
       width=10, height=8, units="in", dpi=300)

## END