rm(list=ls())
gc(TRUE)

# Reload library
library(mgcv)
library(raster)
library(colorspace)
library(BayesX)

source("plot_legend.R")

# Load saved models
Name1 <- "SpatTemp3d_1980_aniso_optim_" 
Opt_type <- "gcv_gcv"

asy_vec_load <- c( "1", "2", "5","10","20","50","80","90","95","98","99")
asy_vec      <- c("01","02","05","10","20","50","80","90","95","98","99")
#asy_vec_load <- c("10","50","90")
#asy_vec <- c("10","50","90")

for(asy in 1:length(asy_vec)) {
    asy1 <- asy_vec_load[asy]
    asy2 <- asy_vec[asy]
    cat("\n", asy1)
    Name_temp <- paste(Name1,asy1,"_",Opt_type,"_1e-04_10_1e-05_10_T_",sep="")
    load(paste(Name_temp,".RData",sep=""))
    assign(paste("rainmodel",asy2,sep=""), rainmodel2)
}


Name_joint <- paste(Name1,Opt_type,"_Supplement",sep="")

rm("rainmodel2","rainmodel1")
gc(TRUE)


# Define Covariates for predictions
excl_elev_day <- c("ti(day)","ti(lon)","ti(lat)","ti(lon,day)","ti(day,lat)","ti(lon,lat)","ti(day,lon,lat)")
incl_elev_Mar <- c("s(elevation)","ti(lon)","ti(lat)","ti(lon,lat)")
incl_elev_day <- c("ti(day)","s(elevation)","ti(lon)","ti(lat)","ti(lon,day)","ti(day,lat)","ti(lon,lat)","ti(day,lon,lat)")


# Load information of stations to plot locations of cities
Meta_Data1 <- read.fwf("KL_Tageswerte_Beschreibung_Stationen.txt",
                       widths=c(nchar("Stations_id"), 
                                nchar(" von_datum"),
                                nchar(" bis_datum"),
                                nchar(" Stationshoehe"),
                                nchar(" geoBreite "),
                                nchar("geoLaenge ")),
                       header=F, stringsAsFactors=F)
Meta_Data2 <- Meta_Data1[-c(1,2),]
colnames(Meta_Data2) <- gsub(pattern=" ", replacement="", x=Meta_Data1[1,])
Meta_Data3 <- Meta_Data2

Meta_Data3$Stations_id   <- as.numeric(Meta_Data2$Stations_id)
Meta_Data3$von_datum     <- as.numeric(Meta_Data2$von_datum)
Meta_Data3$bis_datum     <- as.numeric(Meta_Data2$bis_datum)
Meta_Data3$Stationshoehe <- as.numeric(Meta_Data2$Stationshoehe)
Meta_Data3$geoBreite     <- as.numeric(Meta_Data2$geoBreite)
Meta_Data3$geoLaenge     <- as.numeric(Meta_Data2$geoLaenge)




Meta_Data <- Meta_Data3[,c("Stations_id","Stationshoehe","geoBreite","geoLaenge")]
colnames(Meta_Data)[colnames(Meta_Data) == "geoBreite"] <- "lat"
colnames(Meta_Data)[colnames(Meta_Data) == "geoLaenge"] <- "lon"
colnames(Meta_Data)[colnames(Meta_Data) == "Stationshoehe"] <- "elevation"


# Define ID of cities to plot
STATIONS_ID_Hamburg <- 1975
STATIONS_ID_Berlin  <- 433
STATIONS_ID_Munich <- 3379
STATIONS_ID_Cologne <- 2667

# Define data set for elevation forcasts
newdata_Main_elev <- data.frame("day"      = 0,
                                "lon"      = rep(0      , times=250),
                                "lat"      = rep(0      , times=250),
                                "year"     = rep(0      , times=250),
                                "elevation"= seq(min(homstart$elevation,na.rm=T),max(homstart$elevation,na.rm=T),length=250))

# Define data set for year forcasts
newdata_Main_year <- data.frame("day"      = 0,
                                "lon"      = rep(0      , times=250),
                                "lat"      = rep(0      , times=250),
                                "elevation"= rep(0      , times=250),
                                "year"= seq(min(homstart$year,na.rm=T),max(homstart$year,na.rm=T),length=250))

# Define data set for daily forcasts
newdata_Main <- data.frame("day"      = 1:366,
                           "lon"      = rep(0      , times=366),
                           "lat"      = rep(0      , times=366),
                           "year"     = rep(0      , times=366),
                           "elevation"= rep(0, times=366))


