### This script analyses Diversity Indices of a pan-EU run ###
### written by Maik Billing

rm(list=ls(all=TRUE))
dev.off()

library(ncdf4)
library(ggplot2)
library(gridExtra)
library(ggpubr)
library(scales)
library(cowplot)

#import user-defined functions and graphics parameters
source("C:/Users/maikb/Documents/Dissertation/R/Github/graphics_param.R")
source("C:/Users/maikb/Documents/Dissertation/R/Github/functions.R")

path.in<-"C:/Users/maikb/Documents/Dissertation/Data/LPJmL-Fit/Europe/historical/"
path.out<-"C:/Users/maikb/Documents/Dissertation/Data/LPJmL-Fit/Europe/historical/"

#set startyear of simulation data and timeframe for evaluation
startyear_simulation<-1951
startyear<-2004
endyear<-2013

istart<-startyear-startyear_simulation+1
iend<-endyear-startyear_simulation+1

theme_set(theme_bw(base_size = 30))
colorbar<-guides(fill=guide_colourbar(barwidth = 20,barheight = 2,title.position="top", title.vjust= 10))

#### load grid ####

headersize<-0
file  <- file("C:/Users/maikb/Documents/Dissertation/Data/LPJmL-Fit/Europe/historical/grid_europe_1901_1.bin","rb") # read binary mode
seek(file, where = headersize, origin="start")
x<-readBin(file,integer(),n = 2*5095, size=2)/100
lon <- x[c(1:5095)*2-1]
lat <- x[c(1:5095)*2]

variable<-"vegc"
filenames<-list.files(path.in) #get all .nc filenames in path.output.FIT
#filter filename that contain outputname
filenames<-grep(variable, filenames, value = TRUE) 
filename<-grep(startyear_simulation, filenames, value = TRUE) 
filename<-grep("_1", filename, value = TRUE) 

ncfile<-nc_open(paste(path.in,filename[1],sep=""))
data_vegc<-ncvar_get(ncfile, varid="VegC")
lats<-ncvar_get(ncfile,varid="latitude")
lons<-ncvar_get(ncfile,varid="longitude")
coord<-cbind(rep(lons,length(lats)),rep(lats,each=length(lons)))
nc_close(ncfile)

#### create nodata ####

variable<-"vegc"
filenames<-list.files(path.in) #get all .nc filenames in path.output.FIT
#filter filename that contain outputname
filenames<-grep(variable, filenames, value = TRUE) 
filenames<-grep(startyear_simulation, filenames, value = TRUE) 
data_vegc<-read.mult.runs(filenames,varid = "VegC")

data_vegc<-apply(data_vegc[,,istart:iend],c(1,2),mean,na.rm=T)
data_vegc<-data_vegc/1000

desert<-data_vegc<vegc_thres

#create noheight
variable<-"height"
varid<-"mass_height"
binsname<-"height"
scaling<-1

filenames<-list.files(path.in) #get all .nc filenames in path.output.FIT
#filter filename that contain outputname
filename<-grep(variable, filenames, value = TRUE) 
filename<-grep(startyear_simulation, filename, value = TRUE) 
filenames<-grep("mass", filename, value = TRUE) 

bins<-get.bins(filenames,binsname)
data_height<-read.mult.runs(filenames,varid = varid)

data_pft<-apply(data_height[,,,,istart:iend],c(1,2,3,4),mean,na.rm=T) #mean over years
data_var<-apply(data_pft[,,,],c(1,2,3),sum,na.rm=T) #sum over all pfts
var_mean<-apply(data_var,c(1,2),function(x) sum(x*bins/sum(x,na.rm=T)))

data_bins<-apply(data_pft,c(1,2,4),mean,na.rm=T) #mean over bins
noheight_pft<-data_bins<max(which(bins<height_thres))

#create nopft
variable<-"fpc"
zscalename<-""
filenames<-list.files(path.in) #get all .nc filenames in path.output.FIT
#filter filename that contain outputname
filenames<-grep(variable, filenames, value = TRUE) 
filename<-grep(startyear_simulation, filenames, value = TRUE) 
data_fpc<-read.mult.runs(filename,varid = "FPC")

