library(dplyr)
library(raster)
library(geoR)
library(gstat)
library(class)
library(RStoolbox)
library(ggplot2)
library(RColorBrewer)
library(ggpubr)
library(ggspatial)
library(ggsn)
library(viridis)
library(rgeos)
library(car)
library(tidyr)
library(glmmTMB)
library(DHARMa)
library(sjPlot)
library(spatstat)
library(GGally)
library(ape)
library(performance)
library(visreg)

data = read.csv('aspen data site-level processed 30 Mar 2020.csv')

# count up sites of different types
table(sapply(as.character(data$Site_Code),nchar))

# find pairwise dist
nndists = data %>% filter(Point_Type!="Grid") %>% dplyr::select(X.UTM,Y.UTM) %>% nndist
mean(nndists)
sd(nndists)

# count up ploidy levels
table(data[data$Point_Type!="Grid",]$Ploidy_level)
table(data$Ploidy_level) 
length(unique(data$Genotype))# of clones
table(data$Genotype) %>% mean
table(data$Genotype) %>% sd
table(data$Genotype) %>% max

dr = data[data$Point_Type=="Random",]
# genotypes with more representation

bigclones = dr %>% 
  filter(Genotype %in% names(which(table(dr$Genotype) > 3))) %>% 
  dplyr::select(Genotype, Ploidy_level) %>% 
  unique
bigclones

smallclones = dr %>% 
  filter(Genotype %in% names(which(table(dr$Genotype) == 1))) %>% 
  dplyr::select(Genotype, Ploidy_level) %>% 
  unique
smallclones
table(smallclones$Ploidy_level)



