# Test Script to compute EnsCorr, FairRPSS, FairCRPSS (from SpecsVerification) 
# to ECMWF subseasonal hindcasts 
# Only 1 startdate 4/1/2016
# with and without bia correction
# Uses StartR and chunking on 'time' dimension
# Author: Andrea Manrique  
# April 2018

library(multiApply)

library(startR)
library(easyVerification)  # for veriApply() function
library(SpecsVerification) # for skill scores




cfs.name <- 'ecmwf'
rean.name <- "ERA-Interim"#  "MSWEP" 
forecast.year<-2016
yr1.hind <- 1996  #1996 # first hindcast year
yr2.hind <- 2015 ##2013   # last hindcast year (usually the forecast year -1)
var_name <- 'tas' #'prlr' #
var.name.map <-'2m Temperature' # '10m Wind Speed' #'Precipitation' # '2m Temperature'

cal.method <- 'none'#,#'simple_bias'   # 'none' ¿Como introducirlo somo argumento sin ser un objeto de StartR_header?


ecmwf_path <- paste0('/esarchive/exp/ecmwf/s2s-monthly_ensforhc/weekly_mean/$var$_f6h/$sdate$/$var$_$syear$.nc')

obs_path <- paste0('/esnas/recon/ecmwf/erainterim/weekly_mean/',
                     '$var$-r240x121_f6h/$var$_$syear$$sdate$.nc')

sdates.seq<- format(as.Date(paste(forecast.year,04,07,sep='-')),format='%Y%m%d')

exp <- Start(dat = ecmwf_path, 
              var = var_name,
              sdate = sdates.seq,
              syear = 'all',     # before hdate
              time = 'all',
              ensemble = "all",
              latitude = indices(1:121),
              longitude = indices(1:240),
              retrieve = F,
              syear_depends = 'sdate',
              return_vars = list(latitude = NULL,
                                longitude = NULL)) 

sdate_mmdd<-substr(sdates.seq,5,8)

obs <- Start(dat = obs_path,
              var = var_name,
              sdate = sdate_mmdd,
              syear = sprintf('%4d',yr1.hind:yr2.hind),
              time = 'all',
              latitude = indices(121:1),
              longitude = indices(1:240),
              retrieve = F)



score_calc<-function(forecast,reference,cal.method) {

   source('/esarchive/scratch/amanriqu/scripts_for_p9/Calibration_fcst.R')

        anomaly_simple<-function(data){
	  avg<- Apply(data,c('syear','ensemble'),mean)[[1]]
	  data<-Apply(data,c('sdate'),function(x) x-avg)[[1]]
	  return(data)
	  }   

	default_dims <- c(sdate = 1, syear = 1, ensemble = 1)
	default_dims[names(dim(forecast))] <- dim(forecast)
	dim(forecast) <- default_dims
	default_dims <- c(sdate = 1, syear = 1, ensemble = 1)
	default_dims[names(dim(reference))] <- dim(reference)
	dim(reference) <- default_dims

        cal.method <- 'none'#'simple_bias' #'raw'## this was meant to go as argument
        if (cal.method=='simple_bias'){ 
        calibrated<-BiasCorrection(var_obs=reference[1,,1],var_exp=forecast[1,,])

        forecast <-aperm(calibrated,c(2,1))
        # Add dimension names
        names(dim(forecast))<-list('syear','ensemble') # nombro las que hay
        default_dims[names(dim(forecast))] <- dim(forecast)                                      # añado la que falta
        dim(forecast) <- default_dims
        }

	# Anomalies

	forecast <- anomaly_simple(forecast)
	reference <-anomaly_simple(reference)



        forecast_month <- forecast
        reference_month <- reference
	forecast_month_reshaped <- array(forecast_month, 
			          c(dim(forecast_month)['sdate']*dim(forecast_month)['syear'],
			            dim(forecast_month)['ensemble']))
	reference_month_reshaped <- array(reference_month, 
			          c(dim(reference_month)['sdate']*dim(reference_month)['syear'],
			            dim(reference_month)['ensemble']))

	my.prob=c(1/3,2/3)

	print('Computing FairRpss')

	my.FairRpss <- veriApply("FairRpss", fcst=forecast_month_reshaped, obs=reference_month_reshaped, prob=my.prob, tdim=1, ensdim=2)
	print('Computing FairCrpss')
	my.FairCrpss <- veriApply("FairCrpss", fcst=forecast_month_reshaped, obs=reference_month_reshaped, tdim=1, ensdim=2)
	# Corr of the ensemble mean:
	print('Computing ensemble mean')
	fcst_mean<-Apply(forecast_month_reshaped,'ensemble',mean)[[1]]   
	print('Computing ensemble mean corr')
	my.Enscorr<- cor(fcst_mean,reference_month_reshaped, use="complete.obs")

	my.Enscorr.pvalue<-cor.test(fcst_mean,reference_month_reshaped, use="complete.obs")$p.value
	FairRpss_total<-my.FairRpss$skillscore
		       
	FairRpss.sd<-my.FairRpss$skillscore.sd
		 
	FairCrpss_total<-my.FairCrpss$skillscore

	FairCrpss.sd<-my.FairCrpss$skillscore.sd

       dim(FairRpss_total)<-c(score = 1)
       dim(FairRpss.sd)<-c(score = 1)
       dim(FairCrpss_total)<-c(score = 1)
       dim(FairCrpss.sd)<-c(score = 1)
       dim(my.Enscorr)<-c(score = 1)
       dim(my.Enscorr.pvalue)<-c(score = 1)  

	return(c(list(FairRpss = FairRpss_total),
                 list(FairRpss_sd = FairRpss.sd),
                 list(FairCrpss = FairCrpss_total),
                 list(FairCrpss_sd = FairCrpss.sd),
	         list(Enscorr = my.Enscorr), 
	         list(Enscorr.pvalue = my.Enscorr.pvalue)))
		        
} # end function
#############################################################################