data_pft<-apply(data_fpc[,,2:5,istart:iend],c(1,2,3),mean,na.rm=T)

nopft_pft<-data_pft<fpc_thres

nodata<-nopft_pft|noheight_pft
for(ipft in 1:4){
  nodata[,,ipft]<-nodata[,,ipft]|desert
}

df_nodata<- data.frame(nodata=as.vector(!apply(!nodata,c(1,2),any)),lon=coord[,1],lat=coord[,2])

#### load weighted and non-weighted Data ####
path.in<-"C:/Users/maikb/Documents/Dissertation/Data/LPJmL-Fit/Europe/historical/newInds/"
path.out<-"C:/Users/maikb/Documents/Dissertation/Data/LPJmL-Fit/Europe/historical/"
variable<-"trait_ind"

#weighted
filenames<-list.files(path.in) #get all .nc filenames in path.output.FIT
#filter filename that contain outputname
filenames<-grep(variable, filenames, value = TRUE) 
filenames<-grep(startyear_simulation, filenames, value = TRUE) 
filenames<-grep("quant.Rdata", filenames, value = TRUE) 
filenames<-grep("DivInd", filenames,invert = F, value = TRUE) 
filenames

for(ifile in 1:length(filenames)){
  load(paste(path.in,filenames[ifile],sep=""))
  data.array[which(data.array==-1,arr.ind = T)]<-NA
  if(ifile==1){
    data_trait_ind<-data.array
  }else{
    data_trait_ind<-data_trait_ind+data.array
  }
}
data_weight<-data_trait_ind/length(filenames)
data_weight<-apply(data_weight[,,istart:iend,],c(1,2,4),mean)

#load sd
for(ifile in 1:length(filenames)){
  load(paste(path.in,filenames[ifile],sep=""))
  data.array[which(data.array==-1,arr.ind = T)]<-NA
  if(ifile==1){
    data_sd <- array(NA, dim=c(dim(data.array), length(filenames)))
  }
  data_sd[,,,,ifile]<-data.array
}
data_sd<-apply(data_sd[,,istart:iend,,],c(1,2,4,5),mean)
data_weight_sd<-apply(data_sd,c(1,2,3),sd)


variable<-"trait_ind"
filenames<-list.files(path.in) #get all .nc filenames in path.output.FIT
#filter filename that contain outputname
filenames<-grep(variable, filenames, value = TRUE) 
filenames<-grep(startyear_simulation, filenames, value = TRUE) 
filenames<-grep("quant.Rdata", filenames,invert=T, value = TRUE) 
filenames<-grep("4d", filenames,invert = F, value = TRUE)
#filenames<-grep("noLL", filenames,invert = T, value = TRUE) 
filenames<-grep("DivInd", filenames,invert = F, value = TRUE) 
filenames
#filenames<-filenames[1:(length(filenames)/2)]

for(ifile in 1:length(filenames)){
  load(paste(path.in,filenames[ifile],sep=""))
  print(paste(path.in,filenames[ifile],sep=""))
  data.array[which(data.array==-1,arr.ind = T)]<-NA
  #df<-data.frame(FDiv=as.vector(data.array[,2,63,1]),
  #               lat=lat,lon=lon)
  #p<-ggplot(df,aes(x=lon,y=lat,fill=FDiv))+geom_raster()
  #print(p)
  if(ifile==1){
    data_trait_ind<-data.array
  }else{
    data_trait_ind<-data_trait_ind+data.array
  }
}
data_noweight<-data_trait_ind/length(filenames)
df<-data.frame(FDiv=as.vector(data_noweight[,2,63,1]),
               lat=lat,lon=lon)
p<-ggplot(df,aes(x=lon,y=lat,fill=FDiv))+geom_raster()
print(p)
data_noweight<-apply(data_noweight[,,istart:iend,],c(1,2,4),mean,na.rm=T)

