#library(dplyr) # for lag and lead functions
library(startR)
library(lubridate)
source("/home/Earth/llledo/Repos/WindPower/MJO/MJO_functions.R")

#===================
# Get and check args
#===================
src <- "ERAI" # Bom/ERAI for the MJO
#args <- commandArgs(trailingOnly=TRUE)
args <- c("ERA5","JFM","1981","2017","0")
if(length(args)!=5) { stop("Bad number of arguments") }
model <- args[1]
season <- args[2]
firsty <- as.numeric(args[3])
lasty <- as.numeric(args[4])
lags <- as.numeric(unlist(strsplit(args[5],",")))
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") }

#===================
# 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("Period: ",period))
print(paste0("Season: ",season))
print(paste0("Lags: ",paste(lags,collapse=" ")))

#===================
# Get daily winds 
#===================
print("Loading data")
wind <- Start(dataset=dataset,
                var='sfcWind',
                sdate=format(sdates,"%Y%m"),
                time='all',
                longitude= indices(1:100), #'all',
                latitude= indices(1:100), #'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'))
        )
#lons <- attributes(wind)$Variables$common$longitude
#lats <- attributes(wind)$Variables$common$latitude
	
#===================
# Read MJO and subset to the right period
# Remove 29 February if needed
#===================
MJO <- read_obs_MJO(src)
MJO_period <- MJO[year(MJO$vdate) %in% years & month(MJO$vdate) %in% months,]
if(season=="JFM") { 
MJO_period <- MJO_period[!(month(MJO_period$vdate)==2 & day(MJO_period$vdate)==29),] 
}


verify.fun <- function(wind,MJO_period,season) {
	library(lubridate)
	month.fcst <- lapply(1:8,function(i) { month(MJO_period[with(MJO_period,amplitude>2 & phase==i),"vdate"]) })

	#===================
	# Compute terciles and deciles for 1979-2010 for each month
	#===================
	nmonths <- dim(wind)[3]
	kk <- wind
	dim(kk) <- c(3,nmonths/3,31)
	terciles <- apply(kk,c(1),quantile,1:2/3,na.rm=T)
	names(dim(terciles)) <- c("tercile","month")

	#===================
	# Arrange wind in form (days) to match MJO_period
	#===================
	wind <- aperm(wind,c(1,2,4,3))
	dim(wind) <- 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)
	}
	wind <- wind[daysok]
	dim(wind) <- c(days=length(wind))

	#-------------------
	# Determine observed tercile for each day
	#-------------------
	obs.tercile <- wind*NA
	for (i in 1:3) {
		obs.tercile[month(MJO_period$vdate)==months[i]] <- findInterval(wind[month(MJO_period$vdate)==months[i]],terciles[,i])
	}
	obs.tercile <- t(sapply(obs.tercile+1, tabulate, nbins = 3))
	names(dim(obs.tercile)) <- c("days","tercile")

	#-------------------
	# Determine observed tercile frequencies for each MJO phase
	#-------------------
	tercile.counts <- lapply(1:8,function(i) {
		fcsts <- wind[with(MJO_period,amplitude>2 & phase==i)] 
		tercile.cat <- unlist(lapply(1:3,function(m) {findInterval(fcsts[month.fcst[[i]]==months[m]],terciles[,m])} ))
		tabulate(tercile.cat+1,nbins=3)
	})
	tercile.counts <- data.matrix(dplyr::bind_cols(tercile.counts))
	names(dim(tercile.counts)) <- c("tercile","phase")

	#-------------------
	# Build conditional climatology forecasts and verify them
	# Compute tercile.counts in Cross validation (in this case, just subtract one member in the observed category)
	# (Note: should we remove close days? there is autocorrelation at daily time-scales)
	#-------------------
	fcst.JFM <- array(rep(1/3,length(obs.tercile)),dim=dim(obs.tercile))
	clim.rps <- SpecsVerification::EnsRps(fcst.JFM,obs.tercile,format="members")
	for (i in 1:8) {
		#fcst.JFM[with(MJO_period,amplitude>2 & phase==i),] <- t(tercile.counts[,rep(i,sum(with(MJO_period,amplitude>2 & phase==i)))]) # wo cross-validation
		fcst.JFM[with(MJO_period,amplitude>2 & phase==i),] <- t(tercile.counts[,rep(i,sum(with(MJO_period,amplitude>2 & phase==i)))]) - obs.tercile[with(MJO_period,amplitude>2 & phase==i),]
	}
	condclim.rps <- SpecsVerification::EnsRps(fcst.JFM,obs.tercile,format="members")
	dim(clim.rps) <- c(days=length(clim.rps))
	dim(condclim.rps) <- c(days=length(condclim.rps))

	overall.skill <- SpecsVerification::SkillScore(condclim.rps,clim.rps)[["skillscore"]]
	ampl2.skill <- SpecsVerification::SkillScore(condclim.rps[with(MJO_period,amplitude>2)],clim.rps[with(MJO_period,amplitude>2)])[["skillscore"]] 
	byphase.skill <- abind::abind(lapply(1:8,function(i) { SpecsVerification::SkillScore(condclim.rps[with(MJO_period,amplitude>2 & phase==i)],clim.rps[with(MJO_period,amplitude>2 & phase==i)])[["skillscore"]] }) ,along=0)

	dim(overall.skill) <- c(none=1)
	dim(ampl2.skill) <- c(none=1)
	dim(byphase.skill) <- c(phase=8)
	
	return(list(overall.skill=overall.skill,ampl2.skill=ampl2.skill,byphase.skill=byphase.skill))
}