# Define data set for daily forcasts for specific cities
newdata_Hamburg <- data.frame("day"      = 1:366,
                              "year"     = rep(0      , times=366),
                              "lon"      = rep(Meta_Data[Meta_Data$Stations_id==STATIONS_ID_Hamburg,"lon"]      , times=366),
                              "lat"      = rep(Meta_Data[Meta_Data$Stations_id==STATIONS_ID_Hamburg,"lat"]      , times=366),
                              "elevation"= rep(Meta_Data[Meta_Data$Stations_id==STATIONS_ID_Hamburg,"elevation"], times=366))

newdata_Berlin <- data.frame("day"      = 1:366,
                             "year"     = rep(0      , times=366),
                             "lon"      = rep(Meta_Data[Meta_Data$Stations_id==STATIONS_ID_Berlin,"lon"]      , times=366),
                             "lat"      = rep(Meta_Data[Meta_Data$Stations_id==STATIONS_ID_Berlin,"lat"]      , times=366),
                             "elevation"= rep(Meta_Data[Meta_Data$Stations_id==STATIONS_ID_Berlin,"elevation"], times=366))

newdata_Munich <- data.frame("day"      = 1:366,
                             "year"     = rep(0      , times=366),
                             "lon"      = rep(Meta_Data[Meta_Data$Stations_id==STATIONS_ID_Munich,"lon"]      , times=366),
                             "lat"      = rep(Meta_Data[Meta_Data$Stations_id==STATIONS_ID_Munich,"lat"]      , times=366),
                             "elevation"= rep(Meta_Data[Meta_Data$Stations_id==STATIONS_ID_Munich,"elevation"], times=366))

newdata_Cologne <- data.frame("day"      = 1:366,
                              "year"     = rep(0      , times=366),
                              "lon"      = rep(Meta_Data[Meta_Data$Stations_id==STATIONS_ID_Cologne,"lon"]      , times=366),
                              "lat"      = rep(Meta_Data[Meta_Data$Stations_id==STATIONS_ID_Cologne,"lat"]      , times=366),
                              "elevation"= rep(Meta_Data[Meta_Data$Stations_id==STATIONS_ID_Cologne,"elevation"], times=366))




y_elev_max <- 1800
if(max(homstart$elevation,na.rm=T) < 1500) y_elev_max <- 1500
if(max(homstart$elevation,na.rm=T) < 1000) y_elev_max <- 1000

y_year_max <- 2014

city_vec <- c("Hamburg","Munich","Cologne","Berlin")

# Build predictions of univariate effects and the daily main effect

for(asy1 in asy_vec) {
    cat("\n", asy1)
    rainmodel_temp <- get(paste("rainmodel",asy1,sep=""))
    predict_Main_temp_elev    <- predict(rainmodel_temp$model,newdata=newdata_Main_elev   , type="terms")[,"s(elevation)"]
    predict_Main_temp_year    <- predict(rainmodel_temp$model,newdata=newdata_Main_year   , type="terms")[,"s(year)"]
    predict_Main_temp_day     <- predict(rainmodel_temp$model,newdata=newdata_Main        , type="terms")[,"ti(day)"]
    assign(paste("predict_Main_",asy1,"_elev",sep=""), predict_Main_temp_elev)
    assign(paste("predict_Main_",asy1,"_year",sep=""), predict_Main_temp_year)
    assign(paste("predict_Main_",asy1,"_day" ,sep=""), predict_Main_temp_day)
    
    for(city1 in city_vec){
        data_city1 <- get(paste("newdata_",city1,sep=""))
        predict_t_city_temp     <- predict(rainmodel_temp$model,newdata=data_city1    , type="terms")
        predict_city_temp     <- rowSums(predict_t_city_temp    [,incl_elev_day])
        assign(paste("predict_",city1,"_",asy1 ,sep=""), predict_city_temp)
    }
}
    
    




# Load spatial Metadata and build forcast design matrix for spatial effects
Germany <- readRDS("DEU_adm0.rds")
TopoGermany <- raster("GermanyTopo.tif")

x_min <- Germany@bbox["x","min"]
x_max <- Germany@bbox["x","max"]
y_min <- Germany@bbox["y","min"]
y_max <- Germany@bbox["y","max"]
grid_d <- 500