#load sd
for(ifile in 1:length(filenames)){
  load(paste(path.in,filenames[ifile],sep=""))
  data.array[which(data.array==-1,arr.ind = T)]<-NA
  if(ifile==1){
    data_sd <- array(NA, dim=c(dim(data.array), length(filenames)))
  }
  data_sd[,,,,ifile]<-data.array
}
data_sd<-apply(data_sd[,,istart:iend,,],c(1,2,4,5),mean)
data_noweight_sd<-apply(data_sd,c(1,2,3),sd)
#coefficient of variation
data_noweight_sd <- data_noweight_sd/data_noweight

#### Plot Inidices, all PFTs ####

p.list_noweight<-p.list_weight<-list()

for(iind in 1:3){
  df_noweight<-data.frame(data=data_noweight[,iind,1],lat=lat,lon=lon)
  df_noweight<-merge(df_noweight,df_nodata)
  df_noweight[which(df_noweight$nodata),'data']<-NA
  df_noweight[which(df_noweight$lat>latmax),'data']<-NA
  
  p<-ggplot(data=df_noweight, aes(y=lat, x=lon)) + geom_raster(aes(fill=data)) +coord_fixed(ratio=1)
  p
  
  df_weight<-data.frame(data=data_weight[,iind,1],lat=lat,lon=lon)
  df_weight<-merge(df_weight,df_nodata)
  df_weight[which(df_weight$nodata),'data']<-NA
  df_weight[which(df_weight$lat>latmax),'data']<-NA
  
  lims<-quantile(c(df_noweight$data,df_weight$data), c(.05, .95),na.rm=T)
  #lims<-quantile(c(df_noweight$data), c(.05, .95),na.rm=T)
  #print(max(df_noweight$data,na.rm=T))
  #print(lims)
  #lims<-c(min(min(df_noweight$data,na.rm=T),min(df_weight$data,na.rm=T)),
  #        max(max(df_noweight$data,na.rm=T),max(df_weight$data,na.rm=T)))
  #if(iind==2){
  #  lims<-c(0.6625,0.824)
  #}
  #if(iind==3){
  #  lims<-c(0.621,0.830 )
  #}
  p<-ggplot(data=df_noweight, aes(y=lat, x=lon)) + geom_raster(aes(fill=data)) +coord_fixed(ratio=1)
  if(iind==1){lims<-quantile(df_noweight$data, c(.05, .95),na.rm=T)}
  p<-p+scale_fill_gradientn(oob=squish,limits=lims,name=indices.txt[iind],colors=col.zscale,breaks = pretty(lims, n = 5))
  p<-p+ylim(latmin,latmax)
  p<-p+ theme(panel.background = element_rect(fill = col.plot.bg))
  #if(iind==2){
  #  p<-p+ggtitle("Evenly weighted\n")
  #}else{
  #  p<-p+ggtitle("dummy title\n")
  #  p<-p+theme(plot.title = element_text(color="white"))}
  p<-p+colorbar+theme(legend.position="right")+theme(legend.text = element_text(size=25))
  p<-p+xlab("Longitude [°]")+ylab("Latitude [°]")
  p<-p+guides(fill=guide_colourbar(barwidth = 2,barheight = 20,title.vjust = 0.2))
  p.list_noweight[[iind]]<-p
  print(p)
  
  p<-ggplot(data=df_weight, aes(y=lat, x=lon)) + geom_raster(aes(fill=data)) +coord_fixed(ratio=1)
  if(iind==1){lims<-quantile(df_weight$data, c(.05, .95),na.rm=T)}
  p<-p+scale_fill_gradientn(oob=squish,limits=lims,name=indices.txt[iind],colors=col.zscale,breaks = pretty(lims, n = 5))
  p<-p+ylim(latmin,latmax)
  #if(iind==2){
  #  p<-p+ggtitle("Biomass weighted\n")
  #}else{
  #    p<-p+ggtitle("dummy title\n")
  #    p<-p+theme(plot.title = element_text(color="white"))}
  p<-p+ theme(panel.background = element_rect(fill = col.plot.bg))
  p<-p+colorbar+theme(legend.position="right")+theme(legend.text = element_text(size=25))
  p<-p+xlab("Longitude [°]")+ylab("Latitude [°]")
  p<-p+guides(fill=guide_colourbar(barwidth = 2,barheight = 20,title.vjust = 0.2))
  p.list_weight[[iind]]<-p
  print(p)
}

