library(startR)
# Detect data for 1 grid point:
ini <- as.Date('1979-01-01', '%Y-%m-%d')
end <- as.Date('2018-12-31', '%Y-%m-%d')
dates <- format(seq(ini, end, 'month'), '%Y%m')

hrmin <- Start(dat = "/esarchive/recon/ecmwf/era5/daily/$var$-r1440x721cds/$var$_$date$.nc",
               var = "hursmin",
               date = as.character(dates),
               time = 'all',
               lat = indices(1),
               lon = indices(1),
               time_across = 'date',
               merge_across_dims = TRUE,
               merge_across_dims_narm = TRUE,
               num_procs = 2,
               retrieve = TRUE)

times <- 1:(31*12*40)  
times <- times[-which(is.na(hrmin))]
length(day <- seq(ini, end, 'day')) == length(times)

hrmin <- Start(dat = "/esarchive/recon/ecmwf/era5/daily/$var$-r1440x721cds/$var$_$date$.nc",
               var = "hursmin",
               date = as.character(dates),
               time = indices(times),
               lat = indices(1),
               lon = indices(1),
               time_across = 'date',
               merge_across_dims = TRUE,
               merge_across_dims_narm = TRUE,
               return_vars = list(time = NULL, lat = NULL, lon = NULL),
               num_procs = 2,
               retrieve = TRUE)

tmax <- Start(dat = "/esarchive/recon/ecmwf/era5/daily/$var$-r1440x721cds/$var$_$date$.nc",
               var = "tasmax",
               date = as.character(dates),
               time = indices(times),
               lat = indices(1),
               lon = indices(1),
               time_across = 'date',
               merge_across_dims = TRUE,
               merge_across_dims_narm = TRUE,
               synonims = list(lat = c('lat', 'latitude'),
                               lon = c('lon', 'longitude')),
               return_vars = list(time = NULL, lat = NULL, lon = NULL),
               num_procs = 2,
               retrieve = TRUE)
wind <- Start(dat = "/esarchive/recon/ecmwf/era5/daily_mean/$var$_f1h-r1440x721cds/$var$_$date$.nc",
               var = "sfcWind",
               date = as.character(dates),
               time = indices(times),
               lat = indices(1),
               lon = indices(1),
               time_across = 'date',
               merge_across_dims = TRUE,
               merge_across_dims_narm = TRUE,
               return_vars = list(time = NULL, lat = NULL, lon = NULL),
               num_procs = 2,
               retrieve = TRUE)
ppt <- Start(dat = "/esarchive/recon/ecmwf/era5/daily_mean/$var$_f1h-r1440x721cds/$var$_$date$.nc",
               var = "prlr",
               date = as.character(dates),
               time = indices(times),
               lat = indices(1),
               lon = indices(1),
               time_across = 'date',
               merge_across_dims = TRUE,
               merge_across_dims_narm = TRUE,
               return_vars = list(time = NULL, lat = NULL, lon = NULL),
               num_procs = 2,
               retrieve = TRUE)

# Check metadata:
dim(hrmin)
dim(tmax)
dim(wind)
dim(ppt)
length(times)
dim(attributes(hrmin)$Variable$common$time)
dim(attributes(ppt)$Variable$common$time)
dim(attributes(wind)$Variable$common$time)
dim(attributes(tmax)$Variable$common$time)
## time:
all(as.Date(attributes(hrmin)$Variable$common$time) == 
    as.Date(attributes(ppt)$Variable$common$time), na.rm = T)
all(as.Date(attributes(tmax)$Variable$common$time) == 
    as.Date(attributes(wind)$Variable$common$time), na.rm = T)
all(as.Date(attributes(tmax)$Variable$common$time) == 
    as.Date(attributes(ppt)$Variable$common$time), na.rm = T)
## lat
attributes(hrmin)$Variable$common$lat == attributes(ppt)$Variable$common$lat
attributes(tmax)$Variable$common$lat == attributes(wind)$Variable$common$lat
attributes(hrmin)$Variable$common$lat == attributes(tmax)$Variable$common$lat
## lon
attributes(hrmin)$Variable$common$lon == attributes(ppt)$Variable$common$lon
attributes(tmax)$Variable$common$lon == attributes(wind)$Variable$common$lon
attributes(hrmin)$Variable$common$lon == attributes(tmax)$Variable$common$lon

sum(is.na(hrmin))
sum(is.na(ppt))
sum(is.na(wind))
sum(is.na(tmax))



FWI <- function(temp, rh, ws, prec, date, lat, lon) {
  # dims: 
  # units:  
    # prec: m/s to mm -> * 24*60*60*1000
  prec <- as.vector(prec) * 24 * 60 * 60 * 1000
    # tmax: K to Celsius -> - 272.15
  temp <- as.vector(temp) - 272.15
    # hrmin: % 
  rh <- as.vector(rh)
    # wind: m/s to k/h -> * 3.6
  ws <- as.vector(ws) * 3.6
  res  <- data.frame(temp = temp, rh = rh, ws = ws, prec = prec,
                     lat = lat, long = lon,
                     yr = as.numeric(format(date, "%Y")),
                     mon = as.numeric(format(date, "%m")),
                     day = as.numeric(format(date, "%d")))
  res <- cffdrs::fwi(input = res)
  return(res$FWI)
}
library(tictoc)
tic()
res <-  FWI(tmax, hrmin, wind, ppt, date = day, lat = 90, lon = 0)
toc()


step <- Step(fun = temp_mean,
             target_dims = c('time', 'year'),
             output_dims = NULL)

wf <- AddStep(hrmin, step)

queue_host = 'nord4'   #your own host name for power9
temp_dir =  '/gpfs/scratch/bsc32/bsc32339/startR_hpc/'
ecflow_suite_dir = '/home/Earth/nperez/startR_local/'  #your own local directory

res <- Compute(wf,
              chunks = list(month = 12),
              cluster = list(queue_host = queue_host,
                             queue_type = 'slurm',
                             cores_per_job = 16,
                             temp_dir = temp_dir,
                             polling_period = 10,
                             job_wallclock = '03:00:00',
                             max_jobs = 12,
                             bidirectional = FALSE),
              ecflow_suite_dir = ecflow_suite_dir,
              wait = FALSE)


