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,
               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 = 'all',
               lon = 'all',
               time_across = 'date',
               merge_across_dims = TRUE,
               merge_across_dims_narm = TRUE,
               return_vars = list(time = NULL, lat = NULL, lon = NULL),
               retrieve = FALSE)

tmax <- Start(dat = "/esarchive/recon/ecmwf/era5/daily/$var$-r1440x721cds/$var$_$date$.nc",
               var = "tasmax",
               date = as.character(dates),
               time = indices(times),
               lat = 'all',
               lon = 'all',
               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),
               retrieve = FALSE)
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 = 'all',
               lon = 'all',
               time_across = 'date',
               merge_across_dims = TRUE,
               merge_across_dims_narm = TRUE,
               return_vars = list(time = NULL, lat = NULL, lon = NULL),
              retrieve = FALSE)
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 = 'all',
               lon = 'all',
               time_across = 'date',
               merge_across_dims = TRUE,
               merge_across_dims_narm = TRUE,
               return_vars = list(time = NULL, lat = NULL, lon = NULL),
               retrieve = FALSE)

raw_fwi <- Start(dat = "/esarchive/recon/ecmwf/era5/daily/$var$/$var$_$date$.nc",
                var = "fwiNoOverwintering",
                year = as.character(1979:2018),
                month = sprintf("%02d", 1:12),
                time = 'all',
                lat = 'all',
                lon = 'all',
                synonims = list(lat = c('lat', 'latitude'),
                                lon = c('lon', 'longitude')),
                return_vars = list(lat = NULL, lon = NULL),
                retrieve = FALSE)


# 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
all(attributes(hrmin)$Variable$common$lat == attributes(ppt)$Variable$common$lat)
all(attributes(tmax)$Variable$common$lat == attributes(wind)$Variable$common$lat)
all(attributes(hrmin)$Variable$common$lat == attributes(tmax)$Variable$common$lat)
## lon
all(attributes(hrmin)$Variable$common$lon == attributes(ppt)$Variable$common$lon)
all(attributes(tmax)$Variable$common$lon == attributes(wind)$Variable$common$lon)
all(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 =  NULL, lats, lons, chunk = chunk_indices_apply) {
  # dims: time x lat 
  # units:  
    # prec: m/s to mm -> * 24*60*60*1000ç
  startR:::.warning(chunk)
  startR:::.warning(dim(temp))
  prec <- prec * 24 * 60 * 60 * 1000
    # tmax: K to Celsius -> - 272.15
  temp <- temp - 272.15
    # hrmin: % 
  rh <- rh
    # wind: m/s to k/h -> * 3.6
  ws <- ws * 3.6
  #lat = attributes(temp)$Variables$common$lat
  #lon = attributes(temp)$Variables$common$lon
  result <- array(numeric(), c(time = length(date), lat = length(lats)))
  startR:::.warning(dim(result))
  longitude <- lons[chunk]
  startR:::.warning(longitude)
  for (x in 1:length(lats)) {
    startR:::.warning(x)
    res  <- data.frame(temp = temp[,x], rh = rh[,x], ws = ws[,x], prec = prec[,x],
                       lat = lats[x], long = longitude,
                       yr = as.numeric(format(date, "%Y")),
                       mon = as.numeric(format(date, "%m")),
                       day = as.numeric(format(date, "%d")))
start_time <- Sys.time()
    result[, x] <- cffdrs::fwi(input = res, out = 'fwi')$FWI
end_time <- Sys.time()
startR:::.warning(end_time - start_time)
  }
  return(result)
}

step <- Step(fun = FWI,
             target_dims = list(c('time', 'lat'), 
                                c('time', 'lat'), 
                                c('time', 'lat'), 
                                c('time', 'lat')), 
             output_dims = c('time', 'lat'))

wf <- AddStep(list(tmax, hrmin, wind, ppt), step, date = day,
              lats = attributes(hrmin)$Variable$common$lat,
              lons = attributes(hrmin)$Variable$common$lon)

queue_host = 'nord2'   #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

res2 <- Compute(wf,
              chunks = list(lon = 100),
              threads_load = 1, threads_compute = 16,
              cluster = list(queue_host = queue_host,
                             queue_type = 'lsf',
                             cores_per_job = 16,
                             temp_dir = temp_dir,
                             polling_period = 10,
                             job_wallclock = '03:00',
                             max_jobs = 12,
                             bidirectional = FALSE),
              ecflow_suite_dir = ecflow_suite_dir,
              wait = FALSE)

result <- Collect(res, wait = TRUE)
saveRDS(result,
        file = "/esarchive/scratch/nperez/git/Flor/Landmarc/FWI/2lons_FWI.RDS")
matplot(1:14610, result$output1[,,1,1,1], type = "l")

apply(result$outpu1, 5, 
image(1:14610, sort(as.vector(attributes(hrmin)$Variable$common$lat),
      result$output1[,,1,1,1])
           