#layout_matrix<-rbind(c(1,2,3),c(NA,4,5))

layout_matrix<-rbind(c(1,2,3),c(4,5,6))

#plot.all.ind<-grid.arrange(p.list_noweight[[1]],p.list_noweight[[2]],p.list_weight[[2]],
#                                               p.list_noweight[[3]],p.list_weight[[3]],
#                           layout_matrix=layout_matrix,
#                           labels=c("a)","b)","c)","d)","e)"))


plot.all.ind<-ggarrange(p.list_noweight[[1]],p.list_noweight[[2]],p.list_noweight[[3]],
                        p.list_weight[[1]],p.list_weight[[2]],p.list_weight[[3]],
                           #layout_matrix=layout_matrix,
                           labels=c("a)","b)","c)","","d)","e)"),font.label = list(size = 30))
#plot.all.ind<-ggarrange(p.list_noweight[[1]],p.list_noweight[[2]],p.list_noweight[[3]],
#                        labels=c("a)","b)","c)"),font.label = list(size = 30))

ggsave(paste(path.out,"DivInds_4d_noweight_vs_quant.jpg",sep=""),device = "jpg",dpi = 600,plot.all.ind,scale=3,width=20,height=10, units="cm")



#### Plot Inidices, all PFTs, SD ####

p.list_noweight<-p.list_weight<-list()

for(iind in 1:3){
  df_noweight<-data.frame(data=data_noweight_sd[,iind,1],lat=lat,lon=lon)
  df_noweight<-merge(df_noweight,df_nodata)
  df_noweight[which(df_noweight$nodata),'data']<-NA
  df_noweight[which(df_noweight$lat>latmax),'data']<-NA
  
  p<-ggplot(data=df_noweight, aes(y=lat, x=lon)) + geom_raster(aes(fill=data)) +coord_fixed(ratio=1)
  p
  
  df_weight<-data.frame(data=data_weight_sd[,iind,1],lat=lat,lon=lon)
  df_weight<-merge(df_weight,df_nodata)
  df_weight[which(df_weight$nodata),'data']<-NA
  df_weight[which(df_weight$lat>latmax),'data']<-NA
  
  lims<-quantile(c(df_noweight$data,df_weight$data), c(.05, .95),na.rm=T)
  lims<-quantile(c(df_noweight$data), c(0.05, 0.95),na.rm=T)
  #print(max(df_noweight$data,na.rm=T))
  #print(lims)
  #lims<-c(min(min(df_noweight$data,na.rm=T),min(df_weight$data,na.rm=T)),
  #        max(max(df_noweight$data,na.rm=T),max(df_weight$data,na.rm=T)))
  #if(iind==2){
  #  lims<-c(0.6625,0.824)
  #}
  #if(iind==3){
  #  lims<-c(0.621,0.830 )
  #}
  p<-ggplot(data=df_noweight, aes(y=lat, x=lon)) + geom_raster(aes(fill=data)) +coord_fixed(ratio=1)
  if(iind==1){lims<-quantile(df_noweight$data, c(0.05, 0.95),na.rm=T)}
  p<-p+scale_fill_gradientn(oob=squish,limits=lims,name=paste("COV of ",indices.txt[iind],sep=""),colors=col.zscale,breaks = pretty(lims, n = 5))
  p<-p+ylim(latmin,latmax)
  p<-p+ theme(panel.background = element_rect(fill = col.plot.bg))
  #if(iind==2){
  #  p<-p+ggtitle("Evenly weighted\n")
  #}else{
  #  p<-p+ggtitle("dummy title\n")
  #  p<-p+theme(plot.title = element_text(color="white"))}
  p<-p+colorbar+theme(legend.position="right")+theme(legend.text = element_text(size=25))
  p<-p+xlab("Longitude [°]")+ylab("Latitude [°]")
  p<-p+guides(fill=guide_colourbar(barwidth = 2,barheight = 20,title.vjust = 0.2))
  p.list_noweight[[iind]]<-p
  print(p)
  
  p<-ggplot(data=df_weight, aes(y=lat, x=lon)) + geom_raster(aes(fill=data)) +coord_fixed(ratio=1)
  if(iind==1){lims<-quantile(df_weight$data, c(.05, .95),na.rm=T)}
  p<-p+scale_fill_gradientn(oob=squish,limits=lims,name=indices.txt[iind],colors=col.zscale,breaks = pretty(lims, n = 5))
  p<-p+ylim(latmin,latmax)
  #if(iind==2){
  #  p<-p+ggtitle("Biomass weighted\n")
  #}else{
  #    p<-p+ggtitle("dummy title\n")
  #    p<-p+theme(plot.title = element_text(color="white"))}
  p<-p+ theme(panel.background = element_rect(fill = col.plot.bg))
  p<-p+colorbar+theme(legend.position="right")+theme(legend.text = element_text(size=25))
  p<-p+xlab("Longitude [°]")+ylab("Latitude [°]")
  p<-p+guides(fill=guide_colourbar(barwidth = 2,barheight = 20,title.vjust = 0.2))
  p.list_weight[[iind]]<-p
  #print(p)
}

