#### From Lagrangian to Eulerian perspectives
#### Script: Analysing survey data
#### 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(mgcv)
library(Distance)
library(ggpubr)

workdir <- "your working directory"
setwd(workdir)

#### ---------------------------------- ####
#### --- 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"

## create prediction table
env_stack <- raster::stack(env$rasters)
env_stack <- raster::addLayer(env_stack, dist_col)
env_stack <- raster::addLayer(env_stack, raster::stack(rsce_1$rasters), raster::stack(rsce_2$rasters))

env_df <- as.data.frame(env_stack, xy = T) |> 
  rename(Resource1 = sim1, Resource2 = sim2,
         Suitability_CPF = suitability.1, Suitability_FR = suitability.2) |> 
  mutate(area = 0.5*0.5) 


#### ------------------------------- ####
#### ---- Open and prepare data ---- ####
#### ------------------------------- ####

## Open data -------------------------------------------------------------------
cpf_plane <- read.table(paste0(workdir, "/Virtual_survey_CPF_plane.csv"), h = T, dec = ".", sep = ";")
cpf_plane_obs <- read.table(paste0(workdir, "/Virtual_survey_CPF_plane_obstable.csv"), h = T, dec = ".", sep = ";")
cpf_boat <- read.table(paste0(workdir, "/Virtual_survey_CPF_boat.csv"), h = T, dec = ".", sep = ";")
cpf_boat_obs <- read.table(paste0(workdir, "/Virtual_survey_CPF_boat_obstable.csv"), h = T, dec = ".", sep = ";")

fr_plane <- read.table(paste0(workdir, "/Virtual_survey_FR_plane.csv"), h = T, dec = ".", sep = ";")
fr_plane_obs <- read.table(paste0(workdir, "/Virtual_survey_FR_plane_obstable.csv"), h = T, dec = ".", sep = ";")
fr_boat <- read.table(paste0(workdir, "/Virtual_survey_FR_boat.csv"), h = T, dec = ".", sep = ";")
fr_boat_obs <- read.table(paste0(workdir, "/Virtual_survey_FR_boat_obstable.csv"), h = T, dec = ".", sep = ";")


## Extract environmental conditions ---------------------------------------
cpf_plane[, c("Resource1", "Resource2")] <- terra::extract(env$rasters, cpf_plane[, c("centre_x", "centre_y")])[,2:3]
cpf_plane$dist_col <- terra::extract(rast(dist_col), cpf_plane[, c("centre_x", "centre_y")])[,2]
cpf_plane$Suitability <- terra::extract(rsce_1$rasters, cpf_plane[, c("centre_x", "centre_y")])[,2]

cpf_boat[, c("Resource1", "Resource2")] <- terra::extract(env$rasters, cpf_boat[, c("centre_x", "centre_y")])[,2:3]
cpf_boat$dist_col <- terra::extract(rast(dist_col), cpf_boat[, c("centre_x", "centre_y")])[,2]
cpf_boat$Suitability <- terra::extract(rsce_1$rasters, cpf_boat[, c("centre_x", "centre_y")])[,2]


fr_plane[, c("Resource1", "Resource2")] <- terra::extract(env$rasters, fr_plane[, c("centre_x", "centre_y")])[,2:3]
fr_plane$dist_col <- terra::extract(rast(dist_col), fr_plane[, c("centre_x", "centre_y")])[,2]
fr_plane$Suitability <- terra::extract(rsce_2$rasters, fr_plane[, c("centre_x", "centre_y")])[,2]

fr_boat[, c("Resource1", "Resource2")] <- terra::extract(env$rasters, fr_boat[, c("centre_x", "centre_y")])[,2:3]
fr_boat$dist_col <- terra::extract(rast(dist_col), fr_boat[, c("centre_x", "centre_y")])[,2]
fr_boat$Suitability <- terra::extract(rsce_2$rasters, fr_boat[, c("centre_x", "centre_y")])[,2]


## True maps -------------------------------------------------------------------
CPF_true <- terra::rast(paste0(workdir, "/Simulated_tracks_CPF_KDE_totpop.grd"))
CPF_true_contour <- st_read(paste0(workdir, "/Simulated_tracks_CPF_KDE_totpop_poly.shp"))

