## TEST INTERPOLATION FUNCTION (FIRST TEST)

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')
output_dir <- '/esarchive/scratch/jramon/downscaling/tests_alba'
input_dir <- '/esarchive/scratch/allabres/review_csdownscale/'

# reproduce example in the documentation
    exp <- rnorm(500)
    dim(exp) <- c(member = 5, lat = 4, lon = 5, sdate = 5, time  = 1)
    lons <- 1:5
    lats <- 1:4
    exp <- s2dv_cube(data = exp, lat = lats, lon = lons) 
    res <- CST_Interpolation(exp = exp, method_remap = 'conservative', target_grid = 'r1280x640')
    
    #> dim(exp$data)
    #member    lat    lon  sdate   time
    #     5      4      5      5      1
    #> dim(res$data)
    #member    lat    lon  sdate   time
    #    5     10     14      5      1

    # interpolate to a set of points:
   # exp$source_files <- '/esarchive/exp/ecmwf/system5c3s/daily_mean/tas_f6h/tas_20160501.nc' # alternatively give set the parameter "source_file" when calling the function (needed if the downscaling is to a point location)
   # res2 <- CST_Interpolation(exp = exp, method_point_interp = 'nearest', points = list(lat = c(2,3), lon = c(4,5)))
   # not working with fake data

# test with real data from one grid to another
    load(paste0(input_dir, 'daily_data_to_test_1993-2016_era5land.RData'))
    target_grid <- attr(obs, 'Files')[1]
    rm(obs, lats.ref, lons.ref)
    
    names(dim(exp)) <- replace(names(dim(exp)), which(names(dim(exp))=='longitude'), 'lon')
    names(dim(exp)) <- replace(names(dim(exp)), which(names(dim(exp))=='latitude'), 'lat')
    names(dim(lats.hcst)) <- 'lat'
    names(dim(lons.hcst)) <- 'lon'
    test_data <- s2dv_cube(data = exp, lat = lats.hcst, lon = lons.hcst, Variable = list(varName = var, level = NULL), when = Sys.time(), source_files = attr(exp, 'Files')[1])
    #attr(test_data$lon,"first_lon") <- lons.min
    #attr(test_data$lon,"last_lon") <- lons.max
    #attr(test_data$lat,"first_lat") <- lats.min
    #attr(test_data$lat,"last_lat") <- lats.max

    downscaled <- CST_Interpolation(exp = test_data, method_remap = 'conservative', target_grid = target_grid)

    #> dim(exp)
    #    dat       var     sdate  ensemble      time  latitude longitude
    #    1         1        24        25         1        15        27
    #> dim(downscaled$data)
    #    dat      var    sdate ensemble     time      lat      lon
    #    1        1       24       25        1      141      261

    #data_to_plot = abind(drop(MeanDims(MeanDims(exp, 'ensemble'), 'sdate')), drop(MeanDims(MeanDims(downscaled$data, 'ensemble'), 'sdate')), along = 3)
    #different lat/lon

    original_data_to_plot <- drop(MeanDims(MeanDims(exp, 'ensemble'), 'sdate'))
    same_brks <- seq(as.integer(quantile(original_data_to_plot, 0.01)), as.integer(quantile(original_data_to_plot, 0.99)))
    PlotEquiMap(original_data_to_plot, 
                brks = same_brks,
                lon = lons.hcst, lat = lats.hcst, colNA = 'black', filled.continents = FALSE,
                title_scale = 0.7, intxlon = 5, intylat = 2,
                toptitle = 'Original field ensemble mean Climatology', 
                file = paste0(output_dir, 'test3_interpolation_test_original.png'))

    lons <- downscaled$lon
    if(lons[1] > 180) {
        lons <- lons - 360
    }
    lats <- downscaled$lats
    
    ###########################################################################################################################
    save(test_data, downscaled, lons, lats, lons.min, lons.max, lats.min, lats.max, same_brks, file = paste0(output_dir, 'test3_interolation_test_regrid.RData'))
    ###########################################################################################################################
    
    methods <- c("con", "bil", "bic", "nn", "con2")
    for(method_remap in methods){
        downscaled <- CST_Interpolation(exp = test_data, method_remap = method_remap, target_grid = target_grid)
        PlotEquiMap(drop(MeanDims(MeanDims(downscaled$data, 'ensemble'), 'sdate')), 
                brks = same_brks,
                lon = lons, lat = downscaled$lat, colNA = 'black', filled.continents = FALSE, 
                title_scale = 0.7, intxlon = 5, intylat = 2,
                toptitle = paste0('Downscaled field ensemble mean Climatology (', method_remap, ')'), 
                file = paste0(output_dir, 'test3_interpolation_test_downscaled_', method_remap, '.png'))
                print(paste0('saved plot of method *', method_remap, '* (', which(methods == method_remap), ' out of ', length(methods), ')'))
    }
    
    
