
# Creation:  6/2016
# Author: Nicola Cortesi
# Aim: Measure the impact of Weather Regimes of a chosen variable, from a reanalysis of from a forecast systems.
# I/O: input are all the various "*_series.RData" and "*_psl.RData" created by the script 'weather_regimes.R'
#      its output are "*_.RData" files which are need by the script weather_regimes_maps.R
# Assumption: all input/output files must be located in the 'workdir' folder
# Branch: weather_regimes


library(s2dverification)  # for the function Load()
library(abind)
library(reshape2) # after each undate of s2dverification, it's better to remove and install this package again because sometimes it gets broken! 
#library(TTR)
source('/home/Earth/ncortesi/Downloads/scripts/Rfunctions.R') # for the calendar functions
workdir <- "/scratch/Earth/ncortesi/RESILIENCE/Regimes"       # working dir
# module load NCO # : load it from the terminal before running this script interactively in the SMP machine, if you didn't load it in the .bashrc

# available reanalysis for the geopotential (g500) or the geopotential height (z500 = g500/9.81) and for var data:
ERAint <- '/esnas/reconstructions/ecmwf/eraint/$STORE_FREQ$_mean/$VAR_NAME$_f6h/$VAR_NAME$_$YEAR$$MONTH$.nc'
ERAint.name <- "ERA-Interim"

JRA55 <- '/esnas/reconstructions/jma/jra-55/$STORE_FREQ$_mean/$VAR_NAME$_f6h/$VAR_NAME$_$YEAR$$MONTH$.nc'
JRA55.name <- "JRA-55"

rean <- ERAint #JRA55 #ERAint  # choose one of the two above reanalysis from where to load the input psl data

# available forecast systems for the psl and var data:
#ECMWF_S4 <- list(path = paste0('/esnas/exp/ECMWF/seasonal/0001/s004/m001/$STORE_FREQ$_mean/$VAR_NAME$_f6h/$VAR_NAME$_$START_DATE$.nc'))
#ECMWF_S4 <- '/esnas/exp/ECMWF/seasonal/0001/s004/m001/$STORE_FREQ$_mean/psl_f6h/psl_$START_DATE$.nc'
ECMWF_S4 <- '/esnas/exp/ecmwf/system4_m1/$STORE_FREQ$_mean/$VAR_NAME$_f6h/$VAR_NAME$_$START_DATE$.nc'
ECMWF_S4.name <- "ECMWF-S4"
member.name <- "ensemble"                       # name of the dimension with the ensemble members inside the netCDF files of S4

ECMWF_monthly <- '/esnas/exp/ECMWF/monthly/ensforhc/6hourly/$VAR_NAME$/2014$MONTH$$DAY$00/$VAR_NAME$_$START_DATE$00.nc'
ECMWF_monthly.name <- "ECMWF-Monthly"

forecast <- ECMWF_S4

fields <- rean #forecast #rean # set it to 'rean' to load pressure fields and var from reanalisis, or to 'forecast' to load them from forecast systems

var.num <- 1                                     # Choose a variable. 1: sfcWind  2: tas

var.name <- c("sfcWind","tas")                   # name of the 'predictand' variable of the chosen reanalysis
var.name.full <- c("Wind Speed","Temperature")   # full name of the 'predictand' variable to put in the title of the graphs
var.unit <- c("m/s", "ºC")                       # unit of measure of var (for drawing color scales)

# Only for Reanalysis:
period <- 1:12 #13:16  # 1-12: Jan-Dec; 13-16: Winter-Spring-Summer-Autumn  [bypassed by the optional argument of the script]

# Only for Seasonal forecasts:
startM <- 9                                      # start month (1=January, 12=December)  [bypassed by the optional arguments of the script]
leadM <- 3                                       # select only the data of one lead month:  [bypassed by the optional arguments of the script]
n.members <- 15                                  # number of members of the forecast system to load
n.leadtimes <- 216                               # number of leadtimes of the forecast system to load

