
library(CSTools)
library(startR)
library(s2dv)
library(lubridate)
library(multiApply)
source('/esarchive/scratch/jramon/GitLab_jramon/downscaling/csdownscale/R/Interpolation.R')
source('/esarchive/scratch/jramon/GitLab_jramon/downscaling/csdownscale/R/Intlr.R')
source('/esarchive/scratch/jramon/GitLab_jramon/downscaling/csdownscale/R/Utils.R')

plotpath <- '/esarchive/scratch/jramon/downscaling/plots/ip'
target_grid <- '/esarchive/recon/ecmwf/era5/daily/tasmax-r1440x721cds/tasmax_201702.nc'

#lonmin <- -22
#lonmax <- 45
#latmin <- 27
#latmax <- 72
lonmin <- -11.5
lonmax <- 5.35
latmin <- 35.1
latmax <- 44.1


#sdates <- format(ymd("19811201") + rep(years(0:36),each=1),"%Y%m%d")
sdates <- c('20000201','20010201','20020201','20030201','20040201','20050201','20060201','20070201')

obs1 <- startR::Start(dat = '/esarchive/recon/ecmwf/era5/monthly_mean/$var$_f1h-r1440x721cds/$var$_$sdate$.nc',
                     var = 'tas', time = indices(1), lat = values(list(latmin, latmax)),
                     sdate = format(ymd(sdates), "%Y%m"), lat_reorder = Sort(decreasing = FALSE),
                     lon = values(list(lonmin, lonmax)), lon_reorder = CircularSort(-180, 180),
                     synonims = list(var = c('var','variable'), lon = c('lon', 'longitude'),
                     lat = c('lat', 'latitude')), return_vars = list(lat = 'dat', lon = 'dat'),
                     num_procs = 1, retrieve = TRUE)

exp1 <- startR::Start(dat = '/esarchive/exp/ecmwf/system5c3s/monthly_mean/$var$_f6h/$var$_$sdate$.nc',
                     var = 'tas', time = indices(1), member = indices(1:3), sdate = sdates,
                     lat = values(list(latmin, latmax)), lat_reorder = Sort(decreasing = FALSE),
                     lon = values(list(lonmin, lonmax)), lon_reorder = CircularSort(-180, 180),
                     synonims = list(var = c('var','variable'), lon = c('lon', 'longitude'),
                     lat = c('lat', 'latitude'), member = c('member','ensemble')),
                     return_vars = list(lat = 'dat', lon = 'dat'),
                     num_procs = 1, retrieve = TRUE)

#----------------------------------
# Downscaling with Intlr
#----------------------------------
ind_rdm <- array(NA, dim = c('sdate' = length(sdates), member = 3, 'vars' = 2))
ind_rdm[,,1] <- rnorm(n = prod(length(sdates), 3), mean = 0, sd = 1)
ind_rdm[,,2] <- rnorm(n = prod(length(sdates), 3), mean = 0, sd = 1)

down_1_bas <- Intlr(exp = exp1, obs = obs1, exp_lats = attr(exp1, "Variables")$dat1$lat, 
                    exp_lons = attr(exp1, "Variables")$dat1$lon, obs_lats = attr(obs1, "Variables")$dat1$lat, 
                    obs_lons = attr(obs1, "Variables")$dat1$lon, target_grid = target_grid, lr_method = "basic",
                    int_method = "bilinear", predictors = NULL, loocv = TRUE, ncores = 4)

down_1_lsc <- Intlr(exp = exp1, obs = obs1, exp_lats = attr(exp1, "Variables")$dat1$lat,
                    exp_lons = attr(exp1, "Variables")$dat1$lon, obs_lats = attr(obs1, "Variables")$dat1$lat,
                    obs_lons = attr(obs1, "Variables")$dat1$lon, target_grid = target_grid, lr_method = "large-scale",
                    int_method = "conservative", predictors = ind_rdm, loocv = TRUE, ncores = 4)

down_1_nn <- Intlr(exp = exp1, obs = obs1, exp_lats = attr(exp1, "Variables")$dat1$lat,
                   exp_lons = attr(exp1, "Variables")$dat1$lon, obs_lats = attr(obs1, "Variables")$dat1$lat,
                   obs_lons = attr(obs1, "Variables")$dat1$lon, target_grid = target_grid, lr_method = "4nn",
                   int_method = "conservative", predictors = NULL, loocv = TRUE, ncores = 4)

