library(s2dv)
library(CSTools)
library(FNN)
library(multiApply)
source('/esarchive/scratch/jramon/GitLab_jramon/downscaling/csdownscale/R/Utils.R')
source('/esarchive/scratch/jramon/GitLab_jramon/downscaling/csdownscale/R/Interpolation.R')
source('/esarchive/scratch/jramon/GitLab_jramon/downscaling/csdownscale/R/Intbc.R')
output_dir <- '/esarchive/scratch/allabres/review_csdownscale/'

load(paste0(output_dir, 'daily_data_to_test_1993-2016_era5land.RData'))
exp <- drop(exp) #to remove unnecessary dimensions
obs <- drop(obs) #to remove unnecessary dimensions
target_grid <- attr(obs, 'Files')[1]

int_method <- 'bil' # "con", "bil", "bic", "nn", "con2"
bc_method <- 'simple_bias' # 'simple_bias', 'calibration', 'quantile_mapping'

# test with arrays:
    res <- Intbc(exp = exp, obs = obs, 
                    exp_lats = lats.hcst, exp_lons = lons.hcst, obs_lats = lats.ref, obs_lons = lons.ref, 
                    target_grid = target_grid, bc_method = bc_method, int_method = int_method, 
                    points = NULL, method_point_interp = NULL, 
                    lat_dim = "latitude", lon_dim = "longitude", sdate_dim = "sdate", member_dim = "ensemble", 
                    source_file = NULL, region = NULL, cal_method = "mse_min", ncores = 1)
    # Error in BiasCorrection(exp = exp_interpolated$data, obs = obs_interpolated$data) :
    #     Parameter 'exp' must have the dimensions 'member' and 'sdate'.

    names(dim(exp)) <- c('sdate', 'member', 'latitude', 'longitude')
    names(dim(obs)) <- c('sdate', 'latitude', 'longitude')
    res1 <- Intbc(exp = exp, obs = obs, 
                    exp_lats = lats.hcst, exp_lons = lons.hcst, obs_lats = lats.ref, obs_lons = lons.ref, 
                    target_grid = target_grid, bc_method = bc_method, int_method = int_method, 
                    points = NULL, method_point_interp = NULL, 
                    lat_dim = "latitude", lon_dim = "longitude", sdate_dim = "sdate", member_dim = "member", 
                    source_file = NULL, region = NULL, cal_method = "mse_min", ncores = 1)
    
    #Warning messages:
    #1: In Intbc(exp = exp, obs = obs, exp_lats = lats.hcst, exp_lons = lons.hcst,  :
    #The borders of the downscaling region have not been provided. Assuming the four borders of the downscaling region are defined by the first and last elements of the parametres 'obs_lats' and 'obs_lons'.
    #2: In s2dv::CDORemap(data_array = exp, lats = lats, lons = lons, grid = target_grid,  :
    #CDORemap: Using CDO version 1.9.8.
    #3: In s2dv::CDORemap(data_array = exp, lats = lats, lons = lons, grid = target_grid,  :
    #CDORemap: Using CDO version 1.9.8.
    #4: In BiasCorrection(exp = exp_interpolated$data, obs = obs_interpolated$data) :
    #Parameter 'obs' contains NA values.

    res2 <- Intbc(exp = exp, obs = obs, 
                    exp_lats = lats.hcst, exp_lons = lons.hcst, obs_lats = lats.ref, obs_lons = lons.ref, 
                    target_grid = target_grid, bc_method = bc_method, int_method = int_method, 
                    points = NULL, method_point_interp = NULL, 
                    lat_dim = "latitude", lon_dim = "longitude", sdate_dim = "sdate", member_dim = "member", 
                    source_file = NULL, region = NULL, cal_method = "evmos", ncores = 1) # different cal_method should not affect result (bc_int <- 'simple_bias')

    #> max(abs(res1$data - res2$data), na.rm = T)
    #[1] 0
    
    # test for quantile mapping
    bc_method <- 'quantile_mapping'
    res2bis <- Intbc(exp = exp, obs = obs,
                    exp_lats = lats.hcst, exp_lons = lons.hcst, obs_lats = lats.ref, obs_lons = lons.ref,
                    target_grid = target_grid, bc_method = bc_method, int_method = int_method,
                    points = NULL, method_point_interp = NULL,
                    lat_dim = "latitude", lon_dim = "longitude", sdate_dim = "sdate", member_dim = "ensemble",
                    source_file = NULL, region = NULL, ncores = 1)

