
library(startR)
library(ClimProjDiags)
library(multiApply)
source('/esarchive/scratch/cdelgado/gitlab/c3s-scripts/R/AMV.R')
source('/esarchive/scratch/cdelgado/gitlab/c3s-scripts/R/SPOD.R')
source('/esarchive/scratch/cdelgado/gitlab/c3s-scripts/R/TPI.R')
source('/esarchive/scratch/cdelgado/gitlab/c3s-scripts/R/GSAT.R')
source('/esarchive/scratch/cdelgado/gitlab/c3s-scripts/R/GMST.R')
source('./startR_call_indices.R')

save_output = FALSE
output_path = '/esarchive/scratch/cdelgado/copernicus_outputs/outputs_indices/'
if (save_output){if (!dir.exists(output_path)){dir.create(file.path(output_path), recursive = TRUE)}}

models_dcpp = c('EC-Earth3','HadGEM3','MPI-ESM1.2-HR') #,'CMCC-CM2')
models_hist = c('EC-Earth3_hist','HadGEM3_hist','MPI-ESM1.2-HR_hist') #,'CMCC-CM2_hist')
datasets_tas = c('JRA-55')
datasets_tas_tos = c('GISTEMPv4')
indices = c('AMV','TPI','GSAT','GMST') # ,'SPOD')
types = c('dcpp','hist','obs')
HPCs = c('power9','nord3','local')

HPC = HPCs[2]
type = types[1]
index = indices[4]
data_set = 2

## Function and varible depending on the index
if (index == 'AMV'){
  fun = AMV
  variable = 'tos'
} else if (index == 'SPOD'){
  fun = SPOD
  variable = 'tos'
} else if (index == 'TPI'){
  fun = TPI
  variable = 'tos'
} else if (index == 'GSAT'){
  variable = 'tas'
  fun = GSAT
} else if (index == 'GMST'){
  variable = 'tas_tos'
  fun = GMST
  
  mask_obs = '/esarchive/scratch/cdelgado/copernicus_outputs/outputs_indices/mask_GISTEMPv4_for_indices.rds'
}

indices_for_clim = 10:50 # 1970-2010 as reference period for climatology

## Parameters for the Start() call and Compute() workflow
if (type == 'dcpp'){
  
  members = 'all'
  sdate_min = 1960
  sdate_max = 2018
  fmonth_min = 1
  fmonth_max = 122 # 10 fyears and the first 2 months (init in Nov)
  
  datasets = models_dcpp
  
  retrieve = FALSE
  target_dims = c('sdate','fmonth','lat','lon')
  output_dims = c('fyear','sdate')
  
} else if (type == 'hist'){
  
  members = 'all'
  year_min = 1961
  year_max = 2014
  months = 'all'
  
  datasets = models_hist
  
  retrieve = FALSE
  target_dims = c('year','month','lat','lon')
  output_dims = c('year')
  
} else if (type == 'obs'){
  
  year_min = 1961
  year_max = 2019
  months = 'all'
  
  if (index %in% c('AMV','SPOD','TPI','GMST')){
    datasets = datasets_tas_tos
  } else if (index == 'GSAT'){
    datasets = datasets_tas
  }
  
  retrieve = TRUE
}

dataset = datasets[data_set]
# for (dataset in datasets){

## Start() call (two of them if tas_tos)
if (type == 'dcpp'){
  
  if (variable %in% c('tas','tos')){
    
    data = startR_call_indices(type=type, dataset=dataset, variable=variable, 
                               sdate_min=sdate_min, sdate_max=sdate_max, 
                               fmonth_min=fmonth_min, fmonth_max=fmonth_max,
                               members=members, retrieve=retrieve)
    
  } else if (variable == 'tas_tos'){
    
    data_tas = startR_call_indices(type=type, dataset=dataset, variable='tas', 
                                   sdate_min=sdate_min, sdate_max=sdate_max, 
                                   fmonth_min=fmonth_min, fmonth_max=fmonth_max,
                                   members=members, retrieve=retrieve)
    
    data_tos = startR_call_indices(type=type, dataset=dataset, variable='tos', 
                                   sdate_min=sdate_min, sdate_max=sdate_max, 
                                   fmonth_min=fmonth_min, fmonth_max=fmonth_max,
                                   members=members, retrieve=retrieve)
  }
  
} else if (type == 'hist'){
  
  if (variable %in% c('tas','tos')){
    
    data = startR_call_indices(type=type, dataset=dataset, variable=variable, 
                               year_min=year_min, year_max=year_max, 
                               months=months, members=members, retrieve=retrieve)
    
  } else if (variable == 'tas_tos'){
    
    if (dataset == 'MPI-ESM1.2-HR_hist'){
      
      error = startR_call_indices(type=type, dataset=dataset, variable='tas_tos', 
                          year_min=year_min, year_max=year_max, 
                          months=months, members=members, retrieve=retrieve)
      
    } else {
    
      data_tas = startR_call_indices(type=type, dataset=dataset, variable='tas', 
                                     year_min=year_min, year_max=year_max, 
                                     months=months, members=members, retrieve=retrieve)
      
      data_tos = startR_call_indices(type=type, dataset=dataset, variable='tos', 
                                     year_min=year_min, year_max=year_max, 
                                     months=months, members=members, retrieve=retrieve)
    }
  }
  
} else if (type == 'obs'){
  
  if (variable %in% c('tas','tos')){
    
    # data = startR_call_indices(type=type, dataset=dataset, variable=variable, 
    #                            year_min=year_min, year_max=year_max, 
    #                            months=months, retrieve=retrieve)
    
    if (dataset == 'GISTEMPv4'){
      data_aux = readRDS(paste0(output_path, 'tas_tos_GISTEMPv4_ref1970-2010.rds'))
      data = data_aux$data
      attributes(data)$Variables$dat1$lat = data_aux$lat
      attributes(data)$Variables$dat1$lon = data_aux$lon
    }
    
  } else if (variable == 'tas_tos'){
    
    # data_tas = startR_call_indices(type=type, dataset=dataset, variable=variable, 
    #                                year_min=year_min, year_max=year_max, 
    #                                months=months, retrieve=retrieve)
    # data_tos = data_tas # it is supposed that the observed dataset are already a combination of tas and tos
    
    if (dataset == 'GISTEMPv4'){
      data = readRDS(paste0(output_path, 'tas_tos_GISTEMPv4_ref1970-2010.rds'))
      data_tas = data$data
      data_tos = data$data
      attributes(data_tas)$Variables$dat1$lat = data$lat
      attributes(data_tas)$Variables$dat1$lon = data$lon
      attributes(data_tos)$Variables$dat1$lat = data$lat
      attributes(data_tos)$Variables$dat1$lon = data$lon
    }
    
  }
  
}