Overall_mean <- expand.grid(lon=seq(x_min,x_max,length=grid_d),lat=seq(y_min,y_max,length=grid_d))

Overall_mean$day  <- 0
Overall_mean$year <- 1980

co <- (as.matrix(Overall_mean[,c("lon","lat")]))


attr(co,"dimnames") <- NULL

Overall_mean$elevation <- extract(TopoGermany, co)
Overall_mean$elevation[which(Overall_mean$elevation < 0)] <- 0

Overall_31Jul <- Overall_31Jan <- Overall_mean



Overall_31Jan$day <- 31
Overall_31Jul$day <- 212


spattemp_vec <- c("31Jan","31Jul")



# Forcast spatial effects

for(asy1 in asy_vec) {
    cat("\n", asy1)
    rainmodel_temp <- get(paste("rainmodel",asy1,sep=""))
    
    for(spattemp1 in spattemp_vec){
        data_spattemp1 <- get(paste("Overall_",spattemp1,sep=""))
        predict_t_spattemp_temp     <- predict(rainmodel_temp$model,newdata=data_spattemp1    , type="terms")
        predict_Main_excl_elev_spattemp  <- rowSums(predict_t_spattemp_temp[,excl_elev_day, drop=F])
        Overall_spattemp <- matrix(predict_Main_excl_elev_spattemp,nrow=grid_d, byrow=FALSE)
        
        assign(paste("Overall_",asy1,"_M_excl_",spattemp1 ,sep=""), Overall_spattemp)
    }
}



















# Define some characteristics for the plot design
Length_month <- c("a"=0,"Jan"=31,"Feb"=28,"Mar"=31,"Apr"=30,"May"=31,"Jun"=30,"Jul"=31,"Aug"=31,"Sep"=30,"Oct"=31,"Nov"=30,"Dec"=31)
C_Length_month <- cumsum(Length_month)
names(C_Length_month) <- c(names(Length_month)[-1],"")


hcl.par = list(h=c(260,0), c=185, l=c(30,85))
col_vector <- diverge_hcl(n=length(asy_vec),h=hcl.par$h,c=hcl.par$c,l=hcl.par$l)
col_vector[ceiling(length(col_vector)/2)] <- "darkgrey"

pch_vec <- c(21,23,24,25)




ylim_cities <- c(-10,25)
ylim_cities2 <- c(-10,25)
ylim_elev   <- c(-10,25)
ylim_year   <- c(-10,25)





######################################################################################################
######################################################################################################

# Plot splines for univariate effects
labels_year <- names(C_Length_month)
adj_year <- 1.00

pdf(paste(Name_joint,"2.pdf",sep=""),width=3.5,height=3.5)
    
# Plot of elevation   
par(mar=c(4,4,2,1))
uz <- 1
plot (newdata_Main_elev$elevation,predict_Main_10_elev + rainmodel10$model$coefficient[1], col=col_vector[1],
      lty=1,lwd=2, type="n",ylim=ylim_elev,axes=F,xlab="elevation",ylab=expression(f[tau](elevation)), main="")
for(asy1 in asy_vec) {
    predict_temp    <- get(paste("predict_Main_",asy1,"_elev"   ,sep=""))
    rainmodel_temp  <- get(paste("rainmodel"       ,asy1,sep=""))
    
    lines(newdata_Main_elev$elevation,predict_temp + rainmodel_temp$model$coefficient[1], col=col_vector[uz],lty=1,lwd=2)
    uz <- uz + 1
}
axis(side=1,at=seq(0,y_elev_max,by=250),labels=seq(0,y_elev_max,by=250))
axis(side=2,at=seq(-15,30,by=5),labels=seq(-15,30,by=5))
box()
rug(unique(homstart$elevation))



# Plot of day   
uz <- 1
plot (newdata_Main$day,predict_Main_10_day + rainmodel10$model$coefficient[1], col=col_vector[1], 
      lty=1, lwd=2, type="n", xlim=c(0,366), ylim=ylim_cities, axes=F, xlab="day", ylab=expression(f[tau](day)), main="")
for(asy1 in asy_vec) {
    predict_temp    <- get(paste("predict_Main_",asy1,"_day"   ,sep=""))
    rainmodel_temp  <- get(paste("rainmodel"       ,asy1,sep=""))
    lines(newdata_Main$day, predict_temp + rainmodel_temp$model$coefficient[1], col=col_vector[uz],lty=1,lwd=2)
    uz <- uz + 1
}
axis(side=1,at=C_Length_month[]+15,labels=labels_year,hadj=adj_year,las=2,col.ticks="white")
axis(side=1,at=C_Length_month[],labels=rep("",times=length(C_Length_month)),hadj=adj_year,las=2,col.ticks="black")
axis(side=2,at=seq(-15,30,by=5),labels=seq(-15,30,by=5))
box()