#############################################################################################################################
#ECMWF_monthly <- list(path = paste0('/esnas/exp/ECMWF/monthly/ensforhc/6hourly/psl/2014010200/1994_010200.nc'))
#ECMWF_monthly <- list(path = paste0('/esnas/exp/ECMWF/monthly/ensforhc/6hourly/$VAR_NAME$/2014$MONTH$$DAY$00/$VAR_NAME$_$START_DATE$00.nc'))
rean.name <- unname(ifelse(rean == ERAint, ERAint.name, JRA55.name))
forecast.name <- unname(ifelse(forecast == ECMWF_S4, ECMWF_S4.name, ECMWF_monthly.name))
fields.name <- unname(ifelse(fields == rean, rean.name, forecast.name))
n.years <- year.end - year.start + 1
#if(fields.name == forecast.name) WR.period = 1

# in case the script is run with one argument, it is assumed a reanalysis is being used and the argument becomes the month we want to perform the clustering;
# in case the script is run with two arguments, it is assumed a forecast is used and the first argument represents the startmonth and the second one the leadmonth.
# in case the script is run with no arguments, the values of the variables inside the script are used:
script.arg <- as.integer(commandArgs(TRUE))

if(length(script.arg) == 0){
   start.month <- startM
   WR.period <- start.month
   if(fields.name == forecast.name) lead.month <- leadM
}

# in case the script is run with 1 argument, it is assumed you are using a reanalysis:
if(length(script.arg) == 1){
   fields <- rean
   fields.name <- rean.name
   WR.period <-  script.arg[1]      
}

if(length(script.arg) >= 2){
   fields <- forecast
   fields.name <- forecast.name
   start.month <- script.arg[1]                 
   lead.month <- script.arg[2]
   WR.period <- start.month
}

if(length(script.arg) == 3){ # in this case, we are running the script in the SMP machine and we need to override the variables below:
   source('/gpfs/projects/bsc32/bsc32842/Rfunctions.R') # for the calendar function
   workdir <- "/gpfs/projects/bsc32/bsc32842/RESILIENCE"  # working dir
}

if(length(script.arg) >= 2) {lead.comment <- paste0("Lead month: ", lead.month)} else {lead.comment <- ""}
cat(paste0("Field: ", fields.name, " Period: ", WR.period, " ", lead.comment, "\n"))

if(lat.min >= lat.max) stop("lat.min cannot be greater than lat.max")

days.period <- n.days.period <- period.length <- list()
for (pp in 1:17){
    days.period[[pp]] <- NA
    for(y in year.start:year.end) days.period[[pp]] <- c(days.period[[pp]], n.days.in.a.future.year(year.start, y) + pos.period(y,pp))
    days.period[[pp]] <- days.period[[pp]][-1] # remove the NA at the begin that was introduced only to be able to execute the above command
    # number of days belonging to that period from year.start to year.end:
    n.days.period[[pp]] <- length(days.period[[pp]])
    # Number of days belonging to that period in a single year: 
    period.length[[pp]] <- n.days.in.a.period(pp,1999) #+ ifelse(period==13 | period==2, 1, 0)  # using year 1999 introduce a small error in winter season length because of bisestile years
}


# Load var data (interpolating it to the same reanalysis grid, if necessary):
cat("Loading data. Please wait...\n")