# MAP THE GRID
# this function has a hack in it to make the triploid-only sites plot correctly. 
do_map = function(watershed)
{
  data_ss = data %>% filter(Watershed==watershed & Point_Type=="Grid") %>% 
    mutate(Ploidy_level=as.character(Ploidy_level), Genotype=as.character(Genotype)) %>%
    mutate(Genotype=ifelse(is.na(Genotype),paste("?",1:length(is.na(Genotype))),Genotype)) %>%
    mutate(Ploidy_level=ifelse(is.na(Ploidy_level),"?",Ploidy_level)) %>%
    mutate(Genotype=factor(paste(Genotype,Ploidy_level))) 
  
  
  
  data_ss_clone = SpatialPointsDataFrame(coords=data_ss[,c("X.UTM","Y.UTM")],
                                         data=data_ss[,"Genotype",drop=FALSE],
                                         proj4string = CRS("+proj=utm +zone=13 +ellps=GRS80 +towgs84=0,0,0,0,0,0,0 +units=m +no_defs"))

  data_ss_ploidy = SpatialPointsDataFrame(coords=data_ss[,c("X.UTM","Y.UTM")],
                                          data=data_ss[,c("Ploidy_level","X.UTM","Y.UTM"),drop=FALSE],
                                          proj4string = CRS("+proj=utm +zone=13 +ellps=GRS80 +towgs84=0,0,0,0,0,0,0 +units=m +no_defs"))

  r_clone = extend(raster(res=1,ext=extent(data_ss_clone),crs=crs(data_ss_clone)),25)
  r_clone[] = NA
  
  data_ss_ploidy_for_buffer = data_ss_ploidy
  data_ss_ploidy_for_buffer@data$Ploidy_level = rep(1, length(data_ss_ploidy_for_buffer@data$Ploidy_level))
  r_buffer = buffer(rasterize(data_ss_ploidy_for_buffer,r_clone)[[1]],25)
  
  # code for knn factor classification (nearest neighbor prediction)
  knn_map <- function(raster, xy, val, k=1)
  {
    val <- factor(val)
    val_char <- as.numeric(val)
    val_levels <- as.character(levels(val))
    
    predpoints <- xyFromCell(raster, 1:prod(dim(raster)))
    
    stopifnot(nrow(xy) == length(val))
    
    vals_predicted <- knn(train=xy,test=predpoints, cl=val_char,k=k)
    
    # fill in values
    raster[] <- factor(vals_predicted,levels=val_levels)
    raster@data@values <- as.numeric(vals_predicted)
    
    return(list(raster=raster,levels=val_levels))
  }
  
  genotypes <- knn_map(r_buffer, 
                       xy=data_ss_clone@coords,
                       val=data_ss_clone@data$Genotype)
  genotype_clipped = mask(genotypes[[1]], r_buffer)
  
  ploidy <- knn_map(r_buffer, 
                       xy=data_ss_ploidy@coords,
                       val=data_ss_ploidy@data$Ploidy_level)
  ploidy_clipped = mask(ploidy[[1]], r_buffer)
  
  #data_ss$Clone = as.character(data_ss$Clone)
  #data_ss$Ploidy_level = as.character(data_ss$Ploidy_level)
  #data_ss$Clone[is.na(data_ss$Clone)] = ""
  #data_ss$Ploidy_level[is.na(data_ss$Ploidy_level)] = ""
  #cv = unique(data_ss[,c("Clone","Ploidy_level")])
  
  #colors = hsv(s=ifelse(cv$Ploidy_level=="triploid",1,ifelse(cv$Ploidy_level=="diploid",0.6,0)),
  #              h=seq(0,0.1,length.out=nrow(cv)))
               
               #viridis(length(unique(data_ss_clone@data$Clone)))
  set.seed(1)
  #colors = #sample(c(brewer.pal(9,"Set1"),brewer.pal(8,"Set2")))
  colors = colorRampPalette(brewer.pal(9,"Spectral"))(length(unique(data_ss_clone@data$Genotype)))
  
  print(watershed)
  print(unique(data_ss_ploidy@data$Ploidy_level))
  
  colors_ploidy_map=c("black","red","blue"); names(colors_ploidy_map) = c("?","Triploid","Diploid")
  colors_ploidy = colors_ploidy_map[unique(data_ss_ploidy@data$Ploidy_level)]
  
  g = ggR(genotype_clipped,geom_raster=TRUE,alpha=0.5) + 
    blank() +
    #ggR(ploidy_clipped,geom_raster=FALSE,alpha=0.5) +
    #scale_fill_viridis(discrete=TRUE)
    scale_fill_manual(values=colors,drop=FALSE,guide='none') + 
    #theme_bw() + 
    #geom_text(data=data_ss_ploidy@data,aes(x=X.UTM,y=Y.UTM,label=Ploidy_level,color=Ploidy_level),size=3,alpha=1,fontface='bold') +
    geom_point(data=data_ss_ploidy@data,aes(x=X.UTM,y=Y.UTM,color=Ploidy_level),size=3,alpha=1) +
    scale_color_manual(values=colors_ploidy,drop=FALSE) +
    #theme(legend.position='bottom') +
    ggtitle(watershed) +
    annotation_scale(location = "tl", height = unit(0.1, "cm")) +
    annotation_north_arrow(location = "tr", 
                           style = north_arrow_minimal, 
                           height = unit(0.5, "cm"),
                           width = unit(0.5, "cm")) +
    labs(color='Cytotype') 
  
  return(list(g=g, gt=genotype_clipped, pl=ploidy_clipped))
  #ggsave(g, file=sprintf('clonemap_%s.pdf',watershed))
  #writeRaster(genotype_clipped, file=sprintf('clonemap_%s.grd',watershed),overwrite=TRUE)
  #return(list(g=g, r=genotype_clipped))
}

g_all_clones_grid = lapply(levels(data$Watershed), do_map)
g_clones_grid = ggarrange(plotlist=lapply(g_all_clones_grid, function(x) { x$g }), align='hv',nrow=2,ncol=2,common.legend = TRUE,legend='bottom',labels=c('(a)','(b)','(c)','(d)'))
ggsave(g_clones_grid,file='FIG3_g_clones_grid.pdf',width=7,height=7)



# plot overall site
hill_map = raster('hill.tif')
alt_map = raster('altitudes.tif')
data_plot = data
data_plot$Ploidy_level = as.character(data_plot$Ploidy_level)
data_plot$Ploidy_level[is.na(data_plot$Ploidy_level)] <- "?"
data_plot$Ploidy_level <- factor(data_plot$Ploidy_level,levels=c("?","Triploid","Diploid"),labels=c("?","Triploid","Diploid"))

