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

plotpath <- '/esarchive/scratch/jramon/downscaling/plots/ip'
target_grid <- '/esarchive/recon/ecmwf/era5/monthly_mean/tas_f1h/tas_200002.nc'

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

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

obs <- 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)),
                     lat_reorder = Sort(decreasing = FALSE), lon = values(list(lonmin, lonmax)),
                     lon_reorder = CircularSort(-180, 180), sdate = format(ymd(sdates), "%Y%m"), 
                     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)

exp <- 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 to fields w LogisticReg()
res <- LogisticReg(exp = exp, 
                   obs = obs, 
                   exp_lats = attr(exp, "Variables")$dat1$lat,
                   exp_lons = attr(exp, "Variables")$dat1$lon, 
                   obs_lats = attr(obs, "Variables")$dat1$lat,
                   obs_lons = attr(obs, "Variables")$dat1$lon, 
                   probs_cat = 1:4/5,
                   target_grid = target_grid, 
                   int_method = "bil",
                   loocv = TRUE,
                   return_most_likely_cat = TRUE)

# downscaling to fields w CST_LogisticReg()
obs <- s2dv_cube(obs, lat = attr(obs, "Variables")$dat1$lat, lon = attr(obs, "Variables")$dat1$lon,
                 source_files = attr(obs, "Files")[1,1,])
exp <- s2dv_cube(exp, lat = attr(exp, "Variables")$dat1$lat, lon = attr(exp, "Variables")$dat1$lon,
                 source_files = attr(exp, "Files")[1,1,])

res_cst1 <- CST_LogisticReg(exp = exp,
                            obs = obs,
                            probs_cat = 1:2/3,
                            target_grid = target_grid,
                            int_method = "bil",
                            log_reg_method = "ens_mean",
                            return_most_likely_cat = TRUE,
                            loocv = TRUE)

res_cst2 <- CST_LogisticReg(exp = exp,
                            obs = obs,
                            probs_cat = 1:2/3,
                            target_grid = target_grid,
                            int_method = "bil",
                            log_reg_method = "ens_mean_sd",
                            return_most_likely_cat = TRUE,
                            loocv = TRUE)

res_cst3 <- CST_LogisticReg(exp = exp,
                            obs = obs,
                            probs_cat = 1:2/3,
                            target_grid = target_grid,
                            int_method = "bil",
                            log_reg_method = "sorted_members",
                            return_most_likely_cat = TRUE,
                            loocv = TRUE)


# Plot Most Likely Tercile
colors <- c("#33BFD1", "#b5b5b5", "#FF764D")
PlotEquiMap(res_cst1$data[21,1,1,1,1:29,2:58], lat = res_cst1$lat[1:29], lon = res_cst1$lon[2:58], 
  filled.continents = F, brks = c(0,1,2,3), cols = colors, 
  fileout = "/esarchive/scratch/jramon/downscaling/plots/logreg_ens-mean_202007.png", width = 7, height = 5)

PlotEquiMap(res_cst2$data[21,1,1,1,1:29,2:58], lat = res_cst2$lat[1:29], lon = res_cst2$lon[2:58], 
  filled.continents = F, brks = c(0,1,2,3), cols = colors, 
  fileout = "/esarchive/scratch/jramon/downscaling/plots/logreg_ens-mean-sd_202007.png", width = 7, height = 5)

PlotEquiMap(res_cst3$data[21,1,1,1,1:29,2:58], lat = res_cst3$lat[1:29], lon = res_cst3$lon[2:58],
  filled.continents = F, brks = c(0,1,2,3), cols = colors,
  fileout = "/esarchive/scratch/jramon/downscaling/plots/logreg_sorted-members_202007.png", width = 7, height = 5)


exp_anom <- get_ens_mean_anom(obj_ens = exp, member_dim = "member", sdate_dim = "sdate")
PlotEquiMap(exp_anom[1,1,1,,,21], lat = attr(exp, "Variables")$dat1$lat, lon = attr(exp, "Variables")$dat1$lon,
  filled.continents = F, fileout = "/esarchive/scratch/jramon/downscaling/plots/exp_tas_202007.png", width = 7, 
  height = 5)


# downscaling to point locations
points <- list(lat = c(36, 36.1), lon = c(0.6, 1))
source_file <- '/esarchive/exp/ecmwf/system5c3s/monthly_mean/tas_f6h/tas_20000201.nc'
method_point_interp <- 'bilinear'
res <- LogisticReg(exp = exp, obs = obs, exp_lats = attr(exp, "Variables")$dat1$lat, 
                   exp_lons = attr(exp, "Variables")$dat1$lon, obs_lats = attr(obs, "Variables")$dat1$lat,
                   obs_lons = attr(obs, "Variables")$dat1$lon, target_grid = target_grid, int_method = "bil",
                   points = points, source_file = source_file, method_point_interp = method_point_interp)