step <- Step(score_calc, 
             list(forecast = c('sdate','syear','ensemble'),
                  reference = c('sdate', 'syear')),
             list(FairRpss = c('score'),
                  FairRpss_sd = c('score'),
                  FairCrpss = c('score'),
                  FairCrpss_sd = c('score'),
                  Enscorr = c('score'),
                  Enscorr.pvalue = c('score')),
             use_libraries = c('easyVerification',
                              'SpecsVerification',
                               's2dverification'))

                   
wf <- AddStep(list(exp,obs), step)


# new june2020< Nord III          
## on Nord 3
#-----------modify according to your personal info-----------
  queue_host = 'nord3'   #your own host name for power9
  temp_dir =  '/gpfs/scratch/bsc32/bsc32339/startR_hpc/'
  ecflow_suite_dir = '/home/Earth/nperez/startR_local/'  #your own local directory
 
res <- Compute(wf$FairRpss,    
              chunks = list(time = 4,
                            latitude = 4,
                            longitude = 4),
              threads_load = 1,
              threads_compute = 4,
              cluster = list(queue_host = queue_host,
                   queue_type = 'lsf',
                   temp_dir = temp_dir,
                   cores_per_job = 2,
                   job_wallclock = '05:00',
                   max_jobs = 64,
                   #extra_queue_params = list('#BSUB -q bsc_es'),
                   bidirectional = FALSE,
                   polling_period = 10),
              ecflow_suite_dir = ecflow_suite_dir,
              wait = TRUE)





# test in local:
res <- Compute(wf$FairRpss,    
              chunks = list(time = 4,
                            latitude = 4,
                            longitude = 4))
n.members<-11
FairRpss<-res$FairRpss 
FairRpss_sd<-res$FairRpss_sd
FairCrpss<-res$FairCrpss 
FairCrpss_sd<-res$FairCrpss_sd
Enscorr <- res$Enscorr
Enscorr.pvalue <- res$Enscorr.pvalue
my.prob<-c(1/3,2/3)

# Obtain values of lat and lon and return them for plotting
  
lons<- attr(exp,"Variables")$common$longitude
lats<- attr(exp,"Variables")$common$latitude


dim(lons) <- c('longitude' = length(lons))
dim(lats) <- c('latitude' = length(lats))

leadtime.week <- c('5-11','12-18','19-25','26-32')

print("Saving results...")
save(FairRpss, FairRpss_sd, FairCrpss, FairCrpss_sd, Enscorr, Enscorr.pvalue, yr1.hind, yr2.hind, n.members, my.prob, lons, lats, leadtime.week, var.name.map, cfs.name,
rean.name,
file=paste0(work.dir,'/',cfs.name,'/',var_name,
            '_','scores_july_',cal.method,'_ByStartDate','.RData'))





