# Code to compute a climdex.pcic indicator in Landmarc project
# Author: Núria Pérez-Zanón
# Date: 02-11-2021

library(startR)

# Piece of data for testing:
seq.dates <- seq(as.POSIXct('1850-01-01 12:00:00 UTC', tz = 'UTC'),
                 as.POSIXct('2010-12-31 12:00:00 UTC', tz = 'UTC'), 'day')

hist <- Start(dat = "/esarchive/exp/CMIP6/historical/cnrm-cm6-1/cmip6-historical_i0p1/CMIP/CNRM-CERFACS/CNRM-CM6-1/historical/$member$/day/$var$/gr/v20200101/$var$_day_CNRM-CM6-1_historical_$member$_gr_$period$.nc",
              var = 'tas',
              member = 'r1i1p1f2',
              period = '18500101-20141231',
              lat = 'all',
              lon = 'all',
              time = indices(11324:12053),
                     #values(seq.dates),
                     #values(list(as.POSIXct('1981-01-01 12:00:00 UTC', tz = 'UTC'),
                      #           as.POSIXct('1982-12-31 12:00:00 UTC', tz = 'UTC'))),
              transform = CDORemapper,
              transform_extra_cells = 2,
              transform_params = list(grid = 'r180x90',
                                     method = 'conservative',
                                     crop = FALSE),
              transform_vars = c('lat', 'lon'),
              synonims = list(lon = c('lon', 'longitude'),
                              lat = c('lat', 'latitude')),
              return_vars = list(time = NULL,
                                 lat = 'dat',
                                 lon = 'dat'),
              retrieve =  FALSE)

lon <- attributes(hist)$Varia$dat1$lon
lat <- attributes(hist)$Varia$dat1$lat
time <- attributes(hist)$Variable$common$time

climdex_ind <- function(dat, times) {
  datefactor <- as.numeric(substr(times, 1, 4))
  datefactor <- as.factor(datefactor)
  fd <- climdex.pcic::number.days.op.threshold(dat, datefactor, 273.15, "<")
  return(fd)
}

step <- Step(fun = climdex_ind,
               target_dims = 'time',
               output_dims = 'time')
  
wf <- AddStep(hist, step, times = time)
  
## locally
res <- Compute(wf,
               chunks = list(lat = 2,
                             lon = 2))

s2dv::PlotEquiMap(res$output1[1,1,1,1,1,,], lon = lon, lat = lat, filled.con = F)


# Larger analysis:

hist <- Start(dat = "/esarchive/exp/CMIP6/historical/cnrm-cm6-1/cmip6-historical_i0p1/CMIP/CNRM-CERFACS/CNRM-CM6-1/historical/$member$/day/$var$/gr/v20200101/$var$_day_CNRM-CM6-1_historical_$member$_gr_$period$.nc",
              var = 'tas',
              member = 'r1i1p1f2',
              period = '18500101-20141231',
              lat = 'all',
              lon = 'all',
              time = indices(47848:58804),
                      #values(list(as.POSIXct('1981-01-01 12:00:00 UTC', tz = 'UTC'),
                       #          as.POSIXct('2010-12-31 12:00:00 UTC', tz = 'UTC'))),
              transform = CDORemapper,
              transform_extra_cells = 2,
              transform_params = list(grid = 'r180x90',
                                     method = 'conservative',
                                     crop = FALSE),
              transform_vars = c('lat', 'lon'),
              synonims = list(lon = c('lon', 'longitude'),
                              lat = c('lat', 'latitude')),
              return_vars = list(time = NULL,
                                 lat = 'dat',
                                 lon = 'dat'),
              retrieve = FALSE)
lon <- attributes(hist)$Varia$dat1$lon
lat <- attributes(hist)$Varia$dat1$lat
time <- attributes(hist)$Variable$common$time


wf <- AddStep(hist, step, times = time)

# Run Nord3

  queue_host = 'nord3'   #your own host name for nord3
  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(lat = 2,
                               lon = 2),
                 threads_load = 1,
                 threads_compute = 8,
                 cluster = list(queue_host = queue_host,
                                queue_type = 'lsf',
                                extra_queue_params = list('#BSUB -q bsc_es'),
                                cores_per_job = 8,
                                temp_dir = temp_dir,
                                polling_period = 10,
                                job_wallclock = '02:00',
                                max_jobs = 8,
                                bidirectional = FALSE),
                 ecflow_suite_dir = ecflow_suite_dir,
                 wait = TRUE)


