Sys.setenv(LANG = "en")
Sys.setlocale("LC_TIME","en_GB")
library(startR)
library(s2dverification)
library(abind)
library(CSTools)
library(zeallot)
library(s2dv)
library(easyVerification)
library(ClimProjDiags)

#First of all we define the initialisation date and the variable
variable<-"tas"
fcst.sdate<-20030806


#-------------------------------------------------------------------
# Settings
#-------------------------------------------------------------------
 
#data directories
fcst.dir <- "/esarchive/exp/ncep/cfs-v2/"
hcst.dir <- "/esarchive/exp/ncep/cfs-v2/"
obs.dir  <- "/esarchive/recon/ecmwf/era5/"

##Time aggregation in esarchive
fcst.freq <- "weekly_mean"

# Even though we will deal with weekly aggregation, 
# in /esarchive it is specified the frequency of the original
# data before being aggregated. Here we define these frequencies for 
# the several variables contained at esarchive

  switch (tolower(variable),
      "tas"     = {freq.hcst <- 'f24h';   freq.obs <- 'f1h'},
      "tasmin"  = {freq.hcst <- 'f24h';   freq.obs <- 'f24h'},
      "tasmax"  = {freq.hcst <- 'f24h';   freq.obs <- 'f24h'},
      "psl"     = {freq.hcst <- 'f24h';   freq.obs <- 'f1h'},
      "prlr"    = {freq.hcst <- 'f6h';  freq.obs <- 'f1h'; accum <- TRUE},
      "rsds"    = {freq.hcst <- 's0-24h'; freq.obs <- 'f1h'; accum <- TRUE},
      "sfcwind" = {freq.hcst <- 'f24h';   freq.obs <- 'f1h'}
  )

# grid of the observations. We are taking here
# observations that are already regrided and archived in esarchive.
# You can select the non regrided "-r1440x721cds" 
obs.grid <- "-r384x190cds"
#only for psl
#obs.grid <- "-384x190"


# Let's define the path for each of the files that will be retrieved
# for the observations, hindcast and forecast.
# "var" and  "file_date" will be specified when doing the start call

obs.path  <- paste0(
               obs.dir,fcst.freq,"/$var$_",
               freq.obs,obs.grid,"/$file_date$.nc")

hcst.path <- paste0(
               hcst.dir,fcst.freq,"/s2s","/$var$_",
               freq.hcst,"/$file_date$",
               ".nc")

fcst.path <- paste0(
               fcst.dir,fcst.freq,"/s2s", "/$var$_",
               freq.hcst,"/$file_date$",
               ".nc")


#region limits
lons.min <- 5 
lons.max <- 10 
lats.min <- 5  
lats.max <- 10 

#number of members
hcst.nmember <- 12
fcst.nmember <- 48

# lead times to retrieve  
ltmin <- 1
ltmax <- 4

#define the reference period
# normal is 1999 to 2020
hcst_start_year <- 2015
hcst_end_year <- 2020
#the next two parameters are used to define the arrays containing the dates
hcst.ltmin <- as.numeric(substr(as.character(fcst.sdate),1,4)) - hcst_end_year
hcst.ltmax <- as.numeric(substr(as.character(fcst.sdate),1,4)) - hcst_start_year


# let's select the number of weeks 
# used to calculate the skill. It does not affect calibration, it would be only used
# to make larger the sample and calculate the skill. Let's do it 1 for now
n.skill.weeks <- 9

# We define now the function to obtain the array of the dates that will be retrieved for the hindcast
# This function retrieves also a window containing the last and next monday (dimension "sday")
# to make larger the sample when calibrating. However we will select just the forecast initialisation date 

