## Spatio-temporal models for CTImean, CTIdiv, STBImean, STBIdiv and population trends
## Birds, butterflies, moths, phytoplankton and plants

#set path
try(setwd(''))

suppressPackageStartupMessages({
  library(terra)
  library(tidyverse)
  library(vegan)
  library(reshape)
  library(MASS)
  library(sf)
  library(rnaturalearth)
  library(rnaturalearthdata)
  library(paletteer)
  library(ggrepel)
  library(gridExtra)
  library(RColorBrewer)
  library(INLA)
  library(pracma)
  library(viridis)
  library(ggsci)
  library(wesanderson)
  library(betapart)
  library(parallel)
  library(stars)
  library(adegenet)
  library(spdep)
  library(MoMAColors)
  library(tidyterra)
})

#parallelzing options
options(mc.cores = parallel::detectCores())

#model structures
models = lapply(1:5, function(x) list())
models_temperature = vector(mode = 'list', length = 5)
names(models_temperature) = c('moth', 'butterfly', 'bird', 'plant', 'plankton')

#load data
load('data/Data_repository.RData')

#neighbourhood matrix
nb2knn_matrix = list()

#coordinate system EPSG: 3067
#get a background map
world = ne_countries(scale = "medium", returnclass = "sf")
world_plane = st_transform(world, 3067)
#bioclimatic zones
bioclim = st_read('data/bioclimatic_zones/BioclimaticZonesFinland.shp')

#response variables
response_all = c('CTI', 'CTIdiv', 'CTBI', 'CTBIdiv', 'Warm_sp_abundance', 'Cold_sp_abundance')

#species index
i_sp = 0