ws_boundaries = shapefile("NEON-overflight-4wshed.shp")
ws_boundaries_utm = spTransform(ws_boundaries, CRS=crs(hill_map))
ws_boundaries_utm_cropped = crop(ws_boundaries_utm, hill_map)
ws_boundaries_utm_df <- fortify(ws_boundaries_utm)
ws_boundaries_utm_cropped_df <- fortify(ws_boundaries_utm_cropped)
ws_centers = ws_boundaries_utm_cropped_df %>% group_by(group) %>% summarize(X.UTM=mean(long), Y.UTM=mean(lat))
ws_centers$group = as.character(levels(data$Watershed))[c(4,1,3,2)]


g_map <- ggR(hill_map) + 
  ggR(alt_map,geom_raster=TRUE, ggLayer=TRUE,alpha=0.3) + scale_fill_gradientn(colours = terrain.colors(100), name = "Elevation (m)",limits=c(2400,4000)) +
  blank() +
  xlab("") + ylab("") + annotation_scale(location = "bl", height = unit(0.1, "cm")) +
  annotation_north_arrow(location = "br", 
                         style = north_arrow_minimal, 
                         height = unit(0.5, "cm")) +
  #geom_text(data=data_plot,aes(x=X.UTM,y=Y.UTM,label=Ploidy_level,color=Ploidy_level),size=3,alpha=1,fontface='bold') +
  geom_point(data=data_plot,aes(x=X.UTM,y=Y.UTM,col=Ploidy_level),alpha=0.5) +
  scale_color_manual(values=c("black","red","blue"),drop=FALSE) +
  geom_path(data = ws_boundaries_utm_df, 
          aes(x = long, y = lat, group = group),
          color = 'purple', size = 0.3) +
  geom_text(data=ws_centers, aes(x=X.UTM,y=Y.UTM,label=group)) +
  xlim(extent(alt_map)@xmin, extent(alt_map)@xmax) + 
  ylim(extent(alt_map)@ymin, extent(alt_map)@ymax) +
  labs(color='Cytotype') 

ggsave(g_map,file="FIG2_g_map.png",width=8,height=6)
ggsave(g_map,file="FIG2_g_map.pdf",width=8,height=6)








simplify_geology <- function(unit)
{
  if (unit %in% c("Qal","Qg","Qu","Qlu","Qdf"))
  {
    #return("Quaternary alluvium or outwash gravel or deltas or undifferentiated")
    #return("Quaternary sorted")
    return("Quaternary deposit")
  }
  else if (unit %in% c("Ql","Qlf","Qd","Qdu"))
  {
    #return("Quaternary landslide or debris flow deposit")
    #return("Quaternary unsorted")
    return("Quaternary deposit")
  }
  else if (unit %in% c("Qm","Qmy"))
  {
    #return("Quaternary glacial deposit")
    #return("Quaternary unsorted")
    return("Quaternary deposit")
  }
  else if (unit %in% c("Qt","Qr"))
  {
    #return("Quaternary talus or rock glacier")
    return("Quaternary talus / rock glacier")
  }
  else if (unit %in% c("Tt","Tg","Tf","Tp","Tcd","Tbx","Td","qm"))
  {
    return("Igneous intrusive")
  }
  else if (unit %in% c("Kmu","Kml","Km"))
  {
    return("Shale or limestone")  
  }
  else if (unit %in% c("Kmf"))
  {
    #return("Fort Hays limestone") 
    return("Shale or limestone") # because we have no diploids on limestone - messes with mixed model
  }
  else if (unit %in% c("Kd","Je","Kmv","Kmvo"))
  {
    #return("Sandstone") 
    return("Sandstone or conglomerate or siltstone")
  }
  else if (unit %in% c("Jm","PPm","PPM")) #PPM is probably a typo
  {
    #return("Conglomerate or siltstone")
    return("Sandstone or conglomerate or siltstone")
  }
  else
  {
    return("")
  }
}