# test with s2dv_cubes:
    names(dim(exp)) <- c('sdate', 'member', 'lat', 'lon')
    names(dim(obs)) <- c('sdate', 'lat', 'lon')
    names(dim(lats.hcst)) <- 'lat'
    names(dim(lats.ref)) <- 'lat'
    names(dim(lons.hcst)) <- 'lon'
    names(dim(lons.ref)) <- 'lon'
    exp_s2dvcube <- s2dv_cube(data = exp, lat = lats.hcst, lon = lons.hcst, Variable = list(varName = var, level = NULL), when = Sys.time(), source_files = fcst_path)
    obs_s2dvcube <- s2dv_cube(data = obs, lat = lats.ref, lon = lons.ref, Variable = list(varName = var, level = NULL), when = Sys.time(), source_files = ref_path)
    attr(exp_s2dvcube$lon,"first_lon") <- lons.min
    attr(exp_s2dvcube$lon,"last_lon") <- lons.max
    attr(exp_s2dvcube$lat,"first_lat") <- lats.min
    attr(exp_s2dvcube$lat,"last_lat") <- lats.max
    attr(obs_s2dvcube$lon,"first_lon") <- lons.min
    attr(obs_s2dvcube$lon,"last_lon") <- lons.max
    attr(obs_s2dvcube$lat,"first_lat") <- lats.min
    attr(obs_s2dvcube$lat,"last_lat") <- lats.max

    res3 <- CST_Intbc(exp = exp_s2dvcube, obs = obs_s2dvcube, 
                        target_grid = target_grid, bc_method = bc_method, int_method = int_method, 
                        points = NULL, method_point_interp = NULL, 
                        lat_dim = "lat", lon_dim = "lon", sdate_dim = "sdate", member_dim = "member", 
                        region = NULL, cal_method = "mse_min", ncores = 1)

    #Warning messages:
    #1: In Intbc(exp = exp$data, obs = obs$data, exp_lats = exp[[lat_dim]],  :
    #The borders of the downscaling region have not been provided. Assuming the four borders of the downscaling region are defined by the first and last elements of the parametres 'obs_lats' and 'obs_lons'.
    #2: In s2dv::CDORemap(data_array = exp, lats = lats, lons = lons, grid = target_grid,  :
    #CDORemap: Using CDO version 1.9.8.
    #3: In s2dv::CDORemap(data_array = exp, lats = lats, lons = lons, grid = target_grid,  :
    #CDORemap: Using CDO version 1.9.8.
    #4: In BiasCorrection(exp = exp_interpolated$data, obs = obs_interpolated$data) :
    #Parameter 'obs' contains NA values.

    max(abs(res1$data - res3$data), na.rm = T)
    #[1] 0
    # they are the same -> OK