if (!is.null(attr(data,'NotFilesFound'))){
  stop('there are files not found')
}

if (variable %in% c('tas','tos')){
  
  data_lats = as.vector(attributes(data)$Variables$dat1$lat)
  data_lons = as.vector(attributes(data)$Variables$dat1$lon)
  
} else if (variable == 'tas_tos'){
  
  data_lats_tas = as.vector(attributes(data_tas)$Variables$dat1$lat)
  data_lons_tas = as.vector(attributes(data_tas)$Variables$dat1$lon)
  data_lats_tos = as.vector(attributes(data_tos)$Variables$dat1$lat)
  data_lons_tos = as.vector(attributes(data_tos)$Variables$dat1$lon)
  
  if (identical(round(data_lats_tas,2),round(data_lats_tos,2)) & identical(round(data_lats_tas,2),round(data_lats_tos,2))){
    data_lats = data_lats_tas
    data_lons = data_lons_tas
  } else {
    stop('Latitudes or longitudes of tas and tos are not identical')
  }
  
}

## Workflow with Compute() for dcpp and hist, and locally for obs

if (retrieve == TRUE){
  
  if (index %in% c('AMV','SPOD','TPI','GSAT')){
    
    INDEX = fun(data = data, data_lats = data_lats, data_lons = data_lons, 
                type = type, indices_for_clim = indices_for_clim)
    
  } else if (index == 'GMST'){
    
    INDEX = fun(data_tas = data_tas, data_tos = data_tos, 
                data_lats = data_lats, data_lons = data_lons, 
                type = type, mask_obs = mask_obs, indices_for_clim = indices_for_clim)
  }
  
  INDEX = drop(INDEX)
  
} else if (retrieve == FALSE){
  
  if (index %in% c('AMV','SPOD','TPI','GSAT')){
    
    step = Step(fun = fun,
                use_libraries = c('startR','s2dv','ClimProjDiags'),
                target_dims = target_dims,
                output_dims = output_dims)
    
    wf = AddStep(inputs = data,
                 data_lats = data_lats,
                 data_lons = data_lons,
                 type = type,
                 indices_for_clim = indices_for_clim,
                 step = step)
    
    chunks = list(member = attr(data,'Dimensions')['member'])
    
  } else if (index == 'GMST'){
    
    step = Step(fun = fun,
                use_libraries = c('startR','s2dv','ClimProjDiags'),
                target_dims = list(data_tas = target_dims, data_tos = target_dims),
                output_dims = output_dims)
    
    wf = AddStep(inputs = list(data_tas, data_tos),
                 data_lats = data_lats,
                 data_lons = data_lons,
                 type = type,
                 mask_obs = mask_obs,
                 indices_for_clim = indices_for_clim,
                 step = step)
    
    chunks = list(member = attr(data_tas,'Dimensions')['member'])
  }
  
  if (HPC == 'power9'){
    
    cluster = list(queue_host = 'power',
                   queue_type = 'slurm',
                   temp_dir = '/gpfs/scratch/bsc32/bsc32924/startR_hpc/',
                   r_module = 'R/3.5.0-foss-2018b',
                   CDO_module = 'CDO/1.9.5-foss-2018b',
                   cores_per_job = 4,
                   job_wallclock = '01:00:00',
                   max_jobs = 10,
                   extra_queue_params = list('#SBATCH --mem-per-cpu=60000'), # to avoid memory issues on p9: request 160 cores and add --exclusive
                   bidirectional = FALSE,
                   polling_period = 20)
    
  } else if (HPC == 'nord3'){
    
    cluster = list(queue_host = 'nord1.bsc.es',
                   queue_type = 'lsf',
                   temp_dir = '/gpfs/scratch/bsc32/bsc32924/startR_hpc/',
                   cores_per_job = 2,
                   job_wallclock = '01:00',
                   max_jobs = 10,
                   extra_queue_params = list('#BSUB -x -q bsc_es -M 60000'), # node exclusive: #BSUB -x
                   bidirectional = FALSE,
                   polling_period = 10)
  
  } else if (HPC == 'local'){
    
    cluster = NULL
    
  }
  
  res = Compute(workflow = wf,
                chunks = chunks,
                threads_load = 1,
                threads_compute = 4,
                cluster = cluster,
                ecflow_suite_dir = '/home/Earth/cdelgado/Desktop/startR_local/',
                wait = TRUE)
  
  INDEX = drop(res$output1)
} 

if (save_output){
  saveRDS(INDEX, file = paste0(output_path,index,'_',dataset,'.rds'))
}
print(paste0('FINISHED ', index, ' for ', dataset))

#   }
# }