# simplify the geology
data$Rock_Unit = factor(sapply(data$Geologic.Unit,simplify_geology))
# relabel canopy openness
data$Light_Index = 1 - (as.numeric(data$Canopy_openness) - 1)/4 #data$Canopy_openness # 0 to 1 index
data$Unhealthy_Site = factor(data$Unhealthy_Site, levels=0:1,labels=c("Healthy","Unhealthy"),ordered=TRUE)
data$Cow = factor(data$Cow_Damage, levels=0:1,labels=c("Absent","Present"),ordered=FALSE)
data$Soil_Type = factor(data$Soil.type, levels=c("Soil","Scree","Talus"),ordered=FALSE)
data$Point_Type = as.character(data$Point_Type)
data$Point_Type[sapply(as.character(data$Site_Code), nchar)>5] = "Opportunistic"
data$Cos_Aspect = data$Cos.aspect
data$DBH_Mean = data$DBH.mean
data$Ploidy_Level = factor(data$Ploidy_level,levels=c("Diploid","Triploid"),labels=c("Diploid","Triploid"))

names(data) = gsub("\\.","",gsub("_","",names(data)))

names(data)[which(names(data)=="SoilType")] = "Regolith"


predictors_categorical = c("Cow","Regolith","RockUnit")
predictors_continuous = c("Elevation","SummerInsolation","Slope","DBHMean","LightIndex")
resp_vars = c("Countmedium","Countsmall","Countadultdead", "Countadultdamaged", "UnhealthySite", "PloidyLevel","Genotype")

# write in the NA ploidies
data_no_na = data %>% filter(!is.na(PloidyLevel))

# scale continuous predictors
data_no_na_for_models = data_no_na
data_no_na_for_models[,predictors_continuous] = scale(data_no_na_for_models[,predictors_continuous])

# select only non-NA values of relevant columns
data_no_na_for_models = data_no_na_for_models %>% 
  dplyr::select(c(resp_vars, predictors_categorical, predictors_continuous,"XUTM","YUTM")) %>%
  na.omit %>%
  mutate(Countmortalityintegrative = Countadultdamaged + Countadultdead)
summary(data_no_na_for_models)



# check for predictor correlation
g_pairs = ggpairs(data_no_na_for_models %>% dplyr::select(c(predictors_categorical, predictors_continuous))) +
  theme_bw()
ggsave(g_pairs, width=12,height=12,file='g_pairs.pdf')




# pick photos
data %>% filter(Unhealthy_Site=="Unhealthy" & Ploidy_Level=="Triploid" & Count.small < 5) %>% head(10) #RGBOB01
data %>% filter(Unhealthy_Site=="Healthy" & Ploidy_Level=="Diploid" & Count.small > 1) %>% head(10)

# show distribution of ploidy levels
g_ploidy_count = ggplot(data_no_na,aes(x=PloidyLevel,fill=PloidyLevel)) + 
geom_bar(stat='count') + 
  xlab("Cytotype") + ylab("Count") + 
  theme_bw() + facet_wrap(~PointType) +
  scale_fill_manual(values=c("blue","red")) +
  labs(fill='Cytotype') 
ggsave(g_ploidy_count,file='g_ploidy_count.pdf',width=6,height=5)




plot_categorical <- function(pred)
{
  require(stringr)
  
  data_no_na_this = data_no_na
  data_no_na_this[,pred] = factor(data_no_na_this[,pred])
  
  cols = rev(brewer.pal(n=9,name="Set3"))
  
  g = ggplot(data_no_na_this,aes_string(fill=str_wrap(pred,20),x="PloidyLevel")) + 
    geom_bar(stat="count",position='fill') + 
    theme_bw() +
    scale_fill_manual(values = cols) + 
    ylab("Fraction") + 
    xlab("Cytotype") + theme(legend.position='bottom', legend.direction="vertical",legend.justification = c(0,1))
  
  return(g)
}

g_categorical = lapply(c(predictors_categorical,"LightIndex"), plot_categorical)



plot_continuous <- function(pred)
{
  data_no_na_this = data_no_na
  
  data_no_na_this[,pred] = factor(data_no_na_this[,pred])
  
  g = ggplot(data_no_na,aes_string(x="PloidyLevel",y=pred,group="PloidyLevel")) + 
    geom_violin(alpha=0.5) +
    theme_bw() + theme(legend.position='none') + 
    xlab("Cytotype")
  
  return(g)
}

