## -----------------------------------
## Function for CMIP6 extraction:
## -----------------------------------

# Sebastian Sippel
# 02.09.2021


## Define Raster template:
library(raster)
raster.template = raster(res = 5, xmn = 0, xmx=360, ymn = -90, ymx=90)
areaw=c(matrix(values(raster::area(raster.template)), 72,36)[,36:1]) / sum(c(matrix(values(raster::area(raster.template)), 72,36)[,36:1]))




## Function to repeat vector either by row or column:
rep.row<-function(x,n){
  matrix(rep(x,each=n),nrow=n)
}
rep.col<-function(x,n){
  matrix(rep(x,each=n), ncol=n, byrow=TRUE)
}



## Function to get list of files from CMIP6:
get.CMIP6.file.list <- function(vari= "tas", temp.res = "ann", scen = c("historical", "ssp119", "ssp126", "ssp245", "ssp370", "ssp434", "ssp534", "ssp585"), 
                                CMIP6.dir = "/net/cfc/cmip6/Next_Generation/tas/ann/g025") {
  
  # 1. get  list for each scenario:
  scen.list = list()
  for (scen.idx in 1:length(scen)) {
    scen.list[[scen.idx]] = list.files(path = paste(CMIP6.dir, sep=""), pattern = scen[scen.idx])
  }
  file.name=unlist(scen.list)
  
  # get data.frame with overview of variables:
  vari=rep(vari, length(file.name))
  res=rep(temp.res, length(file.name))
  scen=sapply(strsplit(file.name,"_"),function(x) paste(x[4],collapse="_"))
  mod=sapply(strsplit(file.name,"_"),function(x) paste(x[3],collapse="_"))
  ens.mem=sapply(strsplit(file.name,"_"),function(x) paste(x[5],collapse="_"))
  modall=sapply(strsplit(file.name,"_"),function(x) paste(x[3:5],collapse="_"))
  modcl=sapply(as.character(mod),function(x) substr(x,1,3)[[1]][1])
  
  if ( scen == "piControl" ) {
    period.length=rep(NA, length(file.name))
    
    for(i in 1:length(file.name)) {
      print(i)
      test=nc_open(paste(CMIP6.dir, "/", file.name, sep="")[i])
      # brick(paste(CMIP6.dir, file.name[i], sep=""), varname="tas")
      period.length[i]=dim(ncvar_get(test, varid = vari[1], start=c(1,1,1), count=c(1,1,-1))) # [3]
    }
    return(data.frame(file.name, vari, res, mod, modcl, scen, ens.mem, modall, period.length, stringsAsFactors=F))
  }

  return(data.frame(file.name, vari, res, mod, modcl, scen, ens.mem, modall, stringsAsFactors=F))
}


## Check whether two files are the same:
check.same.2files <- function(f.var1, f.var2) {
  
  ix.var1 = which(f.var1$modall %in% f.var2$modall)
  ix.var2 = which(f.var2$modall %in% f.var1$modall)
  print(all(f.var1$modall[ix.var1] == f.var2$modall[ix.var2]))
  
  # ret.list = list()
  # ret.list[[1]] = f.var1[ix.var1,]
  # ret.list[[2]] = f.var2[ix.var2,]
  return(f.var1[ix.var1,])
}




## Read CMIP6 files: 
read.CMIP6_novar <- function(file.name, var, res="ann", scen, CMIP5.dir, time.count = NA, subtract.drift.piControl = F, rm.years = 100) {
  X = list()
  for (i in 1:length(file.name)) {
    print(i)
    if (scen[i] == "piControl") {  # piControl runs are shortened to 200 years...:
      if(res=="mon") { res.fact<-12 } else res.fact<-1;
      
        if (subtract.drift.piControl == F) {
          temp = apply(get_CMIP5_array(file = paste(CMIP5.dir, "/", file.name[i], sep=""), time.count = time.count[i], var=var), 3, c)
          X[[i]] = temp[,(rm.years+1) :(dim(temp)[2])]
        } else if (subtract.drift.piControl == T) {
          temp = apply(get_CMIP5_array(file = paste(CMIP5.dir, "/", file.name[i], sep=""), time.count = time.count[i], var=var), 3, c)
          X[[i]] = t(apply(X = temp[,(rm.years+1):(dim(temp)[2])], MARGIN=1, FUN=function(x) x - lm(x ~ c(1:(dim(temp)[2]-rm.years)))$fitted)) + rep.col(rowMeans(temp[,(rm.years+1):(dim(temp)[2])]), dim(temp)[2]-rm.years)
        }
    } else {
      X[[i]] = apply(get_CMIP5_array(file = paste(CMIP5.dir, "/", file.name[i], sep=""), var=var), 3, c)
      # get_CMIP5_vec(file = paste(CMIP5.dir, "/", file.name[i], sep=""), var="time")
    }
  }
  return(X)
}