# test several bc_methods:
    bc_method <- 'calibration'
    res4 <- CST_Intbc(exp = exp_s2dvcube, obs = obs_s2dvcube, 
                        target_grid = target_grid, bc_method = bc_method, int_method = int_method, 
                        points = NULL, method_point_interp = NULL, 
                        lat_dim = "lat", lon_dim = "lon", sdate_dim = "sdate", member_dim = "member", 
                        region = NULL, cal_method = "mse_min", ncores = 1)
    res5 <- CST_Intbc(exp = exp_s2dvcube, obs = obs_s2dvcube, 
                        target_grid = target_grid, bc_method = bc_method, int_method = int_method, 
                        points = NULL, method_point_interp = NULL, 
                        lat_dim = "lat", lon_dim = "lon", sdate_dim = "sdate", member_dim = "member", 
                        region = NULL, cal_method = "evmos", ncores = 1)

    bc_method <- 'quantile_mapping'
    res6 <- CST_Intbc(exp = exp_s2dvcube, obs = obs_s2dvcube, 
                        target_grid = target_grid, bc_method = bc_method, int_method = int_method, 
                        points = NULL, method_point_interp = NULL, 
                        lat_dim = "lat", lon_dim = "lon", sdate_dim = "sdate", member_dim = "member", 
                        region = NULL, cal_method = "mse_min", ncores = 1)

    #Error in approx(x = fobj$par$modq[, 1], y = fobj$par$fitq[, 1], xout = x[wet],  :
    #need at least two non-NA values to interpolate
    #In addition: Warning messages:
    #1: In Intbc(exp = exp$data, obs = obs$data, exp_lats = exp[[lat_dim]],  :
    #The borders of the downscaling region have not been provided. Assuming the four borders of the downscaling region are defined by the first and last elements of the parametres 'obs_lats' and 'obs_lons'.
    #2: In s2dv::CDORemap(data_array = exp, lats = lats, lons = lons, grid = target_grid,  :
    #CDORemap: Using CDO version 1.9.8.
    #3: In s2dv::CDORemap(data_array = exp, lats = lats, lons = lons, grid = target_grid,  :
    #CDORemap: Using CDO version 1.9.8.
    #4: In QuantileMapping(exp = exp_interpolated$data, obs = obs_interpolated$data,  :
    #Parameter 'obs' contains NA values.
    #5: In (function (exp, obs, exp_cor = NULL, sample_length = NULL, method = "QUANT",  :
    #The total sum of observed data sample in the sample number 1, is zero and the function may crash.
    obs_s2dvcube$data[which(is.na(obs_s2dvcube$data))] <- 0 # replace na with 0 because of the issue
    res6 <- CST_Intbc(exp = exp_s2dvcube, obs = obs_s2dvcube, 
                        target_grid = target_grid, bc_method = bc_method, int_method = int_method, 
                        points = NULL, method_point_interp = NULL, 
                        lat_dim = "lat", lon_dim = "lon", sdate_dim = "sdate", member_dim = "member", 
                        region = NULL, cal_method = "mse_min", ncores = 1)
    # same error, the issue is not in the NA of the data

# plot results:
    library(abind)
    library(ClimProjDiags)
    data_to_plot <- drop(MeanDims(MeanDims(res3$data, 'member'), 'sdate')) # res3 int_method = bil, bc_method = simple_bias
    data_to_plot <- abind(data_to_plot, drop(MeanDims(MeanDims(res4$data, 'member'), 'sdate')), along = 3) # res4 int_method = bil, bc_method = calibration, cal_method = mse_min
    data_to_plot <- abind(data_to_plot, drop(MeanDims(MeanDims(res5$data, 'member'), 'sdate')), along = 3) # res5 int_method = bil, bc_method = calibration, cal_method = evmos
    #data_to_plot <- abind(data_to_plot, drop(MeanDims(MeanDims(res6$data, 'member'), 'sdate')), along = 3) # res6 obs with 0 instead of NA; bc_method = quantile_mapping

    PlotLayout(fun = PlotEquiMap, # plot layout plots from left to right first and then top to bottom
            plot_dims = c('lon', 'lat'),
            var = ArrayToList(data_to_plot, 3, names = ''),
            lon = res1$lon-360,
            lat = res1$lat,
            colNA = 'black',
            brks = seq(284, 300, 1),
            #bar_scale = 0.5,
            #triangle_ends_scale = 1,
            #triangle_ends = c(TRUE,FALSE),
            intylat = 2.5,
            intxlon = 10,
            filled.continents = FALSE, 
            toptitle = 'downscaled field Intbc (ensemble mean climatology)',
            title_scale = 0.6,
            subtitle_scale = 0.5,
            nrow = 1,
            row_titles = c('Int_method = bil'),
            ncol = 3,
            col_titles = c('bc_method = simple_bias \n cal_method = mse_min', 'bc_method = calibration \n cal_method = mse_min', 'bc_method = calibration \n cal_method = evmos'),
            width = 12,
            height = 3*1, # 3*nrow
            fileout = paste0(output_dir, 'test4_intbc.png')
    )

    #> max(abs(res3$data - res4$data), na.rm = T)
    #Error in res3$data - res4$data : non-conformable arrays
    #> max(abs(res3$data - res5$data), na.rm = T)
    #Error in res3$data - res5$data : non-conformable arrays
    #> max(abs(res4$data - res5$data), na.rm = T)
    #[1] 20.08774
    #> max(abs(res4$data - res6$data), na.rm = T)
    #

    save(var, lats.min, lats.max, lons.min, lons.max, obs_s2dvcube, exp_s2dvcube, output_dir, file = paste0(output_dir, 'data_s2dvcubes_to_test_Intbc_and_Intlr.RData'))library(s2dv)