get.file_dates <- function(startdate, var){

  # This function generates the sdates path to retrieve the hcst
  # files. Outputs a multi-dimensional array with dims:
  # (sweek,sday, syear) containing the character string needed to
  # access each file. Ex: "20201001/tas_20191001.nc"

  sdates <- numeric(0)
  startdate <- as.Date(toString(startdate), "%Y%m%d")
  while (length(sdates) < n.skill.weeks){
    if (format(startdate, "%a") == "Thu" || format(startdate, "%a") == "Mon") {
      sdates <- c(sdates,format(startdate, "%Y%m%d"))
    }
    startdate <- startdate - 1
  }

  file_dates <- array(numeric(),c(0,3,(hcst.ltmax-hcst.ltmin)+1))
  file_dates.fcst <- array(numeric(),c(0,3,(hcst.ltmax-hcst.ltmin)+1))
  for (sdate in sdates) {
    ## TODO: take into account year of the fcst to select hcst years
    fcst.day <- substr(sdate,7,8)
    fcst.month <- substr(sdate,5,6)
    fcst.year <- substr(sdate,1,4)

    startdate <- as.Date(toString(sdate), "%Y%m%d")

    # makes sure of giving a: [M T M || T M T] combination
    if(format(startdate, "%a") == "Thu"){
      hcst.sdays <- c(startdate - 3, startdate, startdate + 4, recursive=T)
    } else {
      hcst.sdays <- c(startdate - 4, startdate, startdate + 3, recursive=T)
    }

    hcst.sdays <- format(as.Date(hcst.sdays,"%Y%m%d"), "%Y%m%d")

    sdays.file_dates <- array(numeric(),c(0,(hcst.ltmax-hcst.ltmin)+1))
    sdays.file_dates.fcst <- array(hcst.sdays,c(3,(hcst.ltmax-hcst.ltmin)+1))


    for (sday.sdate in hcst.sdays) {
      sday.year <- substr(sday.sdate, 1,4)
      sday.mday <- substr(sday.sdate, 5,8)
      sday.years <-
        paste((strtoi(sday.year)-hcst.ltmax):(strtoi(sday.year)-hcst.ltmin), sep = "")
        sday.dates <- apply(expand.grid(sday.years, sday.mday), 1, paste, collapse="")
        sdays.file_dates <- abind(sdays.file_dates, sday.dates,
        along=1)
    }

    names(dim(sdays.file_dates)) <- c('sday','syear')

    file_dates <- abind(file_dates, sdays.file_dates, along=1)
    file_dates.fcst <- abind(file_dates.fcst, sdays.file_dates.fcst, along=1)

  }

  dir_dates <- paste0(file_dates.fcst,"/",var,"_",file_dates)
  dim(dir_dates) <- dim(file_dates)
  names(dim(dir_dates)) <- c('sweek','sday','syear')

  dir_dates <- substr(dir_dates,10,24)
 return(dir_dates)

}

if(n.skill.weeks %% 2 != 0) {
  fcst_sweek_ind <- (n.skill.weeks+1)/2
} else{
  fcst_sweek_ind <- (n.skill.weeks)/2
}

if(fcst_sweek_ind != 1) {
  if(fcst_sweek_ind %% 2 == 0) {
     fcst.sdate_to_start<-as.character(format(as.Date(as.character(fcst.sdate),format = "%Y%m%d")+4+7*((fcst_sweek_ind/2)-1),"%Y%m%d"))
     file_dates <- get.file_dates(fcst.sdate_to_start,variable)
  } else{
    fcst.sdate_to_start<-as.character(format(as.Date(as.character(fcst.sdate),format = "%Y%m%d")+7*((fcst_sweek_ind-1)/2),"%Y%m%d"))
    file_dates <- get.file_dates(fcst.sdate_to_start,variable)
  }
} else {
  file_dates <- get.file_dates(fcst.sdate,variable)  
}

#Let's select the central day of the window (corresponding to the forecast initialisation date)  
#file_dates<-Subset(file_dates, along=c("sday"),indices=c(2),drop=FALSE)

# Don't do the subset and take the 3 days, not just the middle one 
#file_dates <- Subset(file_dates, along=c("sweek","sday"),indices=list(n.skill.weeks:1,2),drop=FALSE)

#Let's give format the file containing the date for the forcast
file_dates.fcst <- paste0(variable, '_',fcst.sdate)
dim(file_dates.fcst) <- dim(Subset(file_dates, c('sweek', 'sday', 'syear'),
                                       list(1,1,1)))