s2dv::PlotEquiMap(res$output1[1,1,1,1,1,,], lon = lon, lat = lat, filled.con = F, brks = seq(0:366, 6))


# Test loading 3 variables to create an climdex object:
ini_date <- as.POSIXct('1981-01-01 12:00:00 UTC', tz = 'UTC')
end_date <- as.POSIXct('2010-12-31 12:00:00 UTC', tz = 'UTC')

file_dates <- seq(as.POSIXct('1850-01-01 12:00:00 UTC', tz = 'UTC'),
                  as.POSIXct('2010-12-31 12:00:00 UTC', tz = 'UTC'), 'day')
pos_ini <- which(file_dates == ini_date)
pos_end <- which(file_dates == end_date)

hist_tx <- Start(dat = "/esarchive/exp/CMIP6/historical/cnrm-cm6-1/cmip6-historical_i0p1/CMIP/CNRM-CERFACS/CNRM-CM6-1/historical/$member$/day/$var$/gr/v20200101/$var$_day_CNRM-CM6-1_historical_$member$_gr_$period$.nc",
              var = 'tas',
              member = 'r1i1p1f2',
              period = '18500101-20141231',
              lat = 'all',
              lon = 'all',
              time = indices(pos_ini:pos_end),
              transform = CDORemapper,
              transform_extra_cells = 2,
              transform_params = list(grid = 'r180x90',
                                     method = 'conservative',
                                     crop = FALSE),
              transform_vars = c('lat', 'lon'),
              synonims = list(lon = c('lon', 'longitude'),
                              lat = c('lat', 'latitude')),
              return_vars = list(time = NULL,
                                 lat = 'dat',
                                 lon = 'dat'),
              retrieve =  FALSE)
lon <- attributes(hist_tx)$Varia$dat1$lon
lat <- attributes(hist_tx)$Varia$dat1$lat
time <- attributes(hist_tx)$Variable$common$time

# -----
file_dates <- seq(as.POSIXct('1950-01-01 12:00:00 UTC', tz = 'UTC'),
                  as.POSIXct('2010-12-31 12:00:00 UTC', tz = 'UTC'), 'day')
pos_ini <- which(file_dates == ini_date)
pos_end <- which(file_dates == end_date)


hist_pr <- Start(dat = "/esarchive/exp/CMIP6/historical/cnrm-cm6-1/cmip6-historical_i0p1/CMIP/CNRM-CERFACS/CNRM-CM6-1/historical/$member$/day/$var$/gr/v20200101/$var$_day_CNRM-CM6-1_historical_$member$_gr_$period$.nc",
              var = 'pr',
              member = 'r1i1p1f2',
              period = '19500101-20141231',
              lat = 'all',
              lon = 'all',
              time = indices(pos_ini:pos_end),
              #values(list(as.POSIXct('1981-01-01 12:00:00 UTC', tz = 'UTC'),
              #           as.POSIXct('1982-12-31 12:00:00 UTC', tz = 'UTC'))),
              transform = CDORemapper,
              transform_extra_cells = 2,
              transform_params = list(grid = 'r180x90',
                                     method = 'conservative',
                                     crop = FALSE),
              transform_vars = c('lat', 'lon'),
              synonims = list(lon = c('lon', 'longitude'),
                              lat = c('lat', 'latitude')),
              return_vars = list(time = NULL,
                                 lat = 'dat',
                                 lon = 'dat'),
              retrieve =  FALSE)
lon_pr <- attributes(hist_pr)$Varia$dat1$lon
lat_pr <- attributes(hist_pr)$Varia$dat1$lat
time_pr <- attributes(hist_pr)$Variable$common$time

#-----
file_dates <- seq(as.POSIXct('1850-01-01 12:00:00 UTC', tz = 'UTC'),
                  as.POSIXct('2010-12-31 12:00:00 UTC', tz = 'UTC'), 'day')
pos_ini <- which(file_dates == ini_date)
pos_end <- which(file_dates == end_date)