library(CSTools)
library(FNN)
library(multiApply)
source('/esarchive/scratch/allabres/gitlab_clones/csdownscale/R/Utils.R')
source('/esarchive/scratch/allabres/gitlab_clones/csdownscale/R/Interpolation.R')
source('/esarchive/scratch/allabres/gitlab_clones/csdownscale/R/Intbc.R')
output_dir <- '/esarchive/scratch/allabres/review_csdownscale/'

load(paste0(output_dir, 'daily_data_to_test_1993-2016_era5land.RData'))
exp <- drop(exp) #to remove unnecessary dimensions
obs <- drop(obs) #to remove unnecessary dimensions
target_grid <- attr(obs, 'Files')[1]

int_method <- 'bil' # "con", "bil", "bic", "nn", "con2"
bc_method <- 'simple_bias' # 'simple_bias', 'calibration', 'quantile_mapping'

# test with arrays:
    res <- Intbc(exp = exp, obs = obs, 
                    exp_lats = lats.hcst, exp_lons = lons.hcst, obs_lats = lats.ref, obs_lons = lons.ref, 
                    target_grid = target_grid, bc_method = bc_method, int_method = int_method, 
                    points = NULL, method_point_interp = NULL, 
                    lat_dim = "latitude", lon_dim = "longitude", sdate_dim = "sdate", member_dim = "ensemble", 
                    source_file = NULL, region = NULL, cal_method = "mse_min", ncores = 1)
    # Error in BiasCorrection(exp = exp_interpolated$data, obs = obs_interpolated$data) :
    #     Parameter 'exp' must have the dimensions 'member' and 'sdate'.

    names(dim(exp)) <- c('sdate', 'member', 'latitude', 'longitude')
    names(dim(obs)) <- c('sdate', 'latitude', 'longitude')
    res1 <- Intbc(exp = exp, obs = obs, 
                    exp_lats = lats.hcst, exp_lons = lons.hcst, obs_lats = lats.ref, obs_lons = lons.ref, 
                    target_grid = target_grid, bc_method = bc_method, int_method = int_method, 
                    points = NULL, method_point_interp = NULL, 
                    lat_dim = "latitude", lon_dim = "longitude", sdate_dim = "sdate", member_dim = "member", 
                    source_file = NULL, region = NULL, cal_method = "mse_min", ncores = 1)
    
    #Warning messages:
    #1: In Intbc(exp = exp, obs = obs, exp_lats = lats.hcst, exp_lons = lons.hcst,  :
    #The borders of the downscaling region have not been provided. Assuming the four borders of the downscaling region are defined by the first and last elements of the parametres 'obs_lats' and 'obs_lons'.
    #2: In s2dv::CDORemap(data_array = exp, lats = lats, lons = lons, grid = target_grid,  :
    #CDORemap: Using CDO version 1.9.8.
    #3: In s2dv::CDORemap(data_array = exp, lats = lats, lons = lons, grid = target_grid,  :
    #CDORemap: Using CDO version 1.9.8.
    #4: In BiasCorrection(exp = exp_interpolated$data, obs = obs_interpolated$data) :
    #Parameter 'obs' contains NA values.

    res2 <- Intbc(exp = exp, obs = obs, 
                    exp_lats = lats.hcst, exp_lons = lons.hcst, obs_lats = lats.ref, obs_lons = lons.ref, 
                    target_grid = target_grid, bc_method = bc_method, int_method = int_method, 
                    points = NULL, method_point_interp = NULL, 
                    lat_dim = "latitude", lon_dim = "longitude", sdate_dim = "sdate", member_dim = "member", 
                    source_file = NULL, region = NULL, cal_method = "evmos", ncores = 1) # different cal_method should not affect result (bc_int <- 'simple_bias')

    #> max(abs(res1$data - res2$data), na.rm = T)
    #[1] 0