# Plot of year   
uz <- 1
plot (newdata_Main_year$year,predict_Main_10_year + rainmodel10$model$coefficient[1] , col=col_vector[1],
      lty=1, lwd=2, type="n", xlim=c(1980,2015), ylim=ylim_year, axes=F, xlab="year", ylab=expression(f[tau](year)), main="")
for(asy1 in asy_vec) {
    predict_temp    <- get(paste("predict_Main_",asy1,"_year"   ,sep=""))
    rainmodel_temp  <- get(paste("rainmodel"       ,asy1,sep=""))
    lines(newdata_Main_year$year, predict_temp + rainmodel_temp$model$coefficient[1], col=col_vector[uz],lty=1,lwd=2)
    abline(h=rainmodel_temp$model$coefficient[1],lty=2,col="black",lwd=0.5)
    
    uz <- uz + 1
}
axis(side=1,at=seq(1980,2015,by=5),labels=seq(1980,2015,by=5))
axis(side=2,at=seq(-15,30,by=5),labels=seq(-15,30,by=5))
box()

rug(unique(homstart$year))

dev.off()

# Plot legend of splines
pdf(paste(Name_joint,"2_legend.pdf",sep=""),width=7,height=1)
par(mar=c(0,0,0,0))
plot.new()
legend("center",legend=c("Hamburg","Berlin","Cologne","Munich"),col=1,pt.bg=6,
       pch=pch_vec,pt.cex=1.8,horiz=T,cex=1)
dev.off()

pdf(paste(Name_joint,"2_legend2.pdf",sep=""),width=11,height=1)
par(mar=c(0,0,0,0))
plot.new()
legend(x="center", lty=1,lwd=2, col=(col_vector), legend=(paste(asy_vec_load,"%",sep="")), bg="white", horiz=T)
dev.off()



# plot spatial effects
ik0_vec <- "excl"
ik1_vec <- asy_vec
ik2_vec <- c("31Jan","31Jul")

zlim1 <- NULL

zlim2 <- matrix(NA,ncol=2,nrow=2)
rownames(zlim2) <- c("excl_31Jan","excl_31Jul")

zlim2["excl_31Jan",] <- c(-18,18)
zlim2["excl_31Jul",] <- c(-18,18)

hcl.par = list(h=c(260,0), c=c(80,290), l=c(15,95))

pdf(paste(Name_joint,"3.pdf",sep=""),width=7,height=7)
for(ik0 in ik0_vec){
    for(ik2 in ik2_vec){
        par(mfrow=c(1,1), mar=c(4, 4, 4, 2) + 0.1, oma=c(1,1,0,0))
        zlim1 <- zlim2[paste(ik0,ik2,sep="_"),]
        col_map <- diverge_hcl(n=zlim1[2]-zlim1[1],h=hcl.par$h,c=hcl.par$c,l=hcl.par$l)
        for(ik1 in ik1_vec){
            Overall_mean_M <- get(paste("Overall",ik1,"M",ik0,ik2,sep="_"))
            if(min(as.numeric(Overall_mean_M),na.rm=T) < zlim1[1]) {
                cat(ik0, " ", ik2, " ", ik1, ": Min: ", min(as.numeric(Overall_mean_M),na.rm=T), " zlim:", zlim1[1], "\n")
            }
            if(max(as.numeric(Overall_mean_M),na.rm=T) > zlim1[2]) {
                cat(ik0, " ", ik2, " ", ik1, ": Max: ", max(as.numeric(Overall_mean_M),na.rm=T), " zlim:", zlim1[2], "\n") 
            }
            ik2_text <- ik2
            if(ik2_text == "Main") ik2_text <- ""
            
            image(x=seq(x_min,x_max,length=grid_d),y=seq(y_min,y_max,length=grid_d),z=Overall_mean_M,
                  main=paste(ik2_text,"",ik1,"%"),useRaster=T,xlab="",ylab="", col=col_map,zlim=zlim1,cex.main=3,cex.axis=2)
            
            if(ik0 == "excl"){
                contour(x=seq(x_min,x_max,length=grid_d),y=seq(y_min,y_max,length=grid_d),z=Overall_mean_M,zlim=zlim1,nlevels=zlim1[2]-zlim1[1],lab.cex=1,add=T,col="black",drawlabels=FALSE)
            }
            
            plot(Germany,add=T)
            
            kk_n <- 1
            for(kk in c("STATIONS_ID_Hamburg","STATIONS_ID_Berlin","STATIONS_ID_Cologne","STATIONS_ID_Munich")) {
                kk2 <- get(kk)
                points(Meta_Data[Meta_Data$Stations_id==kk2,"lon"],Meta_Data[Meta_Data$Stations_id==kk2,"lat"],pch=pch_vec[kk_n],cex=4,
                       bg=6,col=1,lwd=1.2)
                kk_n <- kk_n+1
            }
        }
        mtext(side=1,text="longitude",outer=T,line=-1,cex=1.5)
        mtext(side=2,text="latitude" ,outer=T,line=-0.75,cex=1.5)
        
    }
}