## Reshape CMIP6 files into XAX data structure:
get.XAX_ann <- function(X, ncol = 72*36, M, start.year = 1850, cmip) {
  
  # OVERALL DIMENSION
  s = sum(sapply(X = X, FUN=function(x) dim(x)[2]))
  
  MAX <- data.frame(matrix(nrow=s,ncol=9))
  names(MAX) = c("vari", "res", "file.name", "cmip", "mod", "modcl", "scen", "ens.mem", "year")
  YAX <- data.frame(matrix(nrow=s, ncol=1))
  names(YAX) <- "AGMT"
  
  XAX <- matrix(nrow=s,ncol=ncol)
  
  cc <- 0
  for (k in 1:length(X)){
    cat("\r ",k)
    Xsc <- t(X[[k]])
    # Ysc <- Y[[k]][seq(mon, dim(X[[k]])[3], 12),]
    
    for (pc in 1:(dim(X[[k]])[2])){
      cc <- cc+1
      MAX[cc,] <- as.character(c(M$vari[k], M$res[k], M$file.name[k], cmip, M$mod[k], M$modcl[k], M$scen[k], M$ens.mem[k], start.year-1+as.numeric(pc)))
      YAX[cc,] <- Xsc[pc,] %*% areaw
      XAX[cc,]  <- Xsc[pc,] 
    }
    print(MAX[cc,])
  }
  
  # save to .RData file:
  return(list(X = XAX, Y=YAX, M=MAX))
}


## Adjust and define years in CMIP6:
adjust.years.cmip6 <- function(XAX) {
  file.un = unique(XAX$M$file.name)
  for (i in 1:length(file.un)) {
    # print(i)
    ix = which(XAX$M$file.name == file.un[i])
    if (XAX$M$scen[ix[1]] == "historical") {
      XAX$M$year[ix] = 1850:2014
    } else if (XAX$M$scen[ix[1]] %in% c("hist-GHG", "hist-aer", "hist-nat")) {
      XAX$M$year[ix] = 1850:(1850-1+length(ix))
    } else if (XAX$M$scen[ix[1]] %in% c("ssp119", "ssp126", "ssp245", "ssp370", "ssp434", "ssp534", "ssp585")) {
      XAX$M$year[ix] = 2015:2100
      if (length(XAX$M$year[ix]) != 86) print(i)
    }
  }
  return(XAX)
}


## Center each ensemble member with historical period but over all realizations:
center.ensemble.member.cmip <- function(XAX_scen, ref.scen = "historical") {
  
  mod.phys.un = unique(XAX_scen$M$mod.phys)
  XAX_scen.out = XAX_scen
  mod.phys.bool = rep(NA, length(XAX_scen$M$mod.phys))
  
  # run through each mod.phys.un:
  for (m in 1:length(mod.phys.un)) {
    scen.un = unique(XAX_scen$M$scen[which(XAX_scen$M$mod.phys == mod.phys.un[m])])
    ens.mem.un = unique(XAX_scen$M$ens.mem[which(XAX_scen$M$mod.phys == mod.phys.un[m])])
    print(paste(m, mod.phys.un[m], sep=" "))
    print(scen.un)
    print(ens.mem.un)
    
    # for (em in 1:length(ens.mem.un)) {
      all.ix = which(XAX_scen$M$mod.phys == mod.phys.un[m])
      ref.ix = which(XAX_scen$M$mod.phys == mod.phys.un[m] & XAX_scen$M$scen %in% ref.scen & XAX_scen$M$year %in% 1870:1920)
      
      if (length(ref.ix) == 0) {
        mod.phys.bool[all.ix] = F
        ref.ix = which(XAX_scen$M$mod == substring(mod.phys.un[m], 1, nchar(mod.phys.un[m])-3)  & XAX_scen$M$scen %in% ref.scen & XAX_scen$M$year %in% 1870:1920)
        # if (length(ref.ix) == 0) ref.ix = which(XAX_scen$M$mod == substring(mod.phys.un[m], 1, nchar(mod.phys.un[m])-3) & XAX_scen$M$scen %in% ref.scen & XAX_scen$M$year %in% 1870:1920)
      } else {
        mod.phys.bool[all.ix] = T  # physics preserved in mean-subtraction
      }
      # Center based on ref. period:
      print(length(ref.ix))
      XAX_scen.out$Y$AGMT[all.ix] = XAX_scen$Y$AGMT[all.ix] - mean(XAX_scen$Y$AGMT[ref.ix])
      XAX_scen.out$tas[all.ix,] = XAX_scen$tas[all.ix,] - rep.row(x = colMeans(XAX_scen$tas[ref.ix,]), n = length(all.ix))
      # XAX_scen.out$psl[all.ix,] = XAX_scen$psl[all.ix,] - rep.row(x = colMeans(XAX_scen$psl[ref.ix,]), n = length(all.ix))
  }
  print(any(is.na(mod.phys.bool)))
  
  XAX_scen.out$M$mod.phys.bool = mod.phys.bool
  return(XAX_scen.out)
  
  # plot(cmip5_ann_piControl_5d00_XAX_ct$Y$AGMT[which(cmip5_ann_piControl_5d00_XAX_ct$M$mod == "GISS-E2-H" & cmip5_ann_piControl_5d00_XAX_ct$M$scen == "piControl_DT")])
  # unique(XAX_scen$M$mod)
  # unique(XAX_scen$M$mod[which(XAX_scen$M$scen == "historicalANT")])
  # unique(XAX_scen$M$mod[which(XAX_scen$M$scen == "historicalNat")])
  # unique(XAX_scen$M$mod[which(XAX_scen$M$scen == "historicalGHG")])
}