#library(multiApply)
#out <- Apply(wind,target_dims=c('dataset','var','sdate','time'),fun=verify.fun,ncores=4,MJO_period=MJO_period, season=season)
#out <- Apply(list(wind,terciles),target_dims=list(c("days"),c("tercile","month")),fun=verify.fun,ncores=8)

step <- Step(verify.fun,
             target_dims= c('dataset','var','sdate','time'),
#             target_dims= list(wind = c('dataset','var','sdate','time')),
             output_dims=list( overall.skill = c('none'),
                               ampl2.skill = c('none'),
                               byphase.skill = c('phase') )
#             output_dims=list( overall.skill,
#                               ampl2.skill,
#		               byphase.skill )
)

workflow <- AddStep(wind, step, MJO_period = MJO_period, season = season)

#===================
# Launch distibuted computations
#===================
res <- Compute(workflow$overall.skill,
               chunks = list(latitude = 2,
                             longitude = 2),
               threads_load = 2,
               threads_compute = 4, #2,
               cluster = 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',
                              job_wallclock = '01:30:00',
                              max_jobs = 4,
                              cores_per_job = 2,
                              #temp_dir = '/esarchive/scratch/llledo/MJO/condclim/tmp/',
                              bidirectional = TRUE
                             ),
#
#                 cluster = list(queue_host = 'p1',
#                                queue_type = 'slurm',
#                                temp_dir = '/gpfs/scratch/bsc32/bsc32734/startR_hpc/',
#                                lib_dir = '/gpfs/projects/bsc32/share/R_libs/3.5/',
#                                #r_module = 'R/3.5.0-foss-2018b',
#                                r_module = 'startR/0.1.2-foss-2018b-R-3.5.0',
#                                job_wallclock = '01:30:00',
#                                cores_per_job = 4,
#                                max_jobs = 4,
                                bidirectional = FALSE,
##                                polling_period = 50
#                               ),

                 ecflow_suite_dir = '/home/Earth/aho/startR_local/',
#               ecflow_suite_dir = '/esnas/scratch/llledo/MJO/condclim/ecFlow',
               wait = TRUE)

rds <- paste("/esnas/scratch/llledo/MJO/condclim/bychunk.skill",model,src,period,season,"RDS",sep=".")
saveRDS(res,rds)

#===================
# Gather results together
#===================
#library(startR)
#res <- readRDS(rds)
#out <- Collect(res)
#outrds1 <- paste("/esnas/scratch/llledo/MJO/condclim/overall.skill",model,src,period,season,"RDS",sep=".")
#outrds2 <- paste("/esnas/scratch/llledo/MJO/condclim/ampl2.skill",model,src,period,season,"RDS",sep=".")
#outrds3 <- paste("/esnas/scratch/llledo/MJO/condclim/byphase.skill",model,src,period,season,"RDS",sep=".")
#saveRDS(out$overall.skill,outrds1)
#saveRDS(out$ampl2.skill,outrds2)
#saveRDS(out$byphase.skill,outrds3)