# Load var data:
if(fields.name == rean.name){ # Load daily var data in the reanalysis case:
    vareuFull <-array(NA,c(1,1,n.years,365, n.pos.lat, n.pos.lon))

    #for (y in year.start:year.end){
    #    var <- Load(var = var.name[var.num], exp = NULL, obs = list(fields), paste0(y,'0101'), storefreq = 'daily', leadtimemax = n.days.in.a.year(y), output = 'lonlat',
    #                latmin = lat.min, latmax = lat.max, lonmin = lon.min, lonmax = lon.max)
    #    vareuFull[seq.days.in.a.future.year(year.start, y),,] <- var$obs
    #    rm(var)
    #    gc()
    #}

    vareuFull366 <-  Load(var = var.nmae[var.num], exp = NULL, obs = list(list(path=fields)), paste0(year.start:year.end,'0101'), storefreq = 'daily', leadtimemax = 366, output = 'lonlat', latmin = lat.min, latmax = lat.max, lonmin = lon.min, lonmax = lon.max, nprocs=1)$obs

    # remove bisestile days to have all arrays with the same dimensions:
    cat("Removing bisestile days. Please wait...\n")
    vareuFull[,,,,,] <- vareuFull366[,,,1:365,,]

    for(y in year.start:year.end){
        y <- y - year.start + 1
        if(n.days.in.a.year(y) == 366) vareuFull[,,y,60:365,,] <- vareuFull366[,,y,61:366,,] # take the march to december period removing the 29th of February
    }

    rm(vareuFull366)
    gc()

    # convert psl in daily anomalies with the LOESS filter:
    cat("Calculating anomalies. Please wait...\n")
    varPeriodClim <- apply(vareuFull, c(1,2,4,5,6), mean, na.rm=T)

    varPeriodClimLoess <- varPeriodClim
    for(i in n.pos.lat){
        for(j in n.pos.lon){
             my.data <- data.frame(ens.mean=varPeriodClim[1,1,,i,j], day=1:365)
             my.loess <- loess(ens.mean ~ day, my.data, span=0.35)
             varPeriodClimLoess[1,1,,i,j] <- predict(my.loess)
        }
    }

    rm(varPeriodClim)
    gc()

    varPeriodClim2 <- InsertDim(varPeriodClimLoess, 3, n.years)
        
    ## s4.data <- data.frame(s4=pslPeriodClim[1,], day=1:216)
    ## s4.loess <- loess(s4 ~ day, s4.data, span=0.35)
    ## s4.pred <- predict(s4.loess)
        
    vareuFull <- vareuFull - varPeriodClim2    
    rm(varPeriodClim2)
    gc()
   
}


if(fields.name == ECMWF_S4.name){ # in this case, we can load only the data for 1 month at time (since each month needs ~3 GB of memory)    
    if(start.month <= 9) {start.month.char <- paste0("0",as.character(start.month))} else {start.month.char <- start.month}

    #psleuFull <- Load(var = psl, exp = list(list(path=fields)), obs = NULL, sdates=paste0(years, start.month.char,'01'), dimnames=list(member=member.name), nmember=n.members, leadtimemax=n.leadtimes, storefreq='daily', output = 'lonlat', latmin = lat.min, latmax = lat.max, lonmin = lon.min, lonmax = lon.max, nprocs=1, grid=my.grid, method='bilinear')$mod # not working

    fields2 <- gsub("\\$STORE_FREQ\\$","daily",fields)
    fields3 <- gsub("\\$VAR_NAME\\$",var.name[var.num],fields2)

    years <- c()
    for(y in year.start:year.end){
        fields4 <- gsub("\\$START_DATE\\$",paste0(y, start.month.char,'01'),fields3)

        char.lead <- nchar(system(paste0("ncks -m ",fields4,"| grep -E -i 'psl size' | cut -d '*' -f1"), intern=T))
        # beware, in this way it can only take into account leadtimes from 100 to 999, but it is the only way in the SMP machine:
        #num.lead <- as.integer(substr(system(paste0("ncks -m ",fields4,"| grep -E -i 'psl size' | cut -d '*' -f1"), intern=T),char.lead-3,char.lead)) 
        num.lead <- as.integer(system(paste0("ncks -m ",fields4,"| grep -E -i 'psl size' | cut -d '=' -f2 | cut -d '*' -f1"), intern=T)) # don't work well on the SMP
        num.memb <- as.integer(system(paste0("ncks -m ",fields4,"| grep -E -i 'psl size' | cut -d '=' -f2 | cut -d '*' -f2"), intern=T))
        num.lat  <- as.integer(system(paste0("ncks -m ",fields4,"| grep -E -i 'psl size' | cut -d '=' -f2 | cut -d '*' -f3"), intern=T))
        num.lon  <- as.integer(system(paste0("ncks -m ",fields4,"| grep -E -i 'psl size' | cut -d '=' -f2 | cut -d '*' -f4"), intern=T))

        if(num.lead == n.leadtimes && num.memb >= n.members && num.lat == 181 && num.lon == 360) years <- c(years,y)      
    }

    n.years.full <- length(years) # years really available

    #Load(var = "psl", exp = list(list(path="/esnas/exp/ecmwf/system4_m1/$STORE_FREQ$_mean/$VAR_NAME$_f6h/$VAR_NAME$_$START_DATE$.nc")), obs = NULL, sdates=paste0(1982:2013,'0101'), nmember=15, leadtimemax=216, storefreq='daily', output = 'lonlat', grid=my.grid, nprocs=1)

    vareuFull <- Load(var = var.name[var.num], exp = list(list(path=fields)), obs = NULL, sdates=paste0(years, start.month.char,'01'), dimnames=list(member=member.name), nmember=n.members, leadtimemax=n.leadtimes, storefreq='daily', output = 'lonlat', latmin = lat.min, latmax = lat.max, lonmin = lon.min, lonmax = lon.max, grid=my.grid, method='bilinear', nprocs=1)$mod

     # convert var in daily anomalies with the LOESS filter:
     cat("Calculating anomalies. Please wait...\n")
     varPeriodClim <- apply(vareuFull, c(1,4,5,6), mean, na.rm=T)
     varPeriodClimLoess <- varPeriodClim
     for(i in n.pos.lat){
        for(j in n.pos.lon){
             my.data <- data.frame(ens.mean=varPeriodClim[1,,i,j], day=1:n.leadtimes)
             my.loess <- loess(ens.mean ~ day, my.data, span=0.35)
             varPeriodClimLoess[1,,i,j] <- predict(my.loess)
        }
     }

     rm(varPeriodClim)
     gc()

     varPeriodClim2 <- InsertDim(InsertDim(varPeriodClimLoess, 2, n.years.full), 2, n.members)
        
     ## s4.data <- data.frame(s4=pslPeriodClim[1,], day=1:216)
     ## s4.loess <- loess(s4 ~ day, s4.data, span=0.35)
     ## s4.pred <- predict(s4.loess)
        
     vareuFull <- vareuFull - varPeriodClim2    
     rm(varPeriodClim2)
     gc()
}

