library(s2dverification)
library(startR)
library(RColorBrewer)
devtools::install_git('https://earth.bsc.es/gitlab/external/medscope', branch = 'master')
library(medscope)
devtools::install_git('https://earth.bsc.es/gitlab/ces/multiApply', branch = 'master')
library(multiApply)
mth = '11'
temp = 'tas'
precip = 'prlr'
ini <- 1992
fin <- 2012
start <- as.Date(paste(ini, mth, "01", sep = ""), "%Y%m%d")
end <- as.Date(paste(fin, mth, "01", sep = ""), "%Y%m%d")
dateseq <- format(seq(start, end, by = "year"), "%Y%m%d")
grid <- "256x128"
obs <- "erainterim"
load(file = "/home/Earth/nperez/git/medscope/tas_prlr_toydata.RData")
dim(Temp$mod)
Lat <- Temp$lat
Lon <- Temp$lon
StartData <- list(Temp, Precip)
weight <- c(2, 1)

mvrmse <- MultivarRMSE(StartData, weight)

PlotEquiMap(mvrmse$mvrmse, lon = Lon, lat = Lat, filled.continents = FALSE, 
            toptitle = "Multivariate RMSE tas, prlr 1992 - 2012", colNA = "white", 
            bar_limits = c(0,2.5), cols = brewer.pal(n=5,name='Reds'))
#expected: mod1, mod2, ..., obs1, obs2, ....
#where mod1, mod2 ... can be arrays of whatever number of memebers and start dates and ftimes