g_continuous = lapply(setdiff(predictors_continuous,"LightIndex"), plot_continuous)

g_r1 = ggarrange(plotlist=g_categorical,nrow=1,ncol=length(g_categorical),align='hv',labels=paste("(",letters[5:8],")",sep=""),font.label = list(size = 12))
g_r2 = ggarrange(plotlist=g_continuous,nrow=1,ncol=length(g_continuous),labels=paste("(",letters[1:4],")",sep=""),font.label = list(size = 11))

g_predictors = ggarrange(g_r2, g_r1, nrow=2,ncol=1,heights=c(1,2))

ggsave(g_predictors, file='FIG4_g_predictors.pdf',width=12,height=8)








# do correlation analysis of ploidy
xtabs(~Regolith+PloidyLevel,data=data) %>% prop.table(2)
xtabs(~Regolith+PloidyLevel,data=data) %>% chisq.test

xtabs(~RockUnit+PloidyLevel,data=data) %>% prop.table(2)
xtabs(~RockUnit+PloidyLevel,data=data) %>% chisq.test

xtabs(~Cow+PloidyLevel,data=data) %>% prop.table(2)
xtabs(~Cow+PloidyLevel,data=data) %>% chisq.test

xtabs(~LightIndex+PloidyLevel,data=data) %>% prop.table(2)
xtabs(~LightIndex+PloidyLevel,data=data) %>% chisq.test

t.test(Elevation~PloidyLevel,data=data)
t.test(Slope~PloidyLevel,data=data)
t.test(SummerInsolation~PloidyLevel,data=data)
t.test(DBHmean~PloidyLevel,data=data)








#### build models
do_demography_model <- function(response_variable, family, data, includeE, includeG, includeP, includePxE)
{
  pred_string = NULL # include intercept
  if (includeE==TRUE)
  {
    pred_string=c(pred_string,paste(c(predictors_categorical, predictors_continuous),collapse=" + "))
  }
  if (includeG==TRUE)
  {
    pred_string=c(pred_string,"(1|Genotype)")
  }
  if (includeP==TRUE)
  {
    pred_string=c(pred_string,"PloidyLevel")
  }
  if (includePxE==TRUE)
  {
    pred_string=c(pred_string,paste(paste("PloidyLevel",c(predictors_categorical, predictors_continuous),sep=":"),collapse=" + "))
  }
  pred_string = paste(pred_string, collapse=" + ")
  
  if (pred_string=="")
  {
    pred_string = "1" # intercept-only case
  }
  
  # don't calculate models without the main effects
  if (includePxE == TRUE & (includeP == FALSE | includeE == FALSE))
  {
    return(NULL)
  }
  else
  {
    formula_this = as.formula(sprintf("%s ~ %s", 
                                   response_variable, pred_string),
                              env = globalenv())
    print(formula_this)
    
    m_this = glmmTMB(formula=formula_this, 
                     family=family,
                     data=data)
    return(m_this)
  }
}

do_demography_model_all <- function(response_variable, family, data)
{
  params = expand.grid(includeE=c(T,F),
                       includeP=c(T,F),
                       includeG=c(T,F),
                       includePxE=c(T,F))
  
  models = vector(mode="list",length=nrow(params))
  for (i in 1:nrow(params))
  {
    cat('.')
    models[[i]] = do_demography_model(response_variable = response_variable,
                                      family = family,
                                      data = data,
                                      includeE=params$includeE[i],
                                      includeG=params$includeG[i],
                                      includeP=params$includeP[i],
                                      includePxE=params$includePxE[i]
    )
  }
  names(models) = paste(ifelse(params$includeE,"E",""),ifelse(params$includeG,"G",""),ifelse(params$includeP,"C",""),ifelse(params$includePxE,"CxE",""),sep="_")
  return(models)
}