#layout_matrix<-rbind(c(1,2,3),c(NA,4,5))

layout_matrix<-rbind(c(1,2,3),c(4,5,6))

#plot.all.ind<-grid.arrange(p.list_noweight[[1]],p.list_noweight[[2]],p.list_weight[[2]],
#                                               p.list_noweight[[3]],p.list_weight[[3]],
#                           layout_matrix=layout_matrix,
#                           labels=c("a)","b)","c)","d)","e)"))


plot.all.ind<-ggarrange(p.list_noweight[[1]],p.list_noweight[[2]],p.list_noweight[[3]],
                        p.list_weight[[1]],p.list_weight[[2]],p.list_weight[[3]],
                        #layout_matrix=layout_matrix,
                        labels=c("a)","b)","c)","","d)","e)"),font.label = list(size = 30))
plot.all.ind<-ggarrange(p.list_noweight[[1]],p.list_noweight[[2]],p.list_noweight[[3]],
                        labels=c("a)","b)","c)"),font.label = list(size = 30))

ggsave(paste(path.out,"DivInds_4d_noweight_vs_quant_sd.jpg",sep=""),device = "jpg",dpi = 600,plot.all.ind,scale=3,width=20,height=10, units="cm")



str(data.array)

df<-data.frame(FDiv=as.vector(data.array[,2,63,1]),
               lat=lat,lon=lon)
ggplot(df,aes(x=lon,y=lat,fill=FDiv))+geom_raster()

#ggsave(paste(path.out,"DivInds.jpg",sep=""),device = "jpg",dpi = 300,plot.all.ind,scale=3,width=12,height=12, units="cm")

#p.all.ind_weight<-ggarrange(p.list_weight[[1]],p.list_weight[[2]],p.list_weight[[3]],labels=c("a)","b)","c)"),align="none",font.label = list(size = 30))
#p.all.ind_weight

#ggsave(paste(path.out,"DivInds_weight.jpg",sep=""),device = "jpg",dpi = 300,p.all.ind_weight,scale=3,width=12,height=12, units="cm")
#
#p.all.ind_noweight<-ggarrange(p.list_noweight[[1]],p.list_noweight[[2]],p.list_noweight[[3]],labels=c("a)","b)","c)"),align="none",font.label = list(size = 30))
#p.all.ind_noweight