# test with real data from one grid to a set of point locations
    load(paste0(output_dir, 'test3_interolation_test_regrid.RData'))
    selected_points <- list(lat = c(-15.5), lon = c(-60.5))
    selected_points <- list(lat = c(-18, -18, -10, -14), lon = c(-60, -55, -45, -50))
    downscaled_points <- CST_Interpolation(exp = test_data, method_point_interp = 'bilinear', points = selected_points)
    # ERROR At least one of the points lies outside the model region -> the problem is that the region is not defined 
    downscaled_points <- CST_Interpolation(exp = test_data, method_point_interp = 'bilinear', points = selected_points, region = c(lons.min, lons.max, lats.min, lats.max))
    #ERROR:
        #cdo    sellonlatbox: Processed 1 variable over 216 timesteps [66.81s 809MB].
        #cdo    invertlat: Processed 1 variable over 216 timesteps [1.29s 809MB].
        #Error in laply(1:length(is), function(k) { :
        #could not find function "laply"
        #https://earth.bsc.es/gitlab/es/csdownscale/-/blob/dev-Rpackage/R/Interpolation.R#L687
        #laply instead of lapply
        #NO, the problem is that I need to load the library (library(plyr))

    library(plyr)
    downscaled_points <- CST_Interpolation(exp = test_data, method_point_interp = 'bilinear', points = selected_points, region = c(lons.min, lons.max, lats.min, lats.max))
   
    #cdo    sellonlatbox: Processed 1 variable over 216 timesteps [70.87s 450MB].
    #cdo    invertlat: Processed 1 variable over 216 timesteps [1.24s 456MB].
    #>    dim(downscaled_points$data)
    #location      dat      var    sdate ensemble     time
    #    4        1        1       24       25        1

    #-------------------
    # tests by Jaume
    #-------------------
    # order latitudes from lowest to highest (-90 to 90):
    test_data$lat <- test_data$lat[order(test_data$lat)]     
    test_data$data <- test_data$data[,,,,,order(test_data$lat),,drop = F]
    # since longitudes range from 0 to 360 in the original files, this should be the convention used
    # in test_data and selected_points
    selected_points$lon <- selected_points$lon + 360
    test_data$lon <- test_data$lon + 360
    downscaled_points <- CST_Interpolation(exp = test_data, method_point_interp = 'invdist4nn', points = selected_points)
    # success
    save(downscaled_points, file = paste0(output_dir, 'test3_downscaled_points.RData'))

    #plots:
        load(paste0(output_dir, 'test3_interolation_test_regrid.RData'))
        var <- 'tas'
        methods <- c("nearest", "bilinear", "9point", "invdist4nn", "NE", "NW", "SE", "SW")
        #selected_points <- list(lat = c(-18, -18, -10, -14, -6.4, -8.5), lon = c(-60, -55, -45, -50, -64.9, -53.2))
        
        # will only work if lat and lon of the exp field are CONSECUTIVE INTEGERS
            positions_lats <- 0
            positions_lons <- 0
            for(l in seq(1, length(selected_points$lat),1)){
                positions_lats <- append(positions_lats, which(test_data$lat == as.integer(selected_points$lat[l])))
                positions_lons <- append(positions_lons, which(test_data$lon == as.integer(selected_points$lon[l])))
            }
            positions_lats <- positions_lats[-1]
            positions_lons <- positions_lons[-1]
            original <- MeanDims(MeanDims(drop(test_data$data), 'ensemble'), 'sdate')
            na_structure <- matrix(NA, nrow = dim(original)[1], ncol = dim(original)[2])
            test <- na_structure
            for(p in seq(1,length(positions_lats),1)){
                test[positions_lats[p], positions_lons[p]] <- original[positions_lats[p], positions_lons[p]]
            }
            PlotEquiMap(test, 
                    brks = same_brks,
                    triangle_ends = c(TRUE, TRUE),
                    lon = test_data$lon, lat = test_data$lat, colNA = 'white', filled.continents = FALSE, 
                    title_scale = 0.7, intxlon = 5, intylat = 2,
                    toptitle = paste0('Downscaled points with ', m), 
                    file = paste0(output_dir, 'test3_interpolation_test_points_original.png'))
           

        library(plyr)
        for(m in methods){
            downscaled_points <- CST_Interpolation(exp = test_data, method_point_interp = m, points = selected_points, 
                                                    region = c(lons.min, lons.max, lats.min, lats.max))
          
            data_shaped <- as.data.frame.table(drop(downscaled_points$data))
            dimnames(data_shaped)[[2]] <- c('location', 'sdate', 'ensemble', var)
    
            png(paste0(output_dir, 'test3_interpolation_to_points_', m, '.png'))
                plot(data_shaped$tas~data_shaped$location, ylab = 'temperature (K)', xlab = 'selected points', main = m)
            dev.off()
            
            test <- MeanDims(MeanDims(drop(test_data$data), 'ensemble'), 'sdate')
            na_structure <- matrix(NA, nrow = dim(test)[1], ncol = dim(test)[2])
            test <- na_structure
            for(p in seq(1,length(positions_lats),1)){
                test[positions_lats[p], positions_lons[p]] <- MeanDims(MeanDims(drop(downscaled_points$data), 'ensemble'), 'sdate')[p]
            }
            PlotEquiMap(test, 
                    brks = same_brks,
                    triangle_ends = c(TRUE, TRUE),
                    lon = test_data$lon, lat = test_data$lat, colNA = 'white', filled.continents = FALSE, 
                    title_scale = 0.7, intxlon = 5, intylat = 2,
                    toptitle = paste0('Downscaled points with ', m), 
                    file = paste0(output_dir, 'test3_interpolation_test_points_', m, '.png'))

          
            print(paste0('saved plots of method *', m, '* (', which(m == methods), ' out of ', length(methods), ')'))

        }

        #ERROR with invdist4nn        
        #cdo    sellonlatbox: Processed 1 variable over 216 timesteps [69.24s 1235MB].
        #cdo    invertlat: Processed 1 variable over 216 timesteps [1.23s 1235MB].
        #[1] "saved plots of method *nearest* (1 out of 8)"
        #cdo    sellonlatbox: Processed 1 variable over 216 timesteps [66.20s 561MB].
        #cdo    invertlat: Processed 1 variable over 216 timesteps [1.18s 561MB].
        #[1] "saved plots of method *bilinear* (2 out of 8)"
        #cdo    sellonlatbox: Processed 1 variable over 216 timesteps [66.43s 561MB].
        #cdo    invertlat: Processed 1 variable over 216 timesteps [1.31s 561MB].
        #[1] "saved plots of method *9point* (3 out of 8)"
        #cdo    sellonlatbox: Processed 1 variable over 216 timesteps [65.22s 561MB].
        #cdo    invertlat: Processed 1 variable over 216 timesteps [1.25s 561MB].
        #        ida      idb      idc       idd
        #[1,]      Inf 0.000000 0.000000 0.0000000
        #[2,]      Inf 0.000000 0.000000 0.0000000
        #[3,]      Inf 0.000000 0.000000 0.0000000
        #[4,]      Inf 0.000000 0.000000 0.0000000
        #[5,] 1.643990 2.425356 1.015346 0.9245003
        #[6,] 1.059998 1.059998 1.856953 1.8569534
        #Error in match.fun(FUN) : argument "FUN" is missing, with no default
        #> m
        #[1] "invdist4nn"