# test with s2dv_cubes:
    names(dim(exp)) <- c('sdate', 'member', 'lat', 'lon')
    names(dim(obs)) <- c('sdate', 'lat', 'lon')
    names(dim(lats.hcst)) <- 'lat'
    names(dim(lats.ref)) <- 'lat'
    names(dim(lons.hcst)) <- 'lon'
    names(dim(lons.ref)) <- 'lon'
    exp_s2dvcube <- s2dv_cube(data = exp, lat = lats.hcst, lon = lons.hcst, Variable = list(varName = var, level = NULL), when = Sys.time(), source_files = fcst_path)
    obs_s2dvcube <- s2dv_cube(data = obs, lat = lats.ref, lon = lons.ref, Variable = list(varName = var, level = NULL), when = Sys.time(), source_files = ref_path)
    attr(exp_s2dvcube$lon,"first_lon") <- lons.min
    attr(exp_s2dvcube$lon,"last_lon") <- lons.max
    attr(exp_s2dvcube$lat,"first_lat") <- lats.min
    attr(exp_s2dvcube$lat,"last_lat") <- lats.max
    attr(obs_s2dvcube$lon,"first_lon") <- lons.min
    attr(obs_s2dvcube$lon,"last_lon") <- lons.max
    attr(obs_s2dvcube$lat,"first_lat") <- lats.min
    attr(obs_s2dvcube$lat,"last_lat") <- lats.max

    res3 <- CST_Intbc(exp = exp_s2dvcube, obs = obs_s2dvcube, 
                        target_grid = target_grid, bc_method = bc_method, int_method = int_method, 
                        points = NULL, method_point_interp = NULL, 
                        lat_dim = "lat", lon_dim = "lon", sdate_dim = "sdate", member_dim = "member", 
                        region = NULL, cal_method = "mse_min", ncores = 1)

    #Warning messages:
    #1: In Intbc(exp = exp$data, obs = obs$data, exp_lats = exp[[lat_dim]],  :
    #The borders of the downscaling region have not been provided. Assuming the four borders of the downscaling region are defined by the first and last elements of the parametres 'obs_lats' and 'obs_lons'.
    #2: In s2dv::CDORemap(data_array = exp, lats = lats, lons = lons, grid = target_grid,  :
    #CDORemap: Using CDO version 1.9.8.
    #3: In s2dv::CDORemap(data_array = exp, lats = lats, lons = lons, grid = target_grid,  :
    #CDORemap: Using CDO version 1.9.8.
    #4: In BiasCorrection(exp = exp_interpolated$data, obs = obs_interpolated$data) :
    #Parameter 'obs' contains NA values.

    max(abs(res1$data - res3$data), na.rm = T)
    #[1] 0
    # they are the same -> OK