# We check now if the first year of the hindcast is available or not
    hcst.dir1 <- paste0(
               hcst.dir,fcst.freq,
               "/s2s/",variable,"_", freq.hcst,
               "/",
               paste0(file_dates[1,1,]),".nc") #syears

    first.exists <- FALSE
    n <- 1
    while(!first.exists){
    if(!file.exists(hcst.dir1[n])){
      n <- n + 1
    } else {
      first.exists <- TRUE
      if (n > 1){
        hcst.ltmax <- hcst.ltmax + (1-n)
        warning(paste0("NEW HCST INIT YEAR: ",
                               as.numeric(substr(fcst.sdate,1,4))-hcst.ltmax))

    }}}



#Let's proceed with the loading of the data

    #fcst

    fcst <- Start(dat = fcst.path,
                  var = variable,
                  file_date = file_dates.fcst,
                  time = indices(ltmin:ltmax),
                        latitude = values(list(lats.min, lats.max)),
                        latitude_reorder = Sort(decreasing = TRUE),
                        longitude = values(list(lons.min, lons.max)),
                        longitude_reorder = CircularSort(-180, 180),
                  synonims = list(latitude=c('lat','latitude'),
                                  longitude=c('lon','longitude'),
                                  ensemble=c('member','ensemble')),
                  ensemble = indices(1:fcst.nmember),
                  return_vars = list(lat = 'dat',
                                     lon = 'dat',
                                     time = 'file_date'),
                  split_multiselected_dims = TRUE,
                  retrieve = T)

      #hcst

      hcst<-Start(dat = hcst.path,
              var = variable,
              file_date = file_dates,
              time = indices(ltmin:ltmax),
                        latitude = values(list(lats.min, lats.max)),
                        latitude_reorder = Sort(decreasing = TRUE),
                        longitude = values(list(lons.min, lons.max)),
                        longitude_reorder = CircularSort(-180, 180),#
              synonims = list(latitude=c('lat','latitude'),
                              longitude=c('lon','longitude'),
                              ensemble=c('member','ensemble')),
              ensemble = indices(1:hcst.nmember),
              return_vars = list(latitude = 'dat',
                                 longitude = 'dat',
                                 time = c('file_date')),
              split_multiselected_dims = TRUE,
              num_procs = 1,
              retrieve = T)


# We now extract the dates from the hcst to retrieve the
# observations corresponding to the lead times of the hindcast

  dates <- attr(hcst, 'Variables')$common$time
  dates_file <- sapply(dates, format, '%Y%m%d')
  dates_file <- format(as.Date(dates_file, '%Y%m%d')-3, "%Y%m%d")
  dates_file <- paste0(variable,"_",dates_file)
  dim(dates_file) <- dim(Subset(hcst,
                                along=c('dat','var',
                                'latitude', 'longitude', 'ensemble'),
								list(1,1,1,1,1), drop='selected'))

  # Notice that dates_file is not the same as file_dates.
  # file_dates is the array containing the strings of the initialisations for the hindcast,
  # while dates_file contains the hindcast times  

 #let's obtain the observations 
 obs <- Start(dat = obs.path,
                var = variable,
                file_date = dates_file,
                latitude = values(list(lats.min, lats.max)),
                latitude_reorder = Sort(decreasing = TRUE),
                longitude = values(list(lons.min, lons.max)),
                longitude_reorder = CircularSort(-180, 180),
                synonims = list(latitude=c('lat','latitude'),
                                longitude=c('lon','longitude')),
                return_vars = list(latitude = 'dat',
                                  longitude = 'dat',
                                  time = 'file_date'),
               split_multiselected_dims = TRUE,
                retrieve = TRUE)


obs <- as.s2dv_cube(obs)
hcst <- as.s2dv_cube(hcst)
#data$hcst$data <-multiApply::Apply(data$hcst$data,target_dims=c("sweek","syear"),fun=as.numeric,output_dims="syear")$output1
#data$obs$data <-multiApply::Apply(data$obs$data,target_dims=c("sweek","syear"),fun=as.numeric,output_dims="syear")$output1

# make sure that sday = 3 