hist_tn <- Start(dat = "/esarchive/exp/CMIP6/historical/cnrm-cm6-1/cmip6-historical_i0p1/CMIP/CNRM-CERFACS/CNRM-CM6-1/historical/$member$/day/$var$/gr/v20200101/$var$_day_CNRM-CM6-1_historical_$member$_gr_$period$.nc",
              var = 'tas',
              member = 'r1i1p1f2',
              period = '18500101-20141231',
              lat = 'all',
              lon = 'all',
              time = indices(pos_ini:pos_end),
              transform = CDORemapper,
              transform_extra_cells = 2,
              transform_params = list(grid = 'r180x90',
                                     method = 'conservative',
                                     crop = FALSE),
              transform_vars = c('lat', 'lon'),
              synonims = list(lon = c('lon', 'longitude'),
                              lat = c('lat', 'latitude')),
              return_vars = list(time = NULL,
                                 lat = 'dat',
                                 lon = 'dat'),
              retrieve =  FALSE)
lon_tn <- attributes(hist_tn)$Varia$dat1$lon
lat_tn <- attributes(hist_tn)$Varia$dat1$lat
time_tn <- attributes(hist_tn)$Variable$common$time

# Checks:
time[1]
time_tn[1]
time_pr[1]
time_pr[length(time_pr)]
time_tn[length(time_tn)]
time[length(time)]
all(lat == lat_tn)
all(lat == lat_pr)
all(lon == lon_tn)
all(lon == lon_pr)

climdex_ind <- function(tx, tn, pr, times, lats) {
  library(climdex.pcic)
  times <- as.PCICt(times, cal = 'standard')
  fd <- array(numeric(), c(lat = length(lats), time = 30))
  cdd <- array(numeric(), c(lat = length(lats), time = 30))
print(dim(tx))
print(length(lats))
  for (pos in 1:length(lats)) {
    hemisphere <- ifelse(lats[pos] > 0, TRUE, FALSE)
print(pos)
    data_obj <- climdexInput.raw(tmax = tx[pos,] - 273.15,
                                 tmin = tn[pos,] - 273.15, 
                                 prec = pr[pos,], tmax.dates = times,
                                 tmin.dates = times,
                                 prec.dates = times,
                                 base.range = c(1991, 2010),
                                 n = 5, 
                                 northern.hemisphere = hemisphere)
print(length(climdex.cdd(data_obj)))
    fd[pos, ] <- climdex.fd(data_obj)
    cdd[pos, ] <- climdex.cdd(data_obj)
print(length(climdex.cdd(data_obj)))
    print(sum(is.na(fd[pos,]))) 
  }
  return(list(fd = fd, cdd = cdd))
}

step <- Step(fun = climdex_ind,
               target_dims = list(tx = c('lat', 'time'),
                                  tn = c('lat', 'time'),
                                  pr = c('lat', 'time')),
               output_dims = list(fd = c('lat', 'time'),
                                  cdd = c('lat','time')))

wf <- AddStep(inputs = list(hist_tx, hist_tn, hist_pr), step,
              times = time, lats = lat)

## locally
res <- Compute(wf$fd,
               chunks = list(lon = 2))

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

res <- Compute(wf$fd,
               chunks = list(lon = 2),
                threads_load = 1,
                 threads_compute = 4,
                 cluster = list(queue_host = queue_host,
                                queue_type = 'lsf',
                                extra_queue_params = list('#BSUB -q bsc_es'),
                                cores_per_job = 4,
                                temp_dir = temp_dir,
                                polling_period = 10,
                                job_wallclock = '01:00',
                                max_jobs = 2,
                                bidirectional = FALSE),
                 ecflow_suite_dir = ecflow_suite_dir,
                 wait = TRUE)

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/'

res <- Compute(wf$fd,
               chunks = list(lon = 2),
                threads_load = 2,
                 threads_compute = 4,
                 cluster = list(queue_host = queue_host,
                                queue_type = 'slurm',
                                extra_queue_params = list('#BSUB -q bsc_es'),
                                cores_per_job = 4,
                                temp_dir = temp_dir,
                                polling_period = 10,
                                job_wallclock = '00:01:00',
                                max_jobs = 2,
                                bidirectional = FALSE),
                 ecflow_suite_dir = ecflow_suite_dir,
                 wait = TRUE)

s2dv::PlotEquiMap(res$fd[,1,1,1,1,1,1,], lon = lon, lat = lat, filled.con = F)



#/esarchive/exp/CMIP6/historical/ecearth/cmip6-historical_i0p1/CMIP/EC-Earth-Consortium/EC-Earth3/historical/r2i1p1f1/day/tasmin/gr/v20200120
