#### From Lagrangian to Eulerian perspectives
#### Script: Computing KDEs
#### Last modification: 15 02 2024
#### R version 4.2.3 (2023-03-15 ucrt) -- "Shortstop Beagle"


library(tidyverse)
library(sf)
library(terra)
library(virtualecologist)
library(ggspatial)
library(tidyterra)
library(ggthemes)
library(viridis)
library(raster)

workdir <- "your working directory" 


#### ---------------------------------- ####
#### --- Create virtual environment --- ####
#### ---------------------------------- ####

## Create virtual space
grid <- create_grid()

## Generate env conditions
env <- generate_env_layer(norm = TRUE, return_rasters = TRUE, grid = grid,
                          seed = 2)
plot(env$rasters)

## Generate resource layer for central place forager
rsce_1 <- generate_resource_layer(env_layers = env$rasters,
                                  beta = c(2, -1.5))
plot(rsce_1$rasters)

## Generate resource layer for free ranger
rsce_2 <- generate_resource_layer(env_layers = env$rasters,
                                  beta = c(-2, 0.5))
plot(rsce_2$rasters)

## Compute distance from colonies 
coord_col <- data.frame(Lon = c(30,60),
                        Lat = c(60,30),
                        Col = c("1", "2"))

dist_col <- raster::distanceFromPoints(raster::raster(rsce_1$rasters), xy = coord_col[,1:2])
names(dist_col) <- "dist_col"


#### ------------------------------ ####
#### ---- KDE estimation - CPF ---- ####
#### ------------------------------ ####

#### 1 - Open trajectories -----------------------------------------------------
load(paste0(workdir, "/Simulated_tracks_CPF.rda"))
traj_CPF <- traj_CPF |> 
  mutate(ID = paste0("Col", Colony, ind_id), # one id per bird, pooling both days (BirdId is per track)
         DateTime = ymd_hms(ymd_hms))

### 2 - Transform trajectories as spatial points -------------------------------
traj_proj <- traj_CPF %>% 
  mutate(Longitude = Lon, Latitude = Lat) %>%
  st_as_sf(.,coords = c("Longitude","Latitude"), crs = virtualecologist::custom_proj())

traj_CPF %>% group_by(Colony) %>% summarize(n = length(unique(ID)))
unique_col <- c("Colony 1", "Colony 2")


### 3 - Prepare ID vectors ----------------------------------------------------
library(adehabitatHR)
id.tot <- unique(traj_CPF$ID)

id.col1 <- unique(traj_CPF |> subset(Colony == "1") |> pull("ID"))
id.col2 <- unique(traj_CPF |> subset(Colony == "2") |> pull("ID"))

# transform to sp objects
traj_proj_sp <- as_Spatial(traj_proj[, c("ID", "geometry")])

# make space in memory if necessary
# rm(traj_CPF, traj_proj)


### 4 - Compute KDE for Colony 1 -----------------------------------------------
# convert each traj to raster, stack all (too heavy to do all at once)
rast_stack1 <- NULL
for(i in cli::cli_progress_along(1:8000, "Processing KDE for Colony 1")){
  # compute kernel estimate
  kde <- kernelUD(traj_proj_sp[which(traj_proj_sp@data$ID == id.col1[i]),], 
                  grid = as(raster(env$rasters$sim1), "SpatialPixels"))
  
  # Extract the UD values and coordinates into a data frame
  ud_grid <- data.frame("value" = kde[[1]]$ud, "lon" = kde[[1]]@coords[,1], "lat" = kde[[1]]@coords[,2])
  coordinates(ud_grid)<-~lon+lat
  
  # coerce to SpatialPixelsDataFrame
  gridded(ud_grid) <- TRUE
  
  # coerce to raster
  rast <- raster(ud_grid)
  rast <- rast / cellStats(rast, "sum")
  
  rast_stack1 <- rast(list(rast_stack1, terra::rast(rast)))
}
names(rast_stack1) <- id.col1
# kde_col1 <- raster::overlay(rast_stack1, fun = "mean")
kde_col1 <- mean(rast_stack1, na.rm = T)
plot(kde_col1*8000)

writeRaster(x = kde_col1, paste0(workdir, "/CPF_KDE_Col1.grd"), overwrite = T)

save.image("KDE_Col.RData")



### 5 - Compute KDE for Colony 2 -----------------------------------------------
rast_stack2 <- NULL
for(i in cli::cli_progress_along(1:10000, "Processing KDE for Colony 2")){
  # compute kernel estimate
  kde <- kernelUD(traj_proj_sp[which(traj_proj_sp@data$ID == id.col2[i]),], 
                  grid = as(raster(env$rasters$sim1), "SpatialPixels"))
  # kde[[1]]@h
  
  # Extract the UD values and coordinates into a data frame
  ud_grid <- data.frame("value" = kde[[1]]$ud, "lon" = kde[[1]]@coords[,1], "lat" = kde[[1]]@coords[,2])
  coordinates(ud_grid)<-~lon+lat
  
  # coerce to SpatialPixelsDataFrame
  gridded(ud_grid) <- TRUE
  
  # coerce to raster
  rast <- raster(ud_grid)
  rast <- rast / cellStats(rast, "sum")
  
  rast_stack2 <- rast(list(rast_stack2, terra::rast(rast)))
}
names(rast_stack2) <- id.col2
# kde_col2 <- raster::overlay(rast_stack2, fun = "mean")
kde_col2 <- mean(rast_stack2, na.rm = T)
plot(kde_col2*10000)

writeRaster(x = kde_col2, paste0(workdir, "/CPF_KDE_Col2.grd"), overwrite = T)

save.image("KDE_Col.RData")



### 5 - Combine the results ----------------------------------------------------
kde_col1 <- rast("KDE_Col1.grd")
kde_col2 <- rast("KDE_Col2.grd")

nind_tot <- kde_col1*8000 + kde_col2*10000

plot(nind_tot, col = viridis(50,option = "H"), addfun = c(points(coord_col[, 1:2], col = "white")))  
plot(log(nind_tot), col = viridis(50,option = "H"), addfun = c(points(coord_col[, 1:2], col = "white")))

# convert to polygon contour : convert to % of use
true.kde.reclass <- cut(raster(nind_tot), 
                        breaks = c(0, 
                                   quantile(raster(nind_tot), probs = 0.5),
                                   quantile(raster(nind_tot), probs = 0.8), 
                                   quantile(raster(nind_tot), probs = 0.9), 
                                   quantile(raster(nind_tot), probs = 1)))
contour <- rasterToContour(true.kde.reclass, levels = c(1,2,3,4)) %>%
  st_as_sf() %>%
  mutate(Class = case_when(level == "2" ~ "[50,80%)",
                           level == "3" ~ "[80,90%)",
                           level == "4" ~ "[90,100%)"))

plot(nind_tot) ; lines(as_Spatial(contour))

save.image("KDE_Col.RData")


# save
writeRaster(x = nind_tot,
            filename = paste0(workdir, "/Simulated_tracks_CPF_KDE_totpop.grd"), overwrite = T)
st_write(contour, paste0(workdir, "/Simulated_tracks_CPF_KDE_totpop_poly.shp"))




### 4 - plot -------------------------------------------------------------------
ggarrange(
  ggplot(as.data.frame(kde_col1, xy = T)) +
    geom_tile(aes(x = x, y = y, fill = mean)) +
    scale_fill_viridis(option = "magma", na.value = NA, direction = -1) +
    geom_point(data = coord_col[1,], aes(x = Lon, y = Lat), size = 2, shape = 4, color = "black") +
    coord_sf(xlim = c(0,90), ylim = c(0,90)) +
    theme_bw() + labs(fill = "p. prob", title = "A - UD Colony 1")
  ,
  ggplot(as.data.frame(kde_col2, xy = T)) +
    geom_tile(aes(x = x, y = y, fill = mean)) +
    scale_fill_viridis(option = "magma", na.value = NA, direction = -1) +
    geom_point(data = coord_col[2,], aes(x = Lon, y = Lat), size = 2, shape = 4, color = "black") +
    coord_sf(xlim = c(0,90), ylim = c(0,90)) +
    theme_bw() + labs(fill = "p. prob", title = "B - UD Colony 2")
  ,
  ggplot(as.data.frame(nind_tot, xy = T)) +
    geom_tile(aes(x = x, y = y, fill = mean)) +
    scale_fill_viridis(option = "mako", na.value = NA, direction = -1) +
    geom_sf(data = contour, color = "pink", aes(linetype = Class)) +
    geom_point(data = coord_col, aes(x = Lon, y = Lat), size = 2, shape = 4, color = "black") +
    coord_sf(xlim = c(0,90), ylim = c(0,90)) +
    theme_bw() + labs(fill = "N ind", linetype = "Proportion of population", title = "D - UD total population")
  , ncol = 3, nrow = 1, align = "hv") + border("white") + bgcolor("white")

ggsave(paste0(workdir, "/Simulated_tracks_CPF_KDE_totpop.png"), dpi = 350, width = 15, height = 5)