if(fields.name == ECMWF_monthly.name){
    # Load 6-hourly psl data in the forecast case:    
    psleuFull <-array(NA, c(n.sdates*n.years, n.leadtimes, n.pos.lat, n.pos.lon))  # daily order: 19940102 19950102 ... 20130102 19940109 19950109 ... 20130109 ...
    n.leadtimes <- length(leadtime)
    sdates <- weekly.seq(forecast.year,mes,day)
    n.sdates <- length(sdates)
    
    for (startdate in 1:n.sdates){
       for(lday in leadtime){
          var <- Load(var = psl, exp = NULL, obs = list(list(path=fields)), paste0(year.start:year.end, substr(sdates[startdate],5,6), substr(sdates[startdate],7,8) ), storefreq = 'daily', leadtimemin = lday*4-3, leadtimemax = lday*4, output = 'lonlat', latmin = lat.min, latmax = lat.max, lonmin = lon.min, lonmax = lon.max)
          
          mean.lday <- apply(var$obs, c(3,5,6), mean, na.rm=T)
          rm(var)
          gc()
                           
          psleuFull[(1+(startdate-1)*n.years):(startdate*n.years), lday-leadtime[1]+1,,] <- mean.lday
          rm(mean.lday)
          gc()
       }
    }
}    



if(psl=="g500") psleuFull <- psleuFull/9.81   # convert geopotential to geopotential height (in m)
if(psl=="psl") psleuFull <- psleuFull/100   # convert MSLP in Pascal to MSLP in hPa
gc()

