
library(ClimProjDiags)
library(CSTools)
library(SPEI)

P_PET_Accum <- function(tasmin, tasmax, pr, lat, pet_method, accum_scales, units_tas, units_pr, year_dim, month_dim){
  
  if (anyNA(tasmin) | anyNA(tasmax) | anyNA(pr)){
    
    accum <- list()
    for (accum_scale in accum_scales){accum[[paste0('accum',accum_scale)]] <- NA * s2dv::Reorder(data = tasmin, order = c(month_dim,year_dim))}
    
  } else {
    
    ######################################
    #####----- Correcting units -----#####
    ######################################
    
    ## tasmin and tasmax -> ºC
    if (identical(units_tas,'K')){
      tasmin <- tasmin - 273.15
      tasmax <- tasmax - 273.15
    } else if (!identical(units_tas,'C')){
      stop('Incorrect units for tasmin and tasmax')
    }
    
    ## pr -> mm day-1 (monthly total)
    if (identical(units_pr,'m s-1')){
      pr <- pr * 60*60*24*1000 * 30.4375
    } else if (identical(units_pr,'kg m-2 s-1')){
      pr <- pr * 60*60*24 * 30.4375
    } else if (!identical(units_pr,'mm day-1')){
      stop('Incorrect units for pr')
    }
    
    ########################################
    #####----- Merging dimensions -----#####
    ########################################
    
    if (!is.null(year_dim)){
      tasmin <- CSTools::MergeDims(data = tasmin, merge_dims = c(month_dim, year_dim), rename_dim = month_dim, na.rm = FALSE)
      tasmax <- CSTools::MergeDims(data = tasmax, merge_dims = c(month_dim, year_dim), rename_dim = month_dim, na.rm = FALSE)
      pr <- CSTools::MergeDims(data = pr, merge_dims = c(month_dim, year_dim), rename_dim = month_dim, na.rm = FALSE)
    }
    
    ####################################
    #####----- PET estimation -----#####
    ####################################
    
    ## "hargreaves" method computes the monthly reference evapotranspiration (ET0) of a grass crop 
    ## based on the original Hargreaves equation (1994). However, if precipitation data Pre is provided
    ## a modified form due to Droogers and Allen (2002) will be used; this equation corrects ET0 using 
    ## the amount of rain of each month as a proxy for insolation. The Hargreaves method requires data 
    ## on the mean external radiation, Ra. If such data are not available it can be estimated from the
    ## latitude lat and the month of the year.
    
    if (is.null(lat)){lat <- as.vector(attr(data,'Variables')$common[['lat']])}
    stopifnot(length(lat) == 1)
    
    if (identical(pet_method,'hargreaves')){
      pet <- SPEI::hargreaves(Tmin = as.vector(tasmin), Tmax = as.vector(tasmax), Ra = NULL, lat = lat, Pre = NULL, na.rm = TRUE)
    } else {stop('Incorrect pet_method')}
    
    ##################################
    #####----- Accumulation -----#####
    ##################################
    
    diff_pr_pet <- pr - pet
    
    accum <- list()
    for (accum_scale in accum_scales){
      
      ## P-PET accumulation
      accum[[paste0('accum',accum_scale)]] <- zoo::rollapply(data = diff_pr_pet, width = accum_scale, FUN = sum)
      accum[[paste0('accum',accum_scale)]] <- c(rep(NA, accum_scale-1), accum[[paste0('accum',accum_scale)]])
      accum[[paste0('accum',accum_scale)]] <- array(data = accum[[paste0('accum',accum_scale)]], dim = c(12, length(accum[[paste0('accum',accum_scale)]])/12))
      if (!is.null(year_dim)){
        names(dim(accum[[paste0('accum',accum_scale)]])) <- c(month_dim, year_dim)
      } else {
        names(dim(accum[[paste0('accum',accum_scale)]])) <- c(month_dim, 'fyear') 
      }
      
    }
    
  }
  
  return(accum)
}