#ggsave(paste(path.out,"DivInds_noweight.jpg",sep=""),device = "jpg",dpi = 300,p.all.ind_noweight,scale=3,width=12,height=12, units="cm")

#### Plot Inidices, each PFT ####

pfts<-c("BL-S","BL-E","B-NL","T-NL")

for(ipft in 1:4){
  
  p.list_weight<-p.list_noweight<-list()
  
  for(iind in 1:3){
    df_weight<-data.frame(data=data_weight[,iind,ipft+1],lat=lat,lon=lon)
    df_nodata<- data.frame(nodata=as.vector(nodata[,,ipft]),lon=coord[,1],lat=coord[,2])
    df_weight<-merge(df_weight,df_nodata)
    df_weight[which(df_weight$nodata),'data']<-NA
    
    df_noweight<-data.frame(data=data_noweight[,iind,ipft+1],lat=lat,lon=lon)
    df_noweight<-merge(df_noweight,df_nodata)
    df_noweight[which(df_noweight$nodata),'data']<-NA
    
    lims<-c(quantile(c(df_noweight$data,df_weight$data), c(.05, .95),na.rm=T))
    
    #lims<-c(min(min(df_noweight$data,na.rm=T),min(df_weight$data,na.rm=T)),
    #        max(max(df_noweight$data,na.rm=T),max(df_weight$data,na.rm=T)))
    
    p<-ggplot(data=df_weight, aes(y=lat, x=lon)) + geom_raster(aes(fill=data)) +coord_fixed(ratio=1)
    p<-p+scale_fill_gradientn(oob=squish,limits=lims,name=indices.txt[iind],colors=col.zscale,breaks = pretty(lims, n = 4))
    p<-p+ylim(latmin,latmax)
    p<-p+ theme(panel.background = element_rect(fill = col.plot.bg))
    p<-p+colorbar+theme(legend.position="right")+theme(legend.text = element_text(size=25))
    p<-p+xlab("Longitude [?]")+ylab("Latitude [?]")
    p<-p+guides(fill=guide_colourbar(barwidth = 2,barheight = 20,title.vjust = 0.2))
    #p<-p+ggtitle(pfts[ipft])
    p.list_weight[[iind]]<-p
    print(p)
    
    p<-ggplot(data=df_noweight, aes(y=lat, x=lon)) + geom_raster(aes(fill=data)) +coord_fixed(ratio=1)
    p<-p+scale_fill_gradientn(oob=squish,limits=lims,name=indices.txt[iind],colors=col.zscale,breaks = pretty(lims, n = 4))
    p<-p+ylim(latmin,latmax)
    p<-p+ theme(panel.background = element_rect(fill = col.plot.bg))
    p<-p+colorbar+theme(legend.position="right")+theme(legend.text = element_text(size=25))
    p<-p+xlab("Longitude [?]")+ylab("Latitude [?]")
    p<-p+guides(fill=guide_colourbar(barwidth = 2,barheight = 20,title.vjust = 0.2))
    #p<-p+ggtitle(pfts[ipft])
    p.list_noweight[[iind]]<-p
    print(p)
  }
  
  p.all.ind_weight<-ggarrange(p.list_weight[[1]],p.list_weight[[2]],p.list_weight[[3]],labels=c("a)","b)","c)"),align="none",font.label = list(size = 30))
  p.all.ind_weight
  
  ggsave(paste(path.out,"DivInds_",pfts[ipft],"_weight_4d.jpg",sep=""),device = "jpg",dpi = 300,p.all.ind_weight,scale=4,width=12,height=8, units="cm")
  
  p.all.ind_noweight<-ggarrange(p.list_noweight[[1]],p.list_noweight[[2]],p.list_noweight[[3]],labels=c("a)","b)","c)"),align="none",font.label = list(size = 30))
  p.all.ind_noweight
  
  ggsave(paste(path.out,"DivInds_",pfts[ipft],"_noweight_4d.jpg",sep=""),device = "jpg",dpi = 300,p.all.ind_noweight,scale=4,width=12,height=8, units="cm")

}