## Extract forced response as the ensemble means across >3 members:
extract.fraw2 <- function(XAX_scen, nmem = 3) {
  
  scen.un = unique(XAX_scen$M$scen)
  l = dim(XAX_scen$M)[1]
  
  # Extract "raw" forced response:
  XAX_scen$Y$ftot = rep(NA, l)
  XAX_scen$Y$ftot_l = rep(NA, l)
  XAX_scen$Y$IV.fraw.GMT.10y = rep(NA, l)
  XAX_scen$Y$IV.fraw.GMT.20y = rep(NA, l)
  XAX_scen$Y$IV.fraw.GMT.30y = rep(NA, l)
  XAX_scen$Y$IV.fraw.GMT.50y = rep(NA, l)
  
  
  for (s in 1:length(scen.un)) {
    print(paste(s, scen.un[s], sep=" "))
    cur.scen = scen.un[s]
    
    # Determine number of members per model:
    no.mem.mod = sapply(X = unique(XAX_scen$M$mod), FUN=function(x) length(unique(XAX_scen$M$ens.mem[which(x == XAX_scen$M$mod & XAX_scen$M$scen == cur.scen)])))
    no.mem = sapply(X = unique(XAX_scen$M$mod.phys), FUN=function(x) length(unique(XAX_scen$M$ens.mem[which(x == XAX_scen$M$mod.phys & XAX_scen$M$scen == cur.scen)])))
    # length(which(no.mem.mod >= 3))
    # length(which(no.mem >= 3))
    
    cur.mod.phys.un = names(which(no.mem >= nmem))
    cur.no.mem = no.mem[which(no.mem >= nmem)]

    for (m in 1:length(cur.mod.phys.un)) {
      print(paste(m, cur.mod.phys.un[m], sep=" "))
      
      # all scenarios:
      ix_scen = which(XAX_scen$M$mod.phys == cur.mod.phys.un[m] & XAX_scen$M$scen == cur.scen)
      ens.mat = sapply(X = unique(XAX_scen$M$ens.mem[ix_scen]), FUN=function(cur.ens) XAX_scen$Y$AGMT[which(XAX_scen$M$mod.phys == cur.mod.phys.un[m] & XAX_scen$M$scen == cur.scen & XAX_scen$M$ens.mem == cur.ens)])
      if (cur.no.mem[m] > 2) {
        ftot_IV = sapply(X = 1:dim(ens.mat)[2], FUN=function(k) rowMeans(ens.mat[,-k], na.rm=T))
      } else {
        ftot_IV = matrix(NA, nrow = dim(ens.mat)[1], ncol = dim(ens.mat)[2])
      }
      
      ftot = sapply(X = 1:dim(ens.mat)[2], FUN=function(k) rowMeans(ens.mat, na.rm=T))
      # ftot = sapply(X = 1:dim(ens.mat)[2], FUN=function(k) rowMeans(ens.mat, na.rm=T))
      # print(dim(ftot))
      # Fit LOWESS: fr.loess_0.75 = loess(fr.raw ~ c(1:231), span = 0.75, degree = 2)$fittedn=dim(ens.mat)[1]
      # 231 * 0.25 / n
      n = dim(ens.mat)[1]
      ftot_l = rep.col(loess(rowMeans(ftot) ~ c(1:n), span = 231 * 0.25 / n, degree = 2)$fitted, dim(ens.mat)[2])
      if (all(is.na(ftot[1,]))) ftot_l = rbind(rep(NA, dim(ens.mat)[2]), ftot_l)
      
      if (cur.scen %in% c("ssp119", "ssp126", "ssp245", "ssp370", "ssp434", "ssp585")) {
        ix_hist = c(which(XAX_scen$M$mod.phys == cur.mod.phys.un[m] & XAX_scen$M$scen == "historical"))
        ens.mat.hist = sapply(X = unique(XAX_scen$M$ens.mem[ix_hist]), FUN=function(cur.ens) XAX_scen$Y$AGMT[which(XAX_scen$M$mod.phys == cur.mod.phys.un[m] & XAX_scen$M$scen == "historical" & XAX_scen$M$ens.mem == cur.ens)])
        ftot_hist = c(sapply(X = 1:dim(ens.mat.hist)[2], FUN=function(k) rowMeans(ens.mat.hist, na.rm=T))[,1])
        ftot_long = c(ftot_hist, rowMeans(ftot))
        ftot_l = loess(ftot_long ~ c(1:(165+n)), span = 231 * 0.25 / n, degree = 2)$fitted
        if (is.na(ftot_long[1])) ftot_l = c(NA, ftot_l)
        # if (all(is.na(ftot[1,]))) ftot_l = rbind(rep(NA, dim(ens.mat)[2]), ftot_l)
        ftot_l = rep.col(ftot_l[-c(1:165)], dim(ens.mat)[2])
      }
      
      ## Fill in forced responses and smoothed versions:
      XAX_scen$Y$ftot[ix_scen] = c(ftot)
      XAX_scen$Y$ftot_l[ix_scen] = c(ftot_l)
      
      IV.fraw = ens.mat - ftot_IV
      XAX_scen$Y$IV.fraw.GMT.10y[ix_scen] = c(apply(X = IV.fraw, MARGIN=2, FUN=rollmean, k = 10, fill = NA))
      XAX_scen$Y$IV.fraw.GMT.20y[ix_scen] = c(apply(X = IV.fraw, MARGIN=2, FUN=rollmean, k = 20, fill = NA))
      XAX_scen$Y$IV.fraw.GMT.30y[ix_scen] = c(apply(X = IV.fraw, MARGIN=2, FUN=rollmean, k = 30, fill = NA))
      XAX_scen$Y$IV.fraw.GMT.50y[ix_scen] = c(apply(X = IV.fraw, MARGIN=2, FUN=rollmean, k = 50, fill = NA))
    }
  }
  
  return(XAX_scen)
}