#----------------------------------
# Downscaling to point locations with CST_Intlr
# Note: load first NCO module
#----------------------------------
points <- list(lat = c(36.1, 38.7), lon = c(-1.9, 0.8))
down_2_bas <- Intlr(exp = exp1, obs = obs1, exp_lats = attr(exp1, "Variables")$dat1$lat,
                    exp_lons = attr(exp1, "Variables")$dat1$lon, obs_lats = attr(obs1, "Variables")$dat1$lat,
                    obs_lons = attr(obs1, "Variables")$dat1$lon, target_grid = target_grid, lr_method = "basic",
                    points = points, method_point_interp = 'bilinear', int_method = NULL, predictors = NULL, 
                    source_file_exp = attr(exp1, "Files")[1,1,1], source_file_obs = attr(obs1, "Files")[1,1,1],
                    loocv = TRUE, ncores = 4)

down_2_lsc <- Intlr(exp = exp1, obs = obs1, exp_lats = attr(exp1, "Variables")$dat1$lat,
                    exp_lons = attr(exp1, "Variables")$dat1$lon, obs_lats = attr(obs1, "Variables")$dat1$lat,
                    obs_lons = attr(obs1, "Variables")$dat1$lon, target_grid = target_grid, lr_method = "large-scale",
                    points = points, method_point_interp = 'bilinear', int_method = NULL, predictors = ind_rdm,
                    source_file_exp = attr(exp1, "Files")[1,1,1], source_file_obs = attr(obs1, "Files")[1,1,1],
                    loocv = TRUE, ncores = 4)

down_2_nn <- Intlr(exp = exp1, obs = obs1, exp_lats = attr(exp1, "Variables")$dat1$lat,
                   exp_lons = attr(exp1, "Variables")$dat1$lon, obs_lats = attr(obs1, "Variables")$dat1$lat,
                   obs_lons = attr(obs1, "Variables")$dat1$lon, target_grid = target_grid, lr_method = "4nn",
                   points = points, method_point_interp = 'bilinear', int_method = NULL, predictors = NULL,
                   source_file_exp = attr(exp1, "Files")[1,1,1], source_file_obs = attr(obs1, "Files")[1,1,1],
                   loocv = TRUE, ncores = 4)

#----------------------------------
# Create s2dv objects
#----------------------------------
obs1 <- s2dv_cube(obs1, lat = attr(obs1, "Variables")$dat1$lat, lon = attr(obs1, "Variables")$dat1$lon,
                 source_files = attr(obs, "Files")[1,1,])
attr(obs1$lon,"first_lon") <- obs1$lon[1]
attr(obs1$lon,"last_lon") <- obs1$lon[length(obs1$lon)]
attr(obs1$lat,"first_lat") <- obs1$lat[1]
attr(obs1$lat,"last_lat") <- obs1$lat[length(obs1$lat)]

exp1 <- s2dv_cube(exp1, lat = attr(exp1, "Variables")$dat1$lat, lon = attr(exp1, "Variables")$dat1$lon,
                 source_files = attr(exp1, "Files")[1,1,])
attr(exp1$lon,"first_lon") <- exp1$lon[1]
attr(exp1$lon,"last_lon") <- exp1$lon[length(exp1$lon)]
attr(exp1$lat,"first_lat") <- exp1$lat[1]
attr(exp1$lat,"last_lat") <- exp1$lat[length(exp1$lat)]

#----------------------------------
# Downscaling with CST_Intlr
#----------------------------------
down_3_bas <- CST_Intlr(exp = exp1, obs = obs1, target_grid = target_grid, lr_method = "basic",
                        int_method = "bilinear", predictors = NULL, loocv = TRUE, ncores = 4)

down_3_lsc <- CST_Intlr(exp = exp1, obs = obs1, target_grid = target_grid, lr_method = "large-scale",
                        int_method = "bilinear", predictors = ind_rdm, loocv = TRUE, ncores = 4)

down_3_nn <- CST_Intlr(exp = exp1, obs = obs1, target_grid = target_grid, lr_method = "4nn",
                       int_method = "bilinear", predictors = NULL, loocv = TRUE, ncores = 4)

#----------------------------------
# Downscaling to point locations with Intlr
#----------------------------------
points <- list(lat = c(36.1, 38.7), lon = c(-1.9, 0.8))
obs_point <- Subset(obs1, along = c("lat", "lon"), indices = list(1:2,1), drop = FALSE)
obs_point <- MergeDims(data = obs_point, merge_dims = c("lat", "lon"), rename_dim = 'location')
down_4_bas <- Intlr(exp = exp1, obs = obs_point, exp_lats = attr(exp1, "Variables")$dat1$lat,
                    exp_lons = attr(exp1, "Variables")$dat1$lon, obs_lats = points$lat, obs_lons = points$lon, 
                    target_grid = target_grid, lr_method = "basic", points = points, method_point_interp = 'bilinear', 
                    int_method = 'bilinear', predictors = NULL, source_file_exp = attr(exp1, "Files")[1,1,1], 
                    source_file_obs = attr(obs1, "Files")[1,1,1], loocv = TRUE, ncores = 1)



