## -----------------------------------
## Function for CMIP5 extraction:
## -----------------------------------

# Sebastian Sippel
# 02.09.2021


## Function to get list of CMIP5 files:
get.CMIP5.file.list <- function(vari= "tas", temp.res = "ann", scen = c("rcp85", "rcp45", "rcp26", "piControl"), piControl.length=200, grid = "g025", CMIP5.dir = "/net/atmos/data/cmip5-ng/") {
  
  # 1. get  list for each scenario:
  scen.list = list()
  for (scen.idx in 1:length(scen)) {
    scen.list[[scen.idx]] = grep(pattern = scen[scen.idx], grep(pattern = temp.res, x = list.files(path = paste(CMIP5.dir, sep=""), pattern = grid), value = T), value = T)
    if (scen[scen.idx]=="piControl") {
      piControlLength=sapply(X = scen.list[[scen.idx]], FUN=function(x) get_piControlLength_CMIP5(file=paste(CMIP5.dir, "/", x, sep=""), var = vari, temp.res = temp.res))
      scen.list[[scen.idx]]=scen.list[[scen.idx]][which(piControlLength>(piControl.length-1))]
    }
  }
  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[1] == "piControl") {
    period.length = piControlLength[which(piControlLength>piControl.length)]
  } else {
    period.length = NA
  }
  
  return(data.frame(file.name, vari, res, mod, modcl, scen, ens.mem, modall, period.length, stringsAsFactors=F))
}


## Read CMIP5 files: 
get_CMIP5_array <- function(file, var="tas", time.count=-1) {
  ncfile=nc_open(file)
  ncdata=ncvar_get(nc = ncfile, varid=var, start=c(1,1,1), count=c(-1,-1,time.count))
  nc_close(ncfile)
  return(ncdata)
}


## Read CMIP5 files:
get_CMIP5_vec <- function(file, var="tas", time.count=-1) {
  ncfile=nc_open(file)
  ncdata=ncvar_get(nc = ncfile, varid=var, start=c(1), count=c(time.count))
  nc_close(ncfile)
  return(ncdata)
}


## Read CMIP5 files:
read.CMIP5_novar <- function(file.name, var, res="ann", scen, CMIP5.dir) {
  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;
      X[[i]] = get_CMIP5_array(file = paste(CMIP5.dir, "/", file.name[i], sep=""), time.count = -1, var=var)
    } else {
      X[[i]] = get_CMIP5_array(file = paste(CMIP5.dir, "/", file.name[i], sep=""), var=var)
    }
  }
  return(X)
}


## get length of piControl run in CMIP5
get_piControlLength_CMIP5 <- function(file, var="tas", temp.res="ann") {
  ncfile=nc_open(file)
  ncdata=length(ncvar_get(nc = ncfile, varid=var, start = c(1,1,1), count=c(1,1,-1)))
  # if (grep(pattern = "mon", x = file, value=F)==1, value=T) ncdata=ncdata/12
  if (temp.res=="mon") ncdata=ncdata/12
  nc_close(ncfile)
  return(ncdata)
}