make_aic_table = function(model_list)
{
  aic_vals = sapply(model_list, function(x) { 
    if(is.null(x)) { return(NA)} else { return(AIC(x)) } })
  
  #rmse = sapply(model_list, function(x) { 
  #  if(is.null(x)) { return(NA)} else { return(rmse(x,normalized=TRUE)) } })
  
  df_out = data.frame(aic_vals)#,rmse)
  names(df_out) = c("AIC")#,"RMSE")
  df_out$Model = names(model_list)
  
  df_out$DeltaAIC = df_out$AIC - min(df_out$AIC,na.rm=T)
  
  row.names(df_out) = NULL
  return(df_out)
}



models_recruitment = do_demography_model_all(response_variable = "Countsmall", family=nbinom1(link = 'log'), data=data_no_na_for_models)
models_mortality = do_demography_model_all(response_variable = "Countmortalityintegrative", family=nbinom1(link='log'), data=data_no_na_for_models)

table_recruitment = make_aic_table(models_recruitment)
table_mortality = make_aic_table(models_mortality)

# pick best models
m_recruitment = models_recruitment[[which.min(table_recruitment$DeltaAIC)]]
m_mortality = models_mortality[[which(table_mortality$Model=="E_G_C_")]]







plot_aic <- function(df)
{
  df$ModelPretty = gsub("_", ", ", gsub("_$","",gsub("^_","",gsub("__","_",gsub("__","_",df$Model)))))
  
  # drop non-computed models
  df = na.omit(df)
  
  g = ggplot(df, aes(x=ModelPretty,xend=ModelPretty,y=min(df$AIC,na.rm=T)-5,yend=AIC,col=(DeltaAIC<2))) +
    theme_bw() +
    geom_segment(size=6) +
    #geom_point(size=5,aes(x=ModelPretty,y=AIC)) +
    scale_color_manual(values=c("black","red"),guide=FALSE) +
    ylim(min(df$AIC,na.rm=T)-5,max(df$AIC,na.rm=T)+5) +
    xlab("Model terms") +
    ylab("AIC") + coord_flip()
  
  return(g)
}

g_aic_recruitment = plot_aic(table_recruitment) + ggtitle("Recruitment")
g_aic_mortality = plot_aic(table_mortality) + ggtitle("Mortality")

ggsave(ggarrange(g_aic_recruitment, g_aic_mortality,nrow=2,ncol=1,labels=c("(a))","(b)")),
       file='FIG5_g_aic.pdf',width=6,height=6)











classify_predictor <- function(string)
{
  if (length(grep("\\:",string)) > 0)
  {
    return("CxE")
  }
  else if (length(grep("PloidyLevel",string)) > 0)
  {
    return("C")
  }
  else
  {
    return("E")
  }
}




plot_model_demography = function(m, col_fixed, col_random,heights=c(2,1),title)
{
  df_fixed = data.frame(confint(m))
  names(df_fixed) = c("lo","hi","Estimate")
  df_fixed$Predictor = gsub("cond.","",row.names(df_fixed),fixed=TRUE)
  df_fixed$Type = sapply(df_fixed$Predictor, classify_predictor)
  row.names(df_fixed) = NULL
  df_fixed = df_fixed[,c("Predictor","lo","hi","Estimate","Type")]
  df_fixed = df_fixed %>% filter(!Predictor %in% c("(Intercept)","Genotype.Std.Dev.(Intercept)","sigma"))
  
  df_random = data.frame(ranef(m)$cond)
  names(df_random) = "Estimate"
  df_random$Predictor = gsub("cond.","",row.names(df_random),fixed=TRUE)
  df_random$Type = "G"
  df_random$lo = df_random$Estimate - sqrt(attr(ranef(m)$cond$Genotype,"condVar")[1,1,])
  df_random$hi = df_random$Estimate + sqrt(attr(ranef(m)$cond$Genotype,"condVar")[1,1,])
  row.names(df_random) = NULL
  df_random = df_random[,c("Predictor","lo","hi","Estimate","Type")]

  g_fixed = ggplot(df_fixed,aes(x=reorder(Predictor,Estimate),ymin=lo,ymax=hi,y=Estimate,col=Type)) +
    geom_hline(yintercept = 0) +
    geom_errorbar() + 
    geom_point() + 
    facet_grid(Type~.,scales='free_y',space='free') +
    theme_bw() +
    scale_color_manual(values=col_fixed) +
    coord_flip() +
    theme(legend.position='none') +
    ylab("") + xlab("Fixed effects") +
    ylim(-4,4) + 
    ggtitle(title)
  
  g_random = ggplot(df_random,aes(x=reorder(Predictor,Estimate),ymin=lo,ymax=hi,y=Estimate,col=Type)) +
    geom_hline(yintercept = 0) +
    geom_errorbar() + 
    geom_point() + 
    facet_grid(Type~.,scales='free_y',space='free') +
    theme_bw() +
    scale_color_manual(values=col_random) +
    coord_flip() +
    theme(legend.position='none') +
    theme(axis.text.y = element_text(size=0)) +
    ylab("Estimate") + xlab("Random effects") +
    ylim(-4,4)
  
  g_final = ggarrange(g_fixed, g_random,nrow=2,ncol=1,heights=heights,align='v')
  
  return(g_final)
}