FR_true <- terra::rast(paste0(workdir, "/Simulated_tracks_FR_KDE_totpop.grd"))
FR_true_contour <- st_read(paste0(workdir, "/Simulated_tracks_FR_KDE_totpop_poly.shp"))




#### ----------------------------- ####
#### ----- Habitat modelling ----- ####
#### ----------------------------- ####
source("contrib_function_202303.R")

### 0 - Set the function -------------------------------------------------------
gam_fitting <- function(effort, obs, env_data, outdir, outname, trunc_dist, cell_area = 0.5*0.5, type = "CPF"){
  ## 0 - Estimate ESW 
  region.table <- data.frame(Region.Label = 1, Area = 90*90)
  sample.table <- effort %>% mutate(Sample.Label = seg_id, 
                                    Region.Label = 1, Effort = SegLengthKm)
  obs.table <- obs %>% mutate(size = 1, Region.Label = 1, Sample.Label = seg_id, 
                              distance = dist_seg, object = 1:nrow(obs))
  
  ds.fit.dist <- ds(data = obs.table %>% subset(detected == 1), 
                    obs_table = obs.table %>% subset(detected == 1), 
                    region_table = region.table, sample_table = sample.table,
                    key = "hn",
                    adjustment = NULL,
                    truncation = trunc_dist,
                    quiet = TRUE
  )
  effort$esw <- unlist( predict(ds.fit.dist, newdata = effort, esw = T) )
  
  ds_results <- data.frame(
    esw = mean(effort$esw, na.rm = T),
    Nhat = summary(ds.fit.dist)$dht$individuals$N$Estimate,
    Nhat_low = summary(ds.fit.dist)$dht$individuals$N$lcl,
    Nhat_up = summary(ds.fit.dist)$dht$individuals$N$ucl
  )
  
  ## 1 - Fit the model 
  if(type == "CPF"){
    mod <- gam(N_ind_detected ~ s(Resource1, k = 4) + s(Resource2, k = 4) + s(dist_col, k = 4), 
               data = effort, offset = log(effort$SegLengthKm *2*effort$esw), family = tw())
  }
  if(type == "FR"){
    mod <- gam(N_ind_detected ~ s(Resource1, k = 4) + s(Resource2, k = 4), 
               data = effort, offset = log(effort$SegLengthKm *2*effort$esw), family = tw())
  }
  
  
  gam_results <- contrib_func(model = mod, data = effort) |> 
    as.data.frame() |> 
    mutate(Deviance = summary(mod)$dev.expl)
  
  gam_curves <- predict(mod, newdata = effort, type = "terms")
  gam_curves <- cbind(effort, gam_curves)
  
  ## 2 - Fit the prediction
  pred <- predict(object = mod, newdata = env_data, type = "response", se.fit = T)
  
  pred_df <- cbind(env_df, pred$fit, pred$se.fit) |> 
    rename(fit = `pred$fit`, se.fit = `pred$se.fit`) |> 
    mutate(fit = fit * cell_area, se.fit = se.fit * cell_area)
  
  fitrast <- raster::rasterFromXYZ(pred_df[, c("x", "y", "fit")])
  sefitrast <- raster::rasterFromXYZ(pred_df[, c("x", "y", "se.fit")])
  
  writeRaster(x = fitrast, filename = paste0(outdir, "/", outname, "_abundance_map.grd"), overwrite = T)
  writeRaster(x = fitrast - sefitrast, filename = paste0(outdir, "/", outname, "_abundance_map_CIlow.grd"), overwrite = T)
  writeRaster(x = fitrast + sefitrast, filename = paste0(outdir, "/", outname, "_abundance_map_CIhigh.grd"), overwrite = T)
  
  pred_stack <- raster::stack(fitrast, fitrast - sefitrast, fitrast + sefitrast)
  names(pred_stack) <- c("Abundance", "CIlow", "CIhigh")
  
  ds_results$Pred_mean <- raster::cellStats(pred_stack[[1]], "sum")
  ds_results$Pred_low <- raster::cellStats(pred_stack[[2]], "sum") 
  ds_results$Pred_high <- raster::cellStats(pred_stack[[3]], "sum") 
  
  
  ## 3 - Compare to true maps
  df_list <- list(as.data.frame(fitrast, xy = T), 
                  if(type == "CPF") { as.data.frame(rsce_1$rasters, xy = T) } else {as.data.frame(rsce_2$rasters, xy = T)}, 
                  if(type == "CPF") { as.data.frame(CPF_true, xy = T) } else {as.data.frame(FR_true, xy = T)})
  binded_tables <- df_list %>% reduce(full_join, by =  c("x", "y"))
  
  corr_kde <- Hmisc::rcorr(as.matrix(binded_tables[, c("fit", "mean")]), type = "pearson")
  corr_suit <- Hmisc::rcorr(as.matrix(binded_tables[, c("fit", "suitability")]), type = "pearson")
  
  gam_results$Pearson_kde <- as.vector(corr_kde$r)[2]
  gam_results$Pearson_suitability <- as.vector(corr_suit$r)[2]
  
  ## 4 - Export
  return(list(DistanceSampling = ds_results, 
              ModelStats = gam_results, 
              Model = mod,
              Curves = gam_curves, 
              Prediction = pred_stack,
              Prediction_table = pred_df))
}