#### ----------------------------- ####
#### ---- KDE estimation - FR ---- ####
#### ----------------------------- ####

#### 1 - Open trajectories -----------------------------------------------------
load(paste0(workdir, "/Simulated_tracks_FR.rda"))
traj_FR <- traj_FR |> 
  mutate(ID = ind_id,
         DateTime = ymd_hms(ymd_hms))

### 2 - Transform trajectories as spatial points -------------------------------
traj_FR_proj <- traj_FR %>% 
  mutate(Longitude = Lon, Latitude = Lat) %>%
  st_as_sf(.,coords = c("Longitude","Latitude"), crs = virtualecologist::custom_proj())

traj_FR %>% summarize(n = length(unique(ID)))


### 3 - Prepare ID vectors ----------------------------------------------------
id.tot.fr <- unique(traj_FR$ID)

# transform to sp objects
traj_FR_proj_sp <- as_Spatial(traj_FR_proj[, c("ID", "geometry")])

# make space in memory
# rm(traj_CPF, traj_proj)


### 4 - Compute KDE ------------------------------------------------------------
# convert each traj to raster, stack all
FR_rast_stack1 <- NULL
for(i in cli::cli_progress_along(1:8000, "Processing KDE for FR")){
  # compute kernel estimate
  kde <- kernelUD(traj_FR_proj_sp[which(traj_FR_proj_sp@data$ID == id.tot.fr[i]),], 
                  grid = as(raster(env$rasters$sim1), "SpatialPixels"))
  # kde[[1]]@h
  
  # Extract the UD values and coordinates into a data frame
  ud_grid <- data.frame("value" = kde[[1]]$ud, "lon" = kde[[1]]@coords[,1], "lat" = kde[[1]]@coords[,2])
  coordinates(ud_grid)<-~lon+lat
  
  # coerce to SpatialPixelsDataFrame
  gridded(ud_grid) <- TRUE
  
  # coerce to raster
  rast <- raster(ud_grid)
  rast <- rast / cellStats(rast, "sum")
  
  FR_rast_stack1 <- rast(list(FR_rast_stack1, terra::rast(rast)))
}
names(FR_rast_stack1) <- id.tot.fr
FR_kde <- mean(FR_rast_stack1, na.rm = T)
plot(FR_kde*8000)

writeRaster(x = FR_kde, paste0(workdir, "/FR_KDE.grd"), overwrite = T)

save.image("KDE_Col.RData")



### 5 - Convert to ind numbers -------------------------------------------------
FR_kde <- rast("FR_KDE.grd")

FR_nind_tot <- FR_kde*8000

plot(FR_nind_tot, col = viridis(50,option = "H"))  
plot(log(FR_nind_tot), col = viridis(50,option = "H"))

# convert to polygon contour : convert to % of use 
fr.true.kde.reclass <- cut(raster(FR_nind_tot), 
                        breaks = c(0, 
                                   quantile(raster(FR_nind_tot), probs = 0.5),
                                   quantile(raster(FR_nind_tot), probs = 0.8), 
                                   quantile(raster(FR_nind_tot), probs = 0.9), 
                                   quantile(raster(FR_nind_tot), probs = 1)))
FR_contour <- rasterToContour(fr.true.kde.reclass, levels = c(1,2,3,4)) %>%
  st_as_sf() %>%
  mutate(Class = case_when(level == "2" ~ "[50,80%)",
                           level == "3" ~ "[80,90%)",
                           level == "4" ~ "[90,100%)"))


plot(FR_nind_tot) ; lines(as_Spatial(FR_contour))

save.image("KDE_Col.RData")


# save
writeRaster(x = FR_nind_tot,
            filename = paste0(workdir, "/Simulated_tracks_FR_KDE_totpop.grd"), overwrite = T)
st_write(FR_contour, paste0(workdir, "/Simulated_tracks_FR_KDE_totpop_poly.shp"))




### 4 - plot -------------------------------------------------------------------
ggplot(as.data.frame(FR_nind_tot, xy = T)) +
  geom_tile(aes(x = x, y = y, fill = mean)) +
  scale_fill_viridis(option = "mako", na.value = NA, direction = -1) +
  geom_sf(data = FR_contour, color = "pink", aes(linetype = Class)) +
  coord_sf(xlim = c(0,90), ylim = c(0,90)) +
  theme_bw() + labs(fill = "N ind", linetype = "Proportion of population", title = "UD total population")
ggsave(paste0(workdir, "/Simulated_tracks_FR_KDE_totpop.png"), dpi = 350, width = 5, height = 5)

