
#!/usr/bin/env Rscript
#########################################################
# Compute sfcWind stratifications based on MJO phase.   #
# llledo@bsc.es - April 2018, February 2019             #
#########################################################
library(dplyr) # for lag and lead functions
library(startR)
library(lubridate)
source("~/Repos/WindPower/MJO/MJO_functions.R")

#===================
# Atomic function to compute stratification for each phase with amplitude > 2
# with a T-test.
#===================
# Inputs: 
#       field a vector (time)
#       MJO a datatable with time as rows and phase/amplitude columns
#       time dimensions need to match exactly
# Options:
#       lag 5 := 5 days after MJO status
#       lag -5 := 5 days before MJO status (i.e. lead)
#       ampl: a threshold to select the strength of the events
#       relative: either to compute relative or absolute anomalies
#       signif: the threshold for p.values to consider significant
# Output: a list with
#       strat: a list with the stratification maps for phases 1 to 8
#       sig95: a list with a maps of points that passed significance test for phases 1 to 8
#       ndays: number of daily means in the stratification
#       nseqs: number of consecutive day sequences in the stratifications
stratify_atomic <- function(field,MJO,season=c("JFM","OND"),lag=0,ampl=2,relative=T,signif=0.05) {
        library(plyr)

        #===================
        # Arrange wind in form (days) to match MJO
        #===================
        nmonths <- dim(field)[3]
        field <- aperm(field,c(1,2,4,3))
        dim(field) <- c(31*nmonths)
        if(season=="JFM") {
                daysok <- rep(c(rep(T,31),rep(T,28),rep(F,3),rep(T,31)),nmonths/3)
        } else if (season=="OND") {
                daysok <- rep(c(rep(T,31),rep(T,30),rep(F,1),rep(T,31)),nmonths/3)
        }
        field <- field[daysok]
        dim(field) <- c(days=length(field))

        #===================
        # Check dims of arrays
        #===================
        if(dim(field)[1] != dim(MJO)[1]) {
                stop("MJO indices and wind data have different number of days")
        }

        idx <- function(MJO,phase,ampl,lag) {
                if(lag==0) { return(MJO$phase==phase & MJO$amplitude>ampl) }
                if(lag>0) { return(dplyr::lag(MJO$phase==phase & MJO$amplitude>ampl,lag,default=F)) }
                if(lag<0) { return(dplyr::lead(MJO$phase==phase & MJO$amplitude>ampl,-1*lag,default=F)) }
        }

        strat <- laply(1:8, function(i) {
                idx2 <- idx(MJO,i,ampl,lag)
                if(relative) { return(mean(field[idx2])/mean(field) -1) }
                else { return(mean(field[idx2])-mean(field)) }
        })
        strat.t.test <- laply(1:8, function(i) {
                idx2 <- idx(MJO,i,ampl,lag)
                t.test(field[idx2],field)$p.value
                # TODO: take into account autocorrelation effect
                # See:
                #   * Wilks eq. 5.13 and/or
                #   * Zwiers and VonSotrch 1995, eqs(12,13 and 14). 
        })

        dim(strat) <- c("phase" = length(strat))
        dim(strat.t.test) <- c("phase" = length(strat.t.test))
        return(list(strat=strat,t_test=strat.t.test))
}

#===================
# Get and check args
#===================
RDS_path <- "/esarchive/scratch/llledo/MJO/RDS/"

args <- commandArgs(trailingOnly=TRUE)
#args <- c("ERA5","sfcWind","JFM","1981","2017","0","ERAI")
if(length(args)!=7) { stop("Bad number of arguments") }

model <- args[1]
var <- args[2]
season <- args[3]
firsty <- as.numeric(args[4])
lasty <- as.numeric(args[5])
lag <- as.numeric(args[6])
obs.MJO.src <- args[7]# ERAI/BoM

if(!(model %in% c("ERAI","MERRA2","ERA5"))) { stop("Unknown model") }
if(!(season %in% c("JFM","OND"))) { stop("Unknown season") }
if(firsty>=lasty) { stop("Bad period") }
if(!(obs.MJO.src %in% c("ERAI","BoM"))) { stop("Unknown MJO source") }

#===================
# Prepare sdates and paths
#===================
period <- paste0(firsty,"-",lasty)
years <- firsty:lasty
if(season=="JFM") {
        months <- 1:3
        firstd <- paste0(firsty,"0101")
}
if(season=="OND") {
        months <- 10:12
        firstd <- paste0(firsty,"1001")
}
sdates <- ymd(firstd) + months(0:2) + rep(years(0:(length(years)-1)),each=3)