### 1 - Apply to survey data ---------------------------------------------------
cpf_plane_mod <- gam_fitting(effort = cpf_plane, obs = cpf_plane_obs,
                             env_data = env_df, outdir = workdir, 
                             outname = "CPF_plane",
                             trunc_dist = max(cpf_plane_obs |> subset(detected==1) |> pull(dist_seg)),
                             type = "CPF")
cpf_boat_mod <- gam_fitting(effort = cpf_boat, obs = cpf_boat_obs,
                            env_data = env_df, outdir = workdir, 
                            outname = "CPF_boat",
                            trunc_dist = max(cpf_boat_obs |> subset(detected==1) |> pull(dist_seg)),
                            type = "CPF")

fr_plane_mod <- gam_fitting(effort = fr_plane, obs = fr_plane_obs,
                             env_data = env_df, outdir = workdir, 
                             outname = "FR_plane",
                             trunc_dist = max(fr_plane_obs |> subset(detected==1) |> pull(dist_seg)),
                            type = "FR")
fr_boat_mod <- gam_fitting(effort = fr_boat, obs = fr_boat_obs,
                            env_data = env_df, outdir = workdir, 
                            outname = "FR_boat",
                            trunc_dist = max(fr_boat_obs |> subset(detected==1) |> pull(dist_seg)),
                           type = "FR")




### 2 - Compare results --------------------------------------------------------
## abundance & esw results
ds_results <- rbind(
  cpf_plane_mod$DistanceSampling |> mutate(Species = "Central-place forager", Platform = "Plane"),
  cpf_boat_mod$DistanceSampling |> mutate(Species = "Central-place forager", Platform = "Boat"),
  
  fr_plane_mod$DistanceSampling |> mutate(Species = "Free ranger", Platform = "Plane"),
  fr_boat_mod$DistanceSampling |> mutate(Species = "Free ranger", Platform = "Boat")
  ) |> 
  rename(Nhat_mean = Nhat) |> 
  pivot_longer(cols = c("Nhat_mean", "Nhat_low", "Nhat_up", "Pred_mean", "Pred_low", "Pred_high"),
               names_to = "Origin", values_to = "Abundance") %>% 
  mutate(Source = str_split(Origin, pattern = "_") |> map(`[[`, 1) |> unlist(),
         Stats = str_split(Origin, pattern = "_") |> map(`[[`, 2) |> unlist()) |> 
  select(!c("Origin")) |> 
  mutate(Source = case_when(Source == "Nhat" ~ "Design-based estimate",
                            Source == "Pred" ~ "Model-based estimate"),
         Stats = case_when(Stats %in% c("high", "up") ~ "high",
                           TRUE ~ Stats)) |> 
  pivot_wider(names_from = Stats, values_from = "Abundance")

TrueAb <- data.frame(Species = c("Central-place forager", "Free ranger"),
                     Abundance = c(20000, 8000)) # update the FR value here