# Reshape CMIP6 pi-Control files:
get.XAX_ann_piC <- function(X, ncol = 72*36, M, start.year = 1, center = T, scale = F, cmip) {
  
  # OVERALL DIMENSION
  s = sum(sapply(X = X, FUN=function(x) dim(x)[2]))
  
  MAX <- data.frame(matrix(nrow=s,ncol=9))
  names(MAX) = c("vari", "res", "file.name", "cmip", "mod", "modcl", "scen", "ens.mem", "year")
  YAX <- data.frame(matrix(nrow=s, ncol=5))
  names(YAX) <- c("AGMT", "GMT.10y", "GMT.20y", "GMT.30y", "GMT.50y")
  
  XAX <- matrix(nrow=s,ncol=ncol)
  
  cc <- 0
  for (k in 1:length(X)){
    cat("\r ",k)
    Xsc <- t(X[[k]])
    
    if (center == T & scale == F) {
      Xsc = scale(Xsc, T, F)
    } else if (scale == T & scale == T) {
      Xsc = scale(Xsc, T, T)
    }
    # Ysc <- Y[[k]][seq(mon, dim(X[[k]])[3], 12),]
    
    # do slow averages:
    AGMT = c(Xsc %*% areaw)
    AGMT.10y = rollmean(x = AGMT, k = 10, fill = NA)
    AGMT.20y = rollmean(x = AGMT, k = 20, fill = NA)
    AGMT.30y = rollmean(x = AGMT, k = 30, fill = NA)
    AGMT.50y = rollmean(x = AGMT, k = 50, fill = NA)
    
    for (pc in 1:(dim(X[[k]])[2])){
      cc <- cc+1
      MAX[cc,] <- as.character(c(M$vari[k], M$res[k], M$file.name[k], cmip, M$mod[k], M$modcl[k], M$scen[k], M$ens.mem[k], start.year-1+as.numeric(pc)))
      YAX[cc,1] <- AGMT[pc]
      YAX[cc,2] <- AGMT.10y[pc]
      YAX[cc,3] <- AGMT.20y[pc]
      YAX[cc,4] <- AGMT.30y[pc]
      YAX[cc,5] <- AGMT.50y[pc]
      XAX[cc,]  <- Xsc[pc,] 
    }
    print(MAX[cc,])
  }
  
  # save to .RData file:
  return(list(X = XAX, Y=YAX, M=MAX))
}