dev.off()

# Plot legend of spatial effects
pdf(paste(Name_joint,"3_legend.pdf",sep=""),width=4,height=1)

for(ik0 in ik0_vec){
    for(ik2 in ik2_vec){
        zlim1 <- zlim2[paste(ik0,ik2,sep="_"),]
        col_map <- diverge_hcl(n=zlim1[2]-zlim1[1],h=hcl.par$h,c=hcl.par$c,l=hcl.par$l)
        
        par(mar=c(0,0,0,0))
        plot_legend(ylim_spat=zlim1, xlimits = c(1,3), ylimits = c(1,5), 
                    cex.legend = 1, digits = 2, par_new = FALSE, 
                    shift_text = c(0.5,-0.5),col_map=col_map) 
    }
}
dev.off()


# Plots per City

pdf(paste(Name_joint,"_per_City.pdf",sep=""),width=3.5,height=3.5)

for(asy1 in asy_vec) {
    predict_Main_temp    <- get(paste("predict_Main_"   ,asy1,"_day",sep=""))
    predict_Hamburg_temp <- get(paste("predict_Hamburg_",asy1,sep=""))
    predict_Munich_temp  <- get(paste("predict_Munich_" ,asy1,sep=""))
    predict_Berlin_temp  <- get(paste("predict_Berlin_" ,asy1,sep=""))
    predict_Cologne_temp <- get(paste("predict_Cologne_",asy1,sep=""))
    rainmodel_temp       <- get(paste("rainmodel"       ,asy1,sep=""))
    
    par(mar=c(4,4,2,1))
    
    plot (newdata_Main$day,predict_Hamburg_temp + rainmodel_temp$model$coefficient[1]  , type="n", xlim=c(0,366), 
          ylim=ylim_cities2,axes=F,xlab="day",ylab=expression(f[tau](day)), main=paste(asy1,"%",sep=""))
    lines(newdata_Main$day,predict_Main_temp    + rainmodel_temp$model$coefficient[1]  , col="grey",lty=1,lwd=2)
    lines(newdata_Main$day,predict_Hamburg_temp + rainmodel_temp$model$coefficient[1]  , col=2,lty=2,lwd=2)
    lines(newdata_Main$day,predict_Berlin_temp  + rainmodel_temp$model$coefficient[1]  , col="darkgreen",lty=4,lwd=2)
    lines(newdata_Main$day,predict_Cologne_temp + rainmodel_temp$model$coefficient[1]  , col=4,lty=5,lwd=2)
    lines(newdata_Main$day,predict_Munich_temp  + rainmodel_temp$model$coefficient[1]  , col=6,lty=6,lwd=2)
    axis(side=1,at=C_Length_month[]+15,labels=labels_year,hadj=adj_year,las=2,col.ticks="white")
    axis(side=1,at=C_Length_month[],labels=rep("",times=length(C_Length_month)),hadj=adj_year,las=2,col.ticks="black")
    axis(side=2,at=seq(-15,30,by=5),labels=seq(-15,30,by=5))
    box()
    
}

dev.off()

pdf(paste(Name_joint,"_per_City_legend.pdf",sep=""),width=7,height=1)
par(mar=c(0,0,0,0))
plot.new()
legend("center",legend=c("Main","Hamburg","Berlin","Cologne","Munich"),col=c("grey",2,"darkgreen",4,6),
       lty=c(1,2,4,5,6),lwd=2,horiz=T,cex=1)
dev.off()