#loop over species
for (species_group in c('moth', 'butterfly', 'bird', 'plant', 'phytoplankton')) {
  
  #set index for plotting responses
  i_sp = i_sp+1
  print(i_sp)
  
  #add intercept
  data_temp = data_all[[i_sp]]
  
  #loop over response variables
  for (j in 1:length(response_all)) {
    print(j)
    #set response
    ind_response = which(colnames(data_temp)==response_all[j])

    #define a connectivity matrix for discrete spatial latent effects
    id_all = data_temp$SiteID
    id_unique = id_all[!duplicated(id_all)]
    xy_coord = data.frame(X=data_temp$Longitude, Y=data_temp$Latitude)
    xy_unique[[i_sp]] = xy_coord[!duplicated(id_all),]
    # nb2knn = chooseCN(jitter(as.matrix(xy_unique[[i_sp]]), 1), type = 6, result.type = "nb", k = 6, edit.nb = F)
    # nb2knn_matrix[[i_sp]] = spdep::nb2mat(nb2knn)
    
    #save the adjacency matrix in INLA format
    # nb2INLA(paste0('data/Lattice_k_', species_group, '.graph'), nb2knn)
    Lattice.adj = paste0('data/Lattice_k_', species_group, '.graph')
    
    #linear combinations
    #linear combinations for average effects over all sites
    values = sort(unique(data_temp$SiteID_tr_2))
    m = length(values)
    av_time = rep(1/m, m)
    lc = inla.make.lincomb(SiteID_tr_2 = av_time)
    names(lc) = paste0("Av_time")
    
    #linear combinations for average effects over bc-zones
    bc_unique = sort(unique(data_temp$bc_zone))
    site_df = data.frame(unique(cbind(data_temp$SiteID_tr_2, data_temp$bc_zone)))
    site_df = site_df[sort(site_df[,1], index.return = T)[[2]],]
    for (i in bc_unique) {
      av_time = rep(0,m)
      av_time[site_df$X2==i] = 1/sum(site_df$X2==i)
      lc.i = inla.make.lincomb(SiteID_tr_2 = av_time)
      names(lc.i) = paste0("Av_time_bc_", i)
      lc = c(lc, lc.i)
    }
    
    #linear combinations for temporal change
    values = sort(unique(data_temp$SiteID_tr_2))
    m = length(values)
    av_time = rep(1/m, m)
    for (k in unique(data_temp$Year_continuous)) {
      lc.i = inla.make.lincomb(SiteID_tr_2 = av_time*k)
      names(lc.i) = paste0("Av_change_", which(unique(data_temp$Year_continuous)==k))
      lc = c(lc, lc.i)
    }
    
    #linear combinations for average effects over bc-zones
    bc_unique = sort(unique(data_temp$bc_zone))
    site_df = data.frame(unique(cbind(data_temp$SiteID_tr_2, data_temp$bc_zone)))
    site_df = site_df[sort(site_df[,1], index.return = T)[[2]],]
    for (i in bc_unique) {
      for (k in unique(data_temp$Year_continuous)) {
        av_time = rep(0,m)
        av_time[site_df$X2==i] = 1/sum(site_df$X2==i)
        lc.i = inla.make.lincomb(SiteID_tr_2 = av_time*k)
        names(lc.i) = paste0("Av_change_", which(unique(data_temp$Year_continuous)==k),"_bc_", i)
        lc = c(lc, lc.i)
      }
    }
    
    if (i_sp > 3 & j > 4) {
      
      data_stack = inla.stack(data = list(y = data_temp[,ind_response]+1), A = list(1), 
                              effects = list(data.frame(data_temp)),
                              tag = 'base')
    } else {
    
      data_stack = inla.stack(data = list(y = data_temp[,ind_response]), A = list(1), 
                              effects = list(data.frame(data_temp)),
                              tag = 'base')
    }
  
    #set default priors for the hyperparameters
    hyper.besag = list(prec = list("pc.prec", param = c(1, .01)))
    hyper.bym2 = list(theta1 = list("pc.prec", param = c(1, .01)), theta2 = list("pc.cor1", param = c(.5, .5)))
    
    if (j <= 4) {
      
      formulae = y ~ f(SiteID_tr_1, model = 'bym2', graph = Lattice.adj, constr = T, hyper = hyper.bym2) +
        f(SiteID_tr_2, Year, model = 'bym2', graph = Lattice.adj, constr = F, hyper = hyper.bym2)
      
      if (species_group == 'butterfly') {
        formulae = update.formula(formulae, ~.+ effort + effort_rep)
      }
      
      models[[i_sp]][[j]] = inla(formulae,  data = inla.stack.data(data_stack), lincomb = lc,
                             family = 'normal',
                             control.predictor = list(compute = TRUE, A = inla.stack.A(data_stack)),
                             control.family = list(link = 'identity', hyper=list(prec=list(prior='pc.prec', param=c(1,.1)))),
                             control.inla = list(step.len = 0.1, tolerance = 1e-4),
                             control.compute = list(return.marginals.predictor = T, cpo = T, waic = T, dic = T))
      
    } else if (j > 4 & i_sp <= 3) {
      
      formulae = y ~ f(SiteID_tr_1, model = 'bym2', graph = Lattice.adj, constr = T, hyper = hyper.bym2) +
        f(SiteID_tr_2, Year, model = 'besag', graph = Lattice.adj, constr = F, hyper = hyper.besag)
      
      if (species_group == 'butterfly') {
        formulae = update.formula(formulae, ~.+ effort + effort_rep)
      }
      
      models[[i_sp]][[j]] = inla(formulae,  data = inla.stack.data(data_stack), lincomb = lc,
                                 family = 'poisson',
                                 control.predictor = list(compute = TRUE, A = inla.stack.A(data_stack)), 
                                 control.family = list(link = 'log'),
                                 control.inla = list(step.len = 0.1, tolerance = 1e-4),
                                 control.compute = list(return.marginals.predictor = T, cpo = T, waic = T, dic = T))
      
    } else if (j > 4 & i_sp > 3) {
      
      formulae = y ~ f(SiteID_tr_1, model = 'bym2', graph = Lattice.adj, constr = T, hyper = hyper.bym2) +
        f(SiteID_tr_2, Year, model = 'besag', graph = Lattice.adj, constr = F, hyper = hyper.besag)
      
      models[[i_sp]][[j]] = inla(formulae,  data = inla.stack.data(data_stack), lincomb = lc,
                                 family = 'lognormal',
                                 control.predictor = list(compute = TRUE, A = inla.stack.A(data_stack)), 
                                 control.inla = list(step.len = 0.1, tolerance = 1e-4),
                                 control.compute = list(return.marginals.predictor = T, cpo = T, waic = T, dic = T),
                                 verbose = T)
      
    }
    
    ##Model for temperature change
    if (j == 1) {
      #temperature
      data_stack = inla.stack(data = list(y = data_temp$temp), A = list(1), 
                              effects = list(data.frame(data_temp)),
                              tag = 'base')
      
      #set default priors for the hyperparameters
      formulae = y ~ f(SiteID_tr_1, model = 'bym2', graph = Lattice.adj, constr = T, hyper = hyper.bym2) +
        f(SiteID_tr_2, Year, model = 'bym2', graph = Lattice.adj, constr = F, hyper = hyper.bym2)
      
      models_temperature[[i_sp]] = inla(formulae,  data = inla.stack.data(data_stack), lincomb = lc,
                                 family = 'normal',
                                 control.predictor = list(compute = TRUE, A = inla.stack.A(data_stack)),
                                 control.family = list(link = 'identity', hyper=list(prec=list(prior='pc.prec', param=c(1,.1)))),
                                 control.compute = list(return.marginals.predictor = T, cpo = T, waic = T, dic = T))
    }
  }
}

#save model-structures and data sets
for (i in 1:length(models)) {
  names(models[[i]]) = response_all
}

save(models, models_temperature, data_all, xy_unique, file = 'data/Model_fits.RData', compress = T)