if(model=="ERAI") { dataset <- "/esarchive/recon/ecmwf/erainterim/daily_mean/$var$_f6h/$var$_$sdate$.nc" }
if(model=="MERRA2") { dataset <- "/esarchive/recon/nasa/merra_v2/daily_mean/$var$_f1h/$var$_$sdate$.nc" }
if(model=="ERA5") { dataset <- "/esarchive/recon/ecmwf/era5/daily_mean/$var$_f1h/$var$_$sdate$.nc" }

#===================
# Print work summary
#===================
print(paste0("Model: ",model))
print(paste0("Variable: ",var))
print(paste0("Period: ",period))
print(paste0("Season: ",season))
print(paste0("Lag: ",lag))
print(paste0("MJO obs from: ",obs.MJO.src))

#===================
# Get daily winds 
#===================
print("Loading data")
wind <- Start(dataset=dataset,
                var=var,
                sdate=format(sdates,"%Y%m"),
                time='all',
                longitude='all',
                latitude='all',
                return_vars=list(time=NULL,latitude=NULL,longitude=NULL),
                num_procs=4, retrieve=F,
                synonims=list(longitude = c('lon', 'longitude'), latitude=c('lat', 'latitude'))
        )
#===================
# Save lats/lons in RDS
#===================
lons <- attr(wind,"Variables")$common$longitude
lats <- attr(wind,"Variables")$common$latitude
if(model=="MERRA2") { lons <- ifelse(lons<0,lons+360,lons) }


if(!file.exists(paste0(RDS_path,model,"_lons.RDS"))) {
        saveRDS(lons,paste0(RDS_path,model,"_lons.RDS"))
}
if(!file.exists(paste0(RDS_path,model,"_lats.RDS"))) {
        saveRDS(lats,paste0(RDS_path,model,"_lats.RDS"))
}

#===================
# Read MJO and subset to the right period
# Remove 29 February if needed
#===================
MJO <- read_obs_MJO(obs.MJO.src)
MJO <- MJO[year(MJO$vdate) %in% years & month(MJO$vdate) %in% months,]
if(season=="JFM") { MJO <- MJO[!(month(MJO$vdate)==2 & day(MJO$vdate)==29),] }

#===================
# Prepare StartR step and workflow
#===================
step <- Step(stratify_atomic,
             target_dims=list( field = c('dataset','var','sdate','time') ),
             output_dims=list( strat = c('phase'),
                               t_test = c('phase') )
)

workflow <- AddStep(wind, step, MJO=MJO, season=season, lag=lag)

#===================
# Different cluster configurations
#===================
p9config <- list(
        queue_host = 'p9login1.bsc.es',
        #queue_host = 'p9login1',
        queue_type = 'slurm',
        temp_dir = '/gpfs/scratch/bsc32/bsc32804/startR_hpc/',
        #temp_dir = '/esarchive/scratch/llledo/startR_hpc/',
        lib_dir = '/gpfs/projects/bsc32/share/R_libs/3.5/',
        r_module = 'R/3.5.0-foss-2018b',
        cores_per_job = 2,
        job_wallclock = '10:00:00',
        max_jobs = 4,
        bidirectional = FALSE,
        polling_period = 10
)

fatconfig <- list(
        queue_host = 'bsceslogin01.bsc.es',
        queue_type = 'slurm',
        #lib_dir = '~nmanuben/tmp/startR_mod/startR.Rcheck',
        #r_module = 'R/3.2.0-foss-2015a-bare',
        #ecflow_module = 'ecFlow/4.9.0-foss-2015a',
        temp_dir = "/esarchive/scratch/llledo/MJO/ecFlow/strat/bychunk",
        cores_per_job = 2,
        job_wallclock = '12:00:00',
        max_jobs = 4,
        bidirectional = TRUE
)

#===================
# Launch distibuted computations
#===================
res <- Compute(workflow$strat,
        chunks = list(latitude = 2,
                      longitude = 2),
        threads_load = 2,
        threads_compute = 2,

        cluster = fatconfig,
        #cluster = p9config,

        ecflow_suite_dir = '/esnas/scratch/llledo/MJO/stratifications/ecFlow',
        wait = FALSE
)

lag.str <- ifelse(lag<0,paste0("lead",-1*lag),paste0("lag",lag))
rds <- paste("/esnas/scratch/llledo/MJO/stratifications/bychunk.strat",model,var,obs.MJO.src,period,season,lag.str,"RDS",sep=".")
saveRDS(res,rds)