.MultivarRMSE  <- function(..., weight) {
  data <- list(...)
  if (is.list(unlist(data, recursive = FALSE))) {
    data <- unlist(data, recursive = FALSE)
  }
  #checks 
  if (length(data) %% 2 !=0) {
    stop("A even number of vectors is expected: one containing the forecast
           and other for the observations for each variable.")  
  }
  if (!all(unlist(lapply(data, is.numeric)))) {
    stop("The data must be numeric.")
  }
  num_vars <- length(data) / 2
  names_dim <- lapply(data, function(x){names(dim(x))})
  if (!all(unlist(lapply(names_dim[1 : num_vars], 
                         function(x){x == names_dim[[1]]})))) {
    stop("Dimension names in experiment data must be equal for all variables.")
  }
  if (!all(unlist(lapply(names_dim[(num_vars + 1) : (num_vars * 2)],
                         function(x){x == names_dim[[num_vars + 1]]})))) {
    stop("Dimension names in observed data must be equal for all variables.") 
  }
  pos_sdate_exp <- which(names_dim[[1]] == 'sdate')
  pos_sdate_obs <- which(names_dim[[num_vars + 1]] == 'sdate')
  pos_ftime_exp <- which(names_dim[[1]] == 'ftime')
  pos_ftime_obs <- which(names_dim[[num_vars + 1]] == 'ftime')
  if (length(pos_sdate_exp) == 0 | length(pos_sdate_obs) == 0) {
    stop("A start date dimension called 'sdate' is expected in the 
          experimental and observational arrays.")
  }
  if (length(pos_ftime_exp) == 0 | length(pos_ftime_obs) == 0) {
    stop("A forecast time dimension called 'ftime' is expected in the 
          experimental and observational arrays.")
  }
  if (!all(unlist(lapply(data[1 : num_vars], function(x){
                  dim(x)[pos_sdate_exp] == dim(data[[1]])[pos_sdate_exp]})))) {
    stop("The dimension 'sdate' must be the same length for all forecast 
          variables.")
  }
  if (!all(unlist(lapply(data[(num_vars + 1) : (num_vars * 2)], function(x) {
      dim(x)[pos_sdate_obs] == dim(data[[num_vars + 1]])[pos_sdate_obs]})))) {
    stop("The dimension sdate' must be the same length for all observed
          variables.")
  }
  if (dim(data[[1]])[pos_sdate_exp] != dim(data[[num_vars + 1]])[pos_sdate_obs]) {
    stop("Dimension 'sdate' must be the same length for the forecast and the 
          observed data.")
  }
  if (dim(data[[1]])[pos_ftime_exp] != dim(data[[num_vars + 1]])[pos_ftime_obs]) {
    stop("Dimension 'ftime' must be the same length for the forecast and the 
         observed data.")
  }
 num_sdates <- dim(data[[1]])[pos_sdate_exp]
 num_ftime <- dim(data[[1]])[pos_ftime_exp]
   if (!is.vector(weight)) {
    stop("Parameter 'weight' must be a vector.")
  }

  if (num_vars * 2 != length(data)) {
    stop("Parameter 'weight' must be the same length as the number of 
          variables.")
  }
#Ano_CrossValid(data[[var]], data[[num_vars + var]]) #no funciona porque necesita 4 dim
 #calculating the anomaly using Ano_CrossValid for each var
 dim_exp <- dim(data[[num_vars]])
 dim_obs <- dim(data[[num_vars + 1]])
 dim_exp[pos_sdate_exp] <- 1
 dim_obs[pos_sdate_obs] <- 1
 mvrmse <- 0
 sumweights <- 0
  for (var in 1 : num_vars) {
    ano_exp <- array(dim = dim_exp)
    ano_obs <- array(dim = dim_obs)
    for (jsdate in 1 : num_sdates) {
      exp <- abind::asub(data[[var]], idx = -jsdate, dims = pos_sdate_exp)
      obs <-  abind::asub(data[[var + num_vars]], idx = -jsdate, dims = pos_sdate_obs)
      #calculates the Clim between two arrays of at least 4 dim: returns the clim of obs and mod
      clim_exp <- Mean1Dim(exp, posdim = pos_sdate_exp, narm = TRUE)
      clim_obs <- Mean1Dim(obs, posdim = pos_sdate_obs, narm = TRUE)
      #remove the clim obtained to the modelled and the observed data    
      exp <- abind::asub(data[[var]], idx = jsdate, dims = pos_sdate_exp) - clim_exp
      obs <-  abind::asub(data[[var + num_vars]], idx = jsdate, dims = pos_sdate_obs) - clim_obs
      ano_exp <- abind::abind(ano_exp, exp, along = pos_sdate_exp)
      ano_obs <- abind::abind(ano_obs, obs, along = pos_sdate_obs)
    }
    ano_exp <- abind::asub(ano_exp, idx = -1, dims = pos_sdate_exp)
    names(dim(ano_exp)) <- names_dim[[1]]
    ano_obs <- abind::asub(ano_obs, idx = -1, dims = pos_sdate_obs)
    names(dim(ano_obs)) <- names_dim[[num_vars + 1]]
    #calculating the mean using Mean1Dim for each var and obs or mod
    avg_exp <- apply(ano_exp, pos_sdate_exp, mean, na.rm = TRUE)
    dim(avg_exp) <- dim(data[[var]])[pos_sdate_exp]
    avg_obs <- apply(ano_obs, pos_sdate_obs, mean, na.rm = TRUE)
    dim(avg_obs) <- dim(data[[num_vars + 1]])[pos_sdate_obs]
    #calculating the rmse for each var
    #rmse <- s2dverification::RMS(AvgExp, AvgObs, posloop = 1, 
    #                            posRMS = 2, conf = FALSE)
    
    dif <- avg_exp - avg_obs
    rmse <- mean(dif ** 2, na.rm = TRUE) ** 0.5
    #calculating the standard deviation for obs of each variable
    stdev <- sd(avg_obs)
    #sum of the weighted rmse for all variables
    mvrmse <- mvrmse + (rmse / stdev * as.numeric(weight[var]))
    sumweights <- sumweights + as.numeric(weight[var])
  }
 mvrmse <- mvrmse/sumweights
 return(avg_obs)
}

d1 = 1 : (2 * 3 * 4)
dim(d1) <- c(ftime = 2, sdate = 3, member = 4)
d2 = seq(1, (2 * 3 * 4) * 2, 2) 
dim(d2) <- c(ftime = 2, sdate = 3, member = 4)
obs1 = 1 : (2 * 3)
dim(obs1) <- c(ftime = 2, sdate = 3)
obs2 = rnorm(6, m = 10)
dim(obs2) <- c(ftime = 2, sdate = 3)
data <- list(d1, d2, obs1, obs2)
a <- .MultivarRMSE(d1, d2, obs1, obs2, weight = c(2, 1))           

length(a)
a
is.numeric(a)
length(a)

datos <- list(StartData[[1]]$mod, StartData[[2]]$mod, StartData[[1]]$obs, StartData[[2]]$obs)
Apply(datos, target_dims = list('lat', 'lon'), fun = .MultivarRMSE)

data <- list(datos)
MultivarRMSE_v2 <- function(..., weight) {
  data <- list(...)
  if (is.list(unlist(data, recursive = FALSE))) {
    data <- unlist(data, recursive = FALSE)
  }
  #checks 
  if (length(data) %% 2 !=0) {
    stop("A even number of vectors is expected: one containing the forecast
           and other for the observations for each variable.")  
  }
  if (!all(unlist(lapply(data, is.numeric)))) {
    stop("The data must be numeric.")
  }
  num_vars <- length(data) / 2
  names_dim <- lapply(data, function(x){names(dim(x))})
  if (!all(unlist(lapply(names_dim[1 : num_vars], 
                         function(x){x == names_dim[[1]]})))) {
    stop("Dimension names in experiment data must be equal for all variables.")
  }
  if (!all(unlist(lapply(names_dim[(num_vars + 1) : (num_vars * 2)],
                         function(x){x == names_dim[[num_vars + 1]]})))) {
    stop("Dimension names in observed data must be equal for all variables.") 
  }
  pos_lon_exp <- which(names_dim[[1]] == 'lon' | names_dim[[1]] == 'longitud')
  pos_lon_obs <- which(names_dim[[num_vars + 1]] == 'lon' | 
                         names_dim[[num_vars + 1]] == 'longitud')
  pos_lat_exp <- which(names_dim[[1]] == 'lat' | names_dim[[1]] == 'latitud')
  pos_lat_obs <- which(names_dim[[num_vars + 1]] == 'lat' | names_dim[[num_vars + 1]] == 'latitud') 
  if (length(pos_lon_exp) == 0 | length(pos_lon_obs) == 0) {
    stop("Observed and experiment data must have a longitudinal dimension.")
  }
  if (length(pos_lat_exp) == 0 | length(pos_lat_obs) == 0) {
    stop("Observed and experiment data must have a longitudinal dimension.")
  }
  if (dim(data[[1]])[pos_lon_exp] != dim(data[[num_vars + 1]])[pos_lon_obs]) {
    stop("Longitudinal dimension in experiment and observed data must be the 
          same length.")
  }
  num_lon <- dim(data[[1]])[pos_lon_exp]
  if (dim(data[[1]])[pos_lat_exp] != dim(data[[num_vars + 1]])[pos_lat_obs]) {
    stop("Latitudinal dimension in experiment and observed data must be the 
          same length.")
  }
  num_lat <- dim(data[[1]])[pos_lat_exp]
  # Check dimensions length 1
  re_dim_exp <- which(dim(data[[1]]) == 1)
  re_dim_obs <- which(dim(data[[num_vars + 1]]) == 1)
  output <- NULL
  for (i in 1 : num_lon) {
    loni_exp <- lapply(data[1 : num_vars], function(x) {asub(x, idx = i, 
                                                        dims = pos_lon_exp)})
    loni_obs <- lapply(data[(num_vars + 1) : (num_vars * 2)], function(x) {
                asub(x, idx = i, dims = pos_lon_obs)})
    out <- NULL
    for (j in 1 : num_lat) {
      if (pos_lat_exp > pos_lon_exp) {
        pos_lat_exp <- pos_lat_exp - 1
      }
      latj_exp <- lapply(loni_exp, function(x) {asub(x, idx = j, 
                                    dims = pos_lat_exp - length(re_dim_exp))})
      
      latj_obs <- lapply(loni_obs, function(x) {asub(x, idx = j, 
                                    dims = pos_lat_obs - length(re_dim_obs))})
      for (k in 1 : length(latj_exp)) {
        names(dim(latj_exp[[k]])) <- names_dim[[1]][-c(re_dim_exp, pos_lon_exp, pos_lat_exp)]
      }
      for (k in 1 : length(latj_obs)) {
        names(dim(latj_obs[[k]])) <- names_dim[[num_vars + 1]][-c(re_dim_obs, pos_lon_obs, pos_lat_obs)]
      }
      data_list <- lapply(c(latj_exp, latj_obs), function(x){x})
      out <- c(out, .MultivarRMSE(data_list, weight = weight))
    }
    output <- rbind(output, out)
  }
  names(dim(output)) <- c("lon", "lat")
  return(output)
}
data_plot <- MultivarRMSE_v2(datos, weight = c(2, 1))

str(mvrmse)
str(data_plot)
str(mvrmse$AvgObs[1,1,1,])

PlotEquiMap(data_plot, lon = Lon, lat = Lat, filled.continents = FALSE, 
            toptitle = "Multivariate RMSE tas, prlr 1992 - 2012", colNA = "white", 
            bar_limits = c(0,2.5), cols = brewer.pal(n=5,name='Reds'))