# test several bc_methods:
    bc_method <- 'calibration'
    res4 <- CST_Intbc(exp = exp_s2dvcube, obs = obs_s2dvcube, 
                        target_grid = target_grid, bc_method = bc_method, int_method = int_method, 
                        points = NULL, method_point_interp = NULL, 
                        lat_dim = "lat", lon_dim = "lon", sdate_dim = "sdate", member_dim = "member", 
                        region = NULL, cal_method = "mse_min", ncores = 1)
    res5 <- CST_Intbc(exp = exp_s2dvcube, obs = obs_s2dvcube, 
                        target_grid = target_grid, bc_method = bc_method, int_method = int_method, 
                        points = NULL, method_point_interp = NULL, 
                        lat_dim = "lat", lon_dim = "lon", sdate_dim = "sdate", member_dim = "member", 
                        region = NULL, cal_method = "evmos", ncores = 1)

    bc_method <- 'quantile_mapping'
    res6 <- CST_Intbc(exp = exp_s2dvcube, obs = obs_s2dvcube, 
                        target_grid = target_grid, bc_method = bc_method, int_method = int_method, 
                        points = NULL, method_point_interp = NULL, 
                        lat_dim = "lat", lon_dim = "lon", sdate_dim = "sdate", member_dim = "member", 
                        region = NULL, cal_method = "mse_min", ncores = 1)

    #Error in approx(x = fobj$par$modq[, 1], y = fobj$par$fitq[, 1], xout = x[wet],  :
    #need at least two non-NA values to interpolate
    #In addition: Warning messages:
    #1: In Intbc(exp = exp$data, obs = obs$data, exp_lats = exp[[lat_dim]],  :
    #The borders of the downscaling region have not been provided. Assuming the four borders of the downscaling region are defined by the first and last elements of the parametres 'obs_lats' and 'obs_lons'.
    #2: In s2dv::CDORemap(data_array = exp, lats = lats, lons = lons, grid = target_grid,  :
    #CDORemap: Using CDO version 1.9.8.
    #3: In s2dv::CDORemap(data_array = exp, lats = lats, lons = lons, grid = target_grid,  :
    #CDORemap: Using CDO version 1.9.8.
    #4: In QuantileMapping(exp = exp_interpolated$data, obs = obs_interpolated$data,  :
    #Parameter 'obs' contains NA values.
    #5: In (function (exp, obs, exp_cor = NULL, sample_length = NULL, method = "QUANT",  :
    #The total sum of observed data sample in the sample number 1, is zero and the function may crash.
    obs_s2dvcube$data[which(is.na(obs_s2dvcube$data))] <- 0 # replace na with 0 because of the issue
    res6 <- CST_Intbc(exp = exp_s2dvcube, obs = obs_s2dvcube, 
                        target_grid = target_grid, bc_method = bc_method, int_method = int_method, 
                        points = NULL, method_point_interp = NULL, 
                        lat_dim = "lat", lon_dim = "lon", sdate_dim = "sdate", member_dim = "member", 
                        region = NULL, cal_method = "mse_min", ncores = 1)
    # same error, the issue is not in the NA of the data

# plot results:
    library(abind)
    library(ClimProjDiags)
    data_to_plot <- drop(MeanDims(MeanDims(res3$data, 'member'), 'sdate')) # res3 int_method = bil, bc_method = simple_bias
    data_to_plot <- abind(data_to_plot, drop(MeanDims(MeanDims(res4$data, 'member'), 'sdate')), along = 3) # res4 int_method = bil, bc_method = calibration, cal_method = mse_min
    data_to_plot <- abind(data_to_plot, drop(MeanDims(MeanDims(res5$data, 'member'), 'sdate')), along = 3) # res5 int_method = bil, bc_method = calibration, cal_method = evmos
    #data_to_plot <- abind(data_to_plot, drop(MeanDims(MeanDims(res6$data, 'member'), 'sdate')), along = 3) # res6 obs with 0 instead of NA; bc_method = quantile_mapping

    PlotLayout(fun = PlotEquiMap, # plot layout plots from left to right first and then top to bottom
            plot_dims = c('lon', 'lat'),
            var = ArrayToList(data_to_plot, 3, names = ''),
            lon = res1$lon-360,
            lat = res1$lat,
            colNA = 'black',
            brks = seq(284, 300, 1),
            #bar_scale = 0.5,
            #triangle_ends_scale = 1,
            #triangle_ends = c(TRUE,FALSE),
            intylat = 2.5,
            intxlon = 10,
            filled.continents = FALSE, 
            toptitle = 'downscaled field Intbc (ensemble mean climatology)',
            title_scale = 0.6,
            subtitle_scale = 0.5,
            nrow = 1,
            row_titles = c('Int_method = bil'),
            ncol = 3,
            col_titles = c('bc_method = simple_bias \n cal_method = mse_min', 'bc_method = calibration \n cal_method = mse_min', 'bc_method = calibration \n cal_method = evmos'),
            width = 12,
            height = 3*1, # 3*nrow
            fileout = paste0(output_dir, 'test4_intbc.png')
    )

    #> max(abs(res3$data - res4$data), na.rm = T)
    #Error in res3$data - res4$data : non-conformable arrays
    #> max(abs(res3$data - res5$data), na.rm = T)
    #Error in res3$data - res5$data : non-conformable arrays
    #> max(abs(res4$data - res5$data), na.rm = T)
    #[1] 20.08774
    #> max(abs(res4$data - res6$data), na.rm = T)
    #

    save(var, lats.min, lats.max, lons.min, lons.max, obs_s2dvcube, exp_s2dvcube, output_dir, file = paste0(output_dir, 'data_s2dvcubes_to_test_Intbc_and_Intlr.RData'))