g_recruitment_new = plot_model_demography(models_recruitment[[which.min(table_recruitment$DeltaAIC)]],
                                          col_fixed = brewer.pal(4,"Set1")[1:3],
                                          col_random = brewer.pal(4,"Set1")[4],
                                          title="Recruitment",
                                          heights=c(2,1)
)

ggsave(g_recruitment_new, file='FIG6_g_recruitment_new.pdf',width=8,height=7)


g_mortality_new = plot_model_demography(models_mortality[[which(table_mortality$Model=="E_G_C_")]],
                                        col_fixed = brewer.pal(4,"Set1")[1:2],
                                        col_random = brewer.pal(4,"Set1")[4],
                                        heights=c(1,1),
                                        title="Mortality"
)

ggsave(g_mortality_new, file='FIG7_g_mortality_new.pdf',width=8,height=6)

















### confirm models are OK

data_dists = as.matrix(dist(cbind(data_no_na_for_models$XUTM, data_no_na_for_models$YUTM)))
data_dists.inv = 1.0 / (0.0 + data_dists)
diag(data_dists.inv) = 0 # no self-influence


# spatial autocorrelation

# get distances for Moran's I tests - https://www.seas.upenn.edu/~ese502/lab-content/extra_materials/SPATIAL%20WEIGHT%20MATRICES.pdf
Moran.I(resid(m_recruitment),data_dists.inv)
# model fit
m_recruitment_simres <- simulateResiduals(models_recruitment[[which.min(table_recruitment$DeltaAIC)]])
plot(m_recruitment_simres,quantreg=FALSE)
#testUniformity(m_recruitment_simres)
#testZeroInflation(m_recruitment_simres)
#plotResiduals(m_recruitment_simres)

# spatial autocorrelation
Moran.I(resid(m_mortality),data_dists.inv)
# model fit
m_mortality_simres <- simulateResiduals(models_mortality[[which.min(table_mortality$DeltaAIC)]])
plot(m_mortality_simres,quantreg=FALSE)


# get model tables (need to rerun due to glmmtmb issues with pulling random effects variances from formula call)
m_recruitment_table = glmmTMB(m_recruitment$call$formula,data=data_no_na_for_models,family=nbinom1)
tab_model(m_recruitment_table,show.p=FALSE,transform=NULL)

m_mortality_table = glmmTMB(m_mortality$call$formula,data=data_no_na_for_models,family=nbinom1)
tab_model(m_mortality_table,show.p=FALSE,transform=NULL)


find_big_effects <- function(m)
{
  ci = confint(m)
  df_ci = data.frame(ci)
  df_ci$Variable = row.names(ci)
  row.names(df_ci) = NULL
  names(df_ci) = c("lo","hi","est","var")
  
  df_ci = df_ci %>% filter(ifelse(est < 0, hi<0, lo>0)) %>%
    filter(!var %in% c("cond.(Intercept)","Genotype.cond.Std.Dev.(Intercept)","sigma"))
  return(df_ci)
}



find_big_effects(m_recruitment_table)
find_big_effects(m_mortality_table)






