# impact maps:
for(p in period){
    # load regime data for chosen month (for forecasts, it load only the data correpsonding to the chosen start month p and lead month 'lead.month'::
    if(fields.name == rean.name) {
       load(file=paste0(workdir,"/",fields.name,"_",my.period[p],"_psl.RData"))

       cluster.sequence <- my.cluster[[p]]$cluster
        
       varPeriod <- vareuFull[days.period[[p]],,] # select only var data during the chosen period

       varPeriodClim <- apply(varPeriod, c(2,3), mean, na.rm=T)
       varPeriodClim2 <- InsertDim(varPeriodClim, 1, n.days.period[[p]])
       varPeriodAnom <- varPeriod - varPeriodClim2
       rm(varPeriod, varPeriodClim, varPeriodClim2)
       gc()
    
       #PlotEquiMap2(varPeriodClim[,EU], lon[EU], lat, filled.continents = FALSE, cols=my.cols.var, intxlon=10, intylat=10, cex.lab=1.5) # plot the climatology of the period    
       varPeriodAnom1 <- varPeriodAnom[wr1,,]
       varPeriodAnom2 <- varPeriodAnom[wr2,,]
       varPeriodAnom3 <- varPeriodAnom[wr3,,]
       varPeriodAnom4 <- varPeriodAnom[wr4,,]

       varPeriodAnom1mean <- apply(varPeriodAnom1,c(2,3),mean,na.rm=T)
       varPeriodAnom2mean <- apply(varPeriodAnom2,c(2,3),mean,na.rm=T)
       varPeriodAnom3mean <- apply(varPeriodAnom3,c(2,3),mean,na.rm=T)
       varPeriodAnom4mean <- apply(varPeriodAnom4,c(2,3),mean,na.rm=T)
    
       varPeriodAnomBoth1 <- abind(varPeriodAnom, varPeriodAnom1, along = 1)
       pvalue1 <- apply(varPeriodAnomBoth1, c(2,3), function(x) t.test(x[1:n.days.period[[p]]],x[(n.days.period[[p]]+1):length(x)])$p.value)
       rm(varPeriodAnomBoth1, varPeriodAnom1)
       gc()
    
       varPeriodAnomBoth2 <- abind(varPeriodAnom, varPeriodAnom2, along = 1)
       pvalue2 <- apply(varPeriodAnomBoth2, c(2,3), function(x) t.test(x[1:n.days.period[[p]]],x[(n.days.period[[p]]+1):length(x)])$p.value)
       rm(varPeriodAnomBoth2, varPeriodAnom2)
       gc()
    
       varPeriodAnomBoth3 <- abind(varPeriodAnom, varPeriodAnom3, along = 1)
       pvalue3 <- apply(varPeriodAnomBoth3, c(2,3), function(x) t.test(x[1:n.days.period[[p]]],x[(n.days.period[[p]]+1):length(x)])$p.value)
       rm(varPeriodAnomBoth3, varPeriodAnom3)
       gc()
    
       varPeriodAnomBoth4 <- abind(varPeriodAnom, varPeriodAnom4, along = 1)
       pvalue4 <- apply(varPeriodAnomBoth4, c(2,3), function(x) t.test(x[1:n.days.period[[p]]],x[(n.days.period[[p]]+1):length(x)])$p.value)
       rm(varPeriodAnomBoth4, varPeriodAnom4)
    
       rm(varPeriodAnom)
       gc()

       # save all the data necessary to redraw the graphs when we know the right regime:
       save(lon, lon.max, lon.min, lat, lat.max, lat.min, psl, psl.name, my.period, p, workdir,rean.name,fields.name,var.num,var.name,var.name.full,var.unit,year.start,year.end,WR.period, varPeriodAnom1mean, varPeriodAnom2mean, varPeriodAnom3mean,varPeriodAnom4mean,pvalue1,pvalue2,pvalue3,pvalue4,file=paste0(workdir,"/",fields.name,"_",my.period[p],"_",var.name,".RData"))
    
    }
    
    if(fields.name == forecast.name){
        load(file=paste0(workdir,"/",fields.name,"_",my.period[p],"_leadmonth_",lead.month,"_psl.RData"))


        
    }

     

    rm(pvalue1,pvalue2,pvalue3,pvalue4)
    rm(varPeriodAnom1mean,varPeriodAnom2mean,varPeriodAnom3mean,varPeriodAnom4mean)
    gc()


   