ggplot(ds_results) +
  geom_vline(data = TrueAb, aes(xintercept = Abundance), color = "black", linetype = "dashed") +
  geom_pointrange(aes(x = mean, xmin = low, xmax = high, 
                      y = Platform, color = Source),
                  position = position_dodge(0.5)) +
  scale_color_colorblind() +
  facet_wrap("Species") +
  labs(x = "Estimated population abundance") +
  theme_bw() + 
  theme(axis.title.y = element_blank(),
        strip.background = element_blank(),
        strip.text = element_text(face = "bold.italic", hjust = 0),
        legend.position = "bottom",
        legend.title = element_blank()) +
  coord_cartesian(xlim = c(0,55000))
ggsave("Abundance_estimations.png", dpi = 450, width = 8, height = 5)



## model performances
cpf_plane_mod$ModelStats
cpf_boat_mod$ModelStats

fr_plane_mod$ModelStats
fr_boat_mod$ModelStats

## model summaries
summary(cpf_plane_mod$Model)
summary(cpf_boat_mod$Model)

summary(fr_plane_mod$Model)
summary(fr_boat_mod$Model)


## functional responses
png("CPF_model_curves.png", height = 8, width = 12, res = 450, units = "cm")
par(mfrow = c(2,3), mar = c(2,2,1.1,0.5), oma = c(0.1,1.5,0.8,0.5), mgp = c(2.5,0.5,0), tck = -0.04, cex = 0.7)
plot(cpf_plane_mod$Model, shade = T, shade.col = "#147608ff", scale = 0, select = 1, xlim = c(0,1))
mtext("log(ind)", side = 2, line = 2, cex = 0.8)
mtext("Resource 1", side = 3, line = 0.3, cex = 0.8, adj = 0)
plot(cpf_plane_mod$Model, shade = T, shade.col = "#147608ff", scale = 0, select = 2, xlim = c(0,1))
mtext("Resource 2", side = 3, line = 0.3, cex = 0.8, adj = 0)
plot(cpf_plane_mod$Model, shade = T, shade.col = "#147608ff", scale = 0, select = 3, xlim = c(0,65))
mtext("Distance to colony", side = 3, line = 0.3, cex = 0.8, adj = 0)

plot(cpf_boat_mod$Model, shade = T, shade.col = "#b9e7b3ff", scale = 0, select = 1, xlim = c(0,1))
mtext("log(ind)", side = 2, line = 2, cex = 0.8)
plot(cpf_boat_mod$Model, shade = T, shade.col = "#b9e7b3ff", scale = 0, select = 2, xlim = c(0,1))
plot(cpf_boat_mod$Model, shade = T, shade.col = "#b9e7b3ff", scale = 0, select = 3, xlim = c(0,65))
dev.off()

png("FR_model_curves.png", height = 8, width = 12, res = 450, units = "cm")
par(mfrow = c(2,3), mar = c(2,2,1.1,0.5), oma = c(0.1,1.5,0.8,0.5), mgp = c(2.5,0.5,0), tck = -0.04, cex = 0.7)
plot(fr_plane_mod$Model, shade = T, shade.col = "#8c3eb2ff", scale = 0, select = 1, xlim = c(0,1))
mtext("log(ind)", side = 2, line = 2, cex = 0.8)
mtext("Resource 1", side = 3, line = 0.3, cex = 0.8, adj = 0)
plot(fr_plane_mod$Model, shade = T, shade.col = "#8c3eb2ff", scale = 0, select = 2, xlim = c(0,1))
mtext("Resource 2", side = 3, line = 0.3, cex = 0.8, adj = 0)
plot.new()

plot(fr_boat_mod$Model, shade = T, shade.col = "#dfb1f6ff", scale = 0, select = 1, xlim = c(0,1))
mtext("log(ind)", side = 2, line = 2, cex = 0.8)
plot(fr_boat_mod$Model, shade = T, shade.col = "#dfb1f6ff", scale = 0, select = 2, xlim = c(0,1))
plot.new()
dev.off()





save.image(file = "Survey_analysis.RData")