plot_model_effect <- function(m, inc.xlab=TRUE,title)
{
  
  m1 =  visreg(m,"Elevation",by="PloidyLevel",scale='response',
         xtrans = function(x) { x*sd(data_no_na$Elevation) + mean(data_no_na$Elevation) }, rug=FALSE, gg=TRUE, overlay=TRUE, band=TRUE) +
    theme_bw() + xlab(ifelse(inc.xlab==TRUE,"Elevation (m)","")) + ylab("") + ggtitle(title) + scale_color_manual(values=c("blue","red")) + scale_fill_manual(values=c("#0000FF33","#FF000033")) + labs(color="Cytotype",fill="Cytotype")
  
  m2 = visreg(m,"SummerInsolation",by="PloidyLevel",scale='response',
         xtrans = function(x) { x*sd(data_no_na$SummerInsolation) + mean(data_no_na$SummerInsolation) }, rug=FALSE, gg=TRUE, overlay=TRUE, band=TRUE) +
    theme_bw() + xlab(ifelse(inc.xlab==TRUE,expression(paste("Summer insolation (MJ m"^{-2},")")),"")) + ylab("") + scale_color_manual(values=c("blue","red")) + scale_fill_manual(values=c("#0000FF33","#FF000033")) + labs(color="Cytotype",fill="Cytotype")
  
  m3 = visreg(m,"Regolith",by="PloidyLevel",scale='response',
         rug=FALSE, gg=TRUE, overlay=TRUE, band=TRUE) +
    theme_bw() + theme(axis.text.x = element_text(angle = 15, hjust = 1))  + xlab(ifelse(inc.xlab==TRUE,"Regolith type","")) + ylab("") + scale_color_manual(values=c("blue","red")) + scale_fill_manual(values=c("#0000FF33","#FF000033")) + labs(color="Cytotype",fill="Cytotype")
  
  m4 = visreg(m,"RockUnit",by="PloidyLevel",scale='response',
              rug=FALSE, gg=TRUE, overlay=TRUE, band=TRUE) +
    theme_bw() + theme(axis.text.x = element_text(angle = 15, hjust = 1)) + xlab(ifelse(inc.xlab==TRUE,"Rock unit","")) + ylab("") + scale_color_manual(values=c("blue","red")) + scale_fill_manual(values=c("#0000FF33","#FF000033")) + labs(color="Cytotype",fill="Cytotype")
  
  m5 = visreg(m,"Cow",by="PloidyLevel",scale='response',
              rug=FALSE, gg=TRUE, overlay=TRUE, band=TRUE) +
    theme_bw() + theme(axis.text.x = element_text(angle = 15, hjust = 1)) + xlab(ifelse(inc.xlab==TRUE,"Cow","")) + ylab("") + scale_color_manual(values=c("blue","red")) + scale_fill_manual(values=c("#0000FF33","#FF000033")) + labs(color="Cytotype",fill="Cytotype")
  
  return(list(m1=m1, m2=m2, m3=m3, m4=m4, m5=m5))
}

effects_recruitment = plot_model_effect(m_recruitment_table,inc.xlab = FALSE,title='Recruitment')
effects_mortality = plot_model_effect(m_mortality_table,title='Mortality')

g_climate_impacts = ggarrange(plotlist=c(effects_recruitment,effects_mortality), 
                              nrow=2,ncol=5,common.legend=TRUE,align='hv',
                              legend='bottom',
                              labels=c("(a)",rep("",4),"(b)",rep("",4)))
ggsave(g_climate_impacts, 
       file='FIG8_g_climate_impacts.pdf',
       width=12,height=8)






# count healthy/dead/etc
mean(data$Countadultdead)
sd(data$Countadultdead)

mean(data$Countadultdamaged)
sd(data$Countadultdamaged)

mean(11 - data$Countadultdamaged - data$Countadultdead)
sd(11 - data$Countadultdamaged - data$Countadultdead)

mean(data$Countsmall)
sd(data$Countsmall)

mean(data$Plotradius[is.finite(data$Plotradius)])
sd(data$Plotradius[is.finite(data$Plotradius)])


# count tree level POMs
data_treelevel = read.csv("aspen data tree-level processed 30 Mar 2020.csv")
table(is.na(data_treelevel$POM))