
library(CSTools)
library(startR)
library(s2dv)
library(lubridate)
library(multiApply)
library(FNN)
source('/esarchive/scratch/jramon/GitLab_jramon/downscaling/csdownscale/R/Interpolation.R')
source('/esarchive/scratch/jramon/GitLab_jramon/downscaling/csdownscale/R/Analogs.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
extra <- 1

#sdates <- c('20000201','20010201','20020201','20030201','20040201','20050201','20060201','20070201')
#sdates <- format(seq(ymd("20000101"), ymd("20071201"), '1 month'), "%Y%m%d")
sdates_exp <- format(ymd("20000501") + rep(years(0:8), each=1),"%Y%m%d")
#sdates_obs  <- format(ymd("20000401") + months(0:2) + rep(years(0:2), each=3),"%Y%m")
sdates_obs  <- format(ymd("20000501") + rep(years(0:8), each=1),"%Y%m")

#---------------------------
# Observations
#---------------------------
# dim(obs) <- c('sdate', 'smonth', ...)
# ERA5 observations: used in the training
# change to time = indices(1:28)
obs_era5 <- startR::Start(dat = '/esarchive/recon/ecmwf/era5/daily_mean/$var$_f1h/$var$_$sdate$.nc',
                     var = 'tas', time = indices(1:28), latitude = values(list(latmin - extra, latmax + extra)), 
                     sdate = sdates_obs, latitude_reorder = Sort(decreasing = FALSE),
                     longitude = values(list(lonmin - extra, lonmax + extra)), 
                     longitude_reorder = CircularSort(-180, 180),
                     synonims = list(var = c('var','variable'), longitude = c('lon', 'longitude'), 
                     latitude = c('lat', 'latitude')), return_vars = list(latitude = 'dat', longitude = 'dat'),
                     num_procs = 1, retrieve = TRUE)

# Create 's2dv_cube' object
obs_s2dv <- s2dv_cube(obs_era5, lat = attr(obs_era5, "Variables")$dat1$lat, 
                      lon = attr(obs_era5, "Variables")$dat1$lon, source_files = attr(obs_era5, "Files")[1,1,])

# ERA5-Land observations: high-resolution observations
# change to time = indices(1:28)
obs_hres <- startR::Start(dat = '/esarchive/recon/ecmwf/era5land/daily_mean/$var$_f1h/$var$_$sdate$.nc',
                      var = 'tas', time = indices(1:28), latitude = values(list(latmin - extra, latmax + extra)),
                      sdate = sdates_obs, latitude_reorder = Sort(decreasing = FALSE),
                      longitude = values(list(lonmin - extra, lonmax + extra)),
                      longitude_reorder = CircularSort(-180, 180),
                      synonims = list(var = c('var','variable'), longitude = c('lon', 'longitude'),
                      latitude = c('lat', 'latitude')), return_vars = list(latitude = 'dat', longitude = 'dat'),
                      num_procs = 1, retrieve = TRUE)


# Create window to look for analogues
#dim(obs_s2dv$data) <- c(data = 1, var = 1, time = 28, lat = 39, smonth = 3, sdate = 5, lon = 67) # correct
#obs_new <- Apply(obs_s2dv$data, target_dims = list(c("time", "smonth")), fun = as.vector, output_dims = "time")$output1
#wlen <- 7

#result <- array(NA, dim = c(window = 2*wlen + 1, time = 28, lat = 39, sdate = 7, lon = 67))
#for (i in 1:28) {
#  result[,i,,,] <- obs_new[(28 + i - wlen):(28 + i + wlen),1,1,,,]
#}
#obs_s2dv$data <- result

# > dim(obs$data)
#  data    var   time    lat smonth  sdate    lon 
#     1      1     28     39      3      7     67 
# > dim(result)
#  window   time    lat  sdate    lon 
#      15     28     39      7     67 

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

# Create 's2dv_cube' object
exp_s2dv <- s2dv_cube(exp, lat = attr(exp, "Variables")$dat1$lat, lon = attr(exp, "Variables")$dat1$lon,
                      source_files = attr(exp, "Files")[1,1,])

#---------------------------
# Analogs with no 'window' dimension in the initial object and with cv
#---------------------------
ana_1 <- Analogs(exp = exp, obs = obs_hres, exp_lats = attr(exp, "Variables")$dat1$lat,
                 exp_lons = attr(exp, "Variables")$dat1$lon, obs_lats = attr(obs_hres, "Variables")$dat1$lat,
                 obs_lons = attr(obs_hres, "Variables")$dat1$lon, grid_exp = attr(exp, "Files")[1,1,1],
                 fun_analog = NULL, lat_dim = 'latitude', lon_dim = 'longitude', loocv_window = TRUE, 
                 return_indices = TRUE, ncores = 4)

#---------------------------
# Analogs with no 'window' dimension in the initial object and without cv
#---------------------------
ana_2 <- Analogs(exp = exp, obs = obs_era5, exp_lats = attr(exp, "Variables")$dat1$lat,
                 exp_lons = attr(exp, "Variables")$dat1$lon, obs_lats = attr(obs_era5, "Variables")$dat1$lat,
                 obs_lons = attr(obs_era5, "Variables")$dat1$lon, grid_exp = attr(exp, "Files")[1,1,1],
                 fun_analog = NULL, lat_dim = 'latitude', lon_dim = 'longitude', loocv_window = FALSE, ncores = 4)

#---------------------------
# Analogs with 'window' dimension in the initial object and with cv
#---------------------------
dim(obs_era5) <- c(data = 1, var = 1, time = 28, latitude = 39, smonth = 3, sdate = 3, longitude = 67)
obs_window <- generate_window(obj = obs_era5, sdate_dim = 'sdate', time_dim = 'time', size = 7, loocv = TRUE) 
ana_3 <- Analogs(exp = exp, obs = obs_window, exp_lats = attr(exp, "Variables")$dat1$lat,
                 exp_lons = attr(exp, "Variables")$dat1$lon, obs_lats = attr(obs_era5, "Variables")$dat1$lat,
                 obs_lons = attr(obs_era5, "Variables")$dat1$lon, grid_exp = attr(exp, "Files")[1,1,1],
                 fun_analog = NULL, lat_dim = 'latitude', lon_dim = 'longitude', ncores = 4)

#---------------------------
# Analogs with cv and returning the indices
#---------------------------
ana_4 <- Analogs(exp = exp, obs = obs_era5, exp_lats = attr(exp, "Variables")$dat1$lat,
                 exp_lons = attr(exp, "Variables")$dat1$lon, obs_lats = attr(obs_era5, "Variables")$dat1$lat,
                 obs_lons = attr(obs_era5, "Variables")$dat1$lon, grid_exp = attr(exp, "Files")[1,1,1],
                 fun_analog = NULL, lat_dim = 'latitude', lon_dim = 'longitude', loocv_window = TRUE, 
                 return_indices = TRUE, ncores = 4)

#---------------------------
# Analogs with 'window' dimension and returning the indices
#---------------------------
ana_5 <- Analogs(exp = exp, obs = obs_window, exp_lats = attr(exp, "Variables")$dat1$lat,
                 exp_lons = attr(exp, "Variables")$dat1$lon, obs_lats = attr(obs_era5, "Variables")$dat1$lat,
                 obs_lons = attr(obs_era5, "Variables")$dat1$lon, grid_exp = attr(exp, "Files")[1,1,1],
                 fun_analog = NULL, lat_dim = 'latitude', lon_dim = 'longitude', return_indices = TRUE, ncores = 4)

#---------------------------
# Analogs with NA's in obs
#---------------------------
obs_era5[,,,10:20,,] <- NA
ana_mean_2 <- Analogs(exp = exp, obs = obs_era5, exp_lats = attr(exp, "Variables")$dat1$lat,
                      exp_lons = attr(exp, "Variables")$dat1$lon, obs_lats = attr(obs_era5, "Variables")$dat1$lat,
                      obs_lons = attr(obs_era5, "Variables")$dat1$lon, source_file_exp = attr(exp, "Files")[1,1,1],
                      fun_analog = 'mean', ncores = 4,lat_dim='latitude',lon_dim='longitude')


#---------------------------
# Analogs with two types of observations
#---------------------------
ana_6 <- Analogs(exp = exp, obs = obs_hres, obs2 = obs_era5, exp_lats = attr(exp, "Variables")$dat1$lat,
                 exp_lons = attr(exp, "Variables")$dat1$lon, obs_lats = attr(obs_hres, "Variables")$dat1$lat,
                 obs_lons = attr(obs_hres, "Variables")$dat1$lon, obs2_lats = attr(obs_era5, "Variables")$dat1$lat, 
                 obs2_lons = attr(obs_era5, "Variables")$dat1$lon, grid_exp = attr(exp, "Files")[1,1,1],
                 fun_analog = NULL, lat_dim = 'latitude', lon_dim = 'longitude', loocv_window = TRUE, 
                 return_indices = TRUE, ncores = 4)

#---------------------------
# Analogs with CST_Analogs()
#---------------------------
ana_mean <- CST_Analogs(exp = exp2, obs = obs_s2dv, fun_analog = "mean", ncores = 4)

# test examples
ana_mean <- Analogs(exp = exp, obs = obs, fun_analog = "mean", ncores = 4)
ana_wmean <- Analogs(exp = exp, obs = obs, fun_analog = "wmean", ncores = 4)
ana_max <- Analogs(exp = exp, obs = obs, fun_analog = "max", ncores = 4)
ana_min <- Analogs(exp = exp, obs = obs, fun_analog = "min", ncores = 4)
ana_median <- Analogs(exp = exp, obs = obs, fun_analog = "median", ncores = 4)

# plot examples
s2dv::PlotEquiMap(ana_mean$data[,,1,1,1,1,1], lat = ana_mean$lat, lon = ana_mean$lon, filled.continents = F,
                  toptitle = 'Analogs mean', brks = seq(275,288,1), height = 8, width = 10, 
                  fileout = file.path(plotpath, "ip_ana_mean_2.png"))
s2dv::PlotEquiMap(ana_wmean$data[1,,,1,1,1,1], lat = ana_wmean$lat, lon = ana_wmean$lon, filled.continents = F,
                  toptitle = 'Analogs weighted mean', brks = seq(275,288,1), height = 8, width = 10,
                  fileout = file.path(plotpath, "ip_ana_wmean.png"))
s2dv::PlotEquiMap(ana_max$data[1,,,1,1,1,1], lat = ana_max$lat, lon = ana_max$lon, filled.continents = F,
                  toptitle = 'Analogs max', brks = seq(275,288,1), height = 8, width = 10,
                  fileout = file.path(plotpath, "ip_ana_max.png"))
s2dv::PlotEquiMap(ana_median$data[1,,,1,1,1,1], lat = ana_max$lat, lon = ana_max$lon, filled.continents = F,
                  toptitle = 'Analogs median', brks = seq(275,288,1), height = 8, width = 10,
                  fileout = file.path(plotpath, "ip_ana_median.png"))
s2dv::PlotEquiMap(ana_min$data[1,,,1,1,1,1], lat = ana_max$lat, lon = ana_max$lon, filled.continents = F,
                  toptitle = 'Analogs min', brks = seq(275,288,1), height = 8, width = 10,
                  fileout = file.path(plotpath, "ip_ana_min.png